main.py

import requests
import time
import os
import sys
from pymongo import MongoClient
from pymongo.errors import ServerSelectionTimeoutError

# Configurações
IBGE_URL = "https://servicodados.ibge.gov.br/api/v1/localidades/municipios"
VIA_CEP_URL = "https://viacep.com.br/ws/{UF}/{CIDADE}/{LOGRADOURO}/json/"
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017/")
DELAY_BETWEEN_REQUESTS = 1.5  # Segundos para evitar bloqueio

# Globais para coleções
municipios_col = None
ceps_col = None

def fetch_municipios():
    print("Coletando municípios do IBGE...")
    try:
        response = requests.get(IBGE_URL)
        response.raise_for_status()
        municipios = response.json()
        print(f"Total de {len(municipios)} municípios coletados.")
        return municipios
    except Exception as e:
        print(f"Erro ao coletar municípios: {e}")
        return []

def fetch_ceps(uf, cidade, logradouro="Rua"):
    url = VIA_CEP_URL.format(UF=uf, CIDADE=cidade, LOGRADOURO=logradouro)
    try:
        response = requests.get(url)
        if response.status_code == 200:
            return response.json()
        elif response.status_code == 429:
            print("Limite de requisições atingido (429). Aguardando mais tempo...")
            time.sleep(60)
        else:
            print(f"Erro {response.status_code} para {uf}/{cidade}")
    except Exception as e:
        print(f"Erro na requisição ViaCEP: {e}")
    return []

def main():
    # Conexão com MongoDB
    try:
        client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
        # Testa a conexão
        client.admin.command('ping')
        db = client.cep_database
        global municipios_col, ceps_col
        municipios_col = db.municipios
        ceps_col = db.ceps
        
        # Criar índices para melhorar performance
        municipios_col.create_index("ibge_id", unique=True)
        ceps_col.create_index("ibge_id")
        ceps_col.create_index("cep", unique=True)
        print("Conectado ao MongoDB com sucesso.")
    except ServerSelectionTimeoutError:
        print(f"Erro: Não foi possível conectar ao MongoDB no URI: {MONGO_URI}")
        print("Certifique-se de que o banco de dados está rodando.")
        sys.exit(1)
    except Exception as e:
        print(f"Erro inesperado ao conectar ao MongoDB: {e}")
        sys.exit(1)

    municipios = fetch_municipios()
    
    for mun in municipios:
        try:
            nome_municipio = mun['nome']
            uf = mun.get('microrregiao', {}).get('mesorregiao', {}).get('UF', {}).get('sigla')
            
            # Caso a estrutura seja diferente, tenta buscar na regiao-imediata
            if not uf:
                uf = mun.get('regiao-imediata', {}).get('regiao-intermediaria', {}).get('UF', {}).get('sigla')
            
            if not uf:
                print(f"Aviso: Não foi possível encontrar a UF para o município {nome_municipio}")
                continue

            ibge_id = mun['id']
            
            # Salvar município no banco se não existir
            if not municipios_col.find_one({"ibge_id": ibge_id}):
                municipios_col.insert_one({
                    "ibge_id": ibge_id,
                    "nome": nome_municipio,
                    "uf": uf
                })
                print(f"Salvo município: {nome_municipio} - {uf}")

            # Verificar se já processamos os CEPs deste município
            if ceps_col.find_one({"ibge_id": ibge_id}):
                print(f"CEPs de {nome_municipio} já estão no banco. Pulando...")
                continue

            print(f"Consultando CEPs de {nome_municipio}...")
            
            # O ViaCEP exige logradouro de pelo menos 3 caracteres. 
            # Como o objetivo é pedagógico/base pequena, usaremos "Rua" como exemplo.
            ceps = fetch_ceps(uf, nome_municipio, "Rua")
            
            if ceps and isinstance(ceps, list):
                for cep_data in ceps:
                    cep_data['ibge_id'] = ibge_id  # Relacionar com o município
                    ceps_col.update_one({"cep": cep_data['cep']}, {"$set": cep_data}, upsert=True)
                print(f"Encontrados {len(ceps)} CEPs para {nome_municipio}")
            
            # Delay obrigatório para evitar bloqueio
            time.sleep(DELAY_BETWEEN_REQUESTS)
        except Exception as e:
            print(f"Erro ao processar município {mun.get('nome', 'desconhecido')}: {e}")

if __name__ == "__main__":
    main()