app/services/ibge_api.py

"""
IBGE API Service
Access to Brazilian geographic and demographic data
"""
import httpx
from typing import Optional, Dict, Any, List
from dataclasses import dataclass


IBGE_BASE_URL = "https://servicodados.ibge.gov.br/api/v1"


@dataclass
class Estado:
    """Brazilian state data"""
    id: int
    sigla: str
    nome: str
    regiao: str


@dataclass
class Municipio:
    """Brazilian municipality data"""
    id: int
    nome: str
    estado_sigla: str
    estado_nome: str
    regiao: str
    # Optional enriched data
    populacao: Optional[int] = None
    area_km2: Optional[float] = None


async def listar_estados() -> List[Estado]:
    """List all Brazilian states"""
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(f"{IBGE_BASE_URL}/localidades/estados")
            
            if response.status_code != 200:
                return []
            
            data = response.json()
            estados = []
            
            for item in data:
                estados.append(Estado(
                    id=item["id"],
                    sigla=item["sigla"],
                    nome=item["nome"],
                    regiao=item.get("regiao", {}).get("nome", "")
                ))
            
            return sorted(estados, key=lambda x: x.nome)
            
    except Exception as e:
        print(f"IBGE estados error: {e}")
        return []


async def listar_municipios(uf: str) -> List[Municipio]:
    """List all municipalities in a state"""
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(
                f"{IBGE_BASE_URL}/localidades/estados/{uf}/municipios"
            )
            
            if response.status_code != 200:
                return []
            
            data = response.json()
            municipios = []
            
            for item in data:
                municipios.append(Municipio(
                    id=item["id"],
                    nome=item["nome"],
                    estado_sigla=uf.upper(),
                    estado_nome=item.get("microrregiao", {}).get("mesorregiao", {}).get("UF", {}).get("nome", ""),
                    regiao=item.get("microrregiao", {}).get("mesorregiao", {}).get("UF", {}).get("regiao", {}).get("nome", "")
                ))
            
            return sorted(municipios, key=lambda x: x.nome)
            
    except Exception as e:
        print(f"IBGE municipios error: {e}")
        return []


async def buscar_municipio(nome: str, uf: Optional[str] = None) -> List[Municipio]:
    """Search for municipalities by name"""
    try:
        # If UF provided, search only that state
        if uf:
            municipios = await listar_municipios(uf)
            return [m for m in municipios if nome.lower() in m.nome.lower()]
        
        # Otherwise search all states (slower)
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(f"{IBGE_BASE_URL}/localidades/municipios")
            
            if response.status_code != 200:
                return []
            
            data = response.json()
            results = []
            
            for item in data:
                if nome.lower() in item["nome"].lower():
                    uf_info = item.get("microrregiao", {}).get("mesorregiao", {}).get("UF", {})
                    results.append(Municipio(
                        id=item["id"],
                        nome=item["nome"],
                        estado_sigla=uf_info.get("sigla", ""),
                        estado_nome=uf_info.get("nome", ""),
                        regiao=uf_info.get("regiao", {}).get("nome", "")
                    ))
            
            return results[:20]  # Limit results
            
    except Exception as e:
        print(f"IBGE search error: {e}")
        return []


async def obter_municipio_por_id(id_municipio: int) -> Optional[Municipio]:
    """Get municipality by IBGE code"""
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(
                f"{IBGE_BASE_URL}/localidades/municipios/{id_municipio}"
            )
            
            if response.status_code != 200:
                return None
            
            item = response.json()
            uf_info = item.get("microrregiao", {}).get("mesorregiao", {}).get("UF", {})
            
            return Municipio(
                id=item["id"],
                nome=item["nome"],
                estado_sigla=uf_info.get("sigla", ""),
                estado_nome=uf_info.get("nome", ""),
                regiao=uf_info.get("regiao", {}).get("nome", "")
            )
            
    except Exception as e:
        print(f"IBGE municipio error: {e}")
        return None


async def enriquecer_localizacao(cidade: str, uf: Optional[str] = None) -> Dict[str, Any]:
    """
    Enrich a location name with IBGE data.
    Useful for adding context to extracted locations.
    """
    resultado = {
        "cidade_original": cidade,
        "encontrado": False,
        "ibge_codigo": None,
        "cidade": None,
        "estado": None,
        "estado_sigla": None,
        "regiao": None
    }
    
    municipios = await buscar_municipio(cidade, uf)
    
    if municipios:
        # Take best match (exact or first)
        melhor = None
        for m in municipios:
            if m.nome.lower() == cidade.lower():
                melhor = m
                break
        
        if not melhor:
            melhor = municipios[0]
        
        resultado.update({
            "encontrado": True,
            "ibge_codigo": melhor.id,
            "cidade": melhor.nome,
            "estado": melhor.estado_nome,
            "estado_sigla": melhor.estado_sigla,
            "regiao": melhor.regiao
        })
    
    return resultado