app/services/tse_api.py

"""
TSE (Tribunal Superior Eleitoral) API Service
Access to Brazilian electoral data - candidates, assets, donations
"""
import httpx
from typing import Optional, Dict, Any, List
from dataclasses import dataclass, field


# DivulgaCand API (unofficial but functional)
TSE_DIVULGACAND_URL = "https://divulgacandcontas.tse.jus.br/divulga/rest/v1"


@dataclass
class Candidato:
    """Electoral candidate data"""
    id: int
    nome: str
    nome_urna: str
    cpf_parcial: str = ""  # TSE only shows partial
    numero: str = ""
    cargo: str = ""
    partido_sigla: str = ""
    partido_nome: str = ""
    coligacao: str = ""
    situacao: str = ""
    
    # Location
    uf: str = ""
    municipio: str = ""
    
    # Personal
    data_nascimento: str = ""
    genero: str = ""
    grau_instrucao: str = ""
    ocupacao: str = ""
    
    # Assets
    total_bens: float = 0.0
    bens: List[Dict[str, Any]] = field(default_factory=list)
    
    # Campaign
    total_receitas: float = 0.0
    total_despesas: float = 0.0


@dataclass
class Eleicao:
    """Election metadata"""
    id: int
    ano: int
    descricao: str
    turno: int = 1


async def listar_eleicoes() -> List[Eleicao]:
    """List available elections"""
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            response = await client.get(f"{TSE_DIVULGACAND_URL}/eleicao/ordinarias")
            
            if response.status_code != 200:
                return []
            
            data = response.json()
            eleicoes = []
            
            for item in data:
                eleicoes.append(Eleicao(
                    id=item.get("id", 0),
                    ano=item.get("ano", 0),
                    descricao=item.get("descricaoEleicao", ""),
                    turno=item.get("turno", 1)
                ))
            
            return sorted(eleicoes, key=lambda x: x.ano, reverse=True)
            
    except Exception as e:
        print(f"TSE eleicoes error: {e}")
        return []


async def buscar_candidatos(
    nome: str,
    ano: int = 2024,
    uf: Optional[str] = None,
    cargo: Optional[str] = None
) -> List[Candidato]:
    """
    Search for candidates by name.
    
    Args:
        nome: Candidate name to search
        ano: Election year (default 2024)
        uf: State filter (optional)
        cargo: Position filter (optional)
    """
    try:
        # First get the election ID for the year
        eleicoes = await listar_eleicoes()
        eleicao = next((e for e in eleicoes if e.ano == ano), None)
        
        if not eleicao:
            # Try common election IDs
            eleicao_id = {2024: 546, 2022: 544, 2020: 426, 2018: 295}.get(ano, 546)
        else:
            eleicao_id = eleicao.id
        
        # Build search URL
        base_url = f"{TSE_DIVULGACAND_URL}/candidatura/listar/{ano}/{eleicao_id}"
        
        params = {"nomeCompleto": nome}
        if uf:
            params["uf"] = uf.upper()
        if cargo:
            params["cargo"] = cargo
        
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(base_url, params=params)
            
            if response.status_code != 200:
                return []
            
            data = response.json()
            candidatos_data = data.get("candidatos", [])
            
            candidatos = []
            for item in candidatos_data:
                candidatos.append(Candidato(
                    id=item.get("id", 0),
                    nome=item.get("nomeCompleto", ""),
                    nome_urna=item.get("nomeUrna", ""),
                    cpf_parcial=item.get("cpf", "")[:3] + ".***.***-**" if item.get("cpf") else "",
                    numero=str(item.get("numero", "")),
                    cargo=item.get("cargo", {}).get("nome", "") if isinstance(item.get("cargo"), dict) else str(item.get("cargo", "")),
                    partido_sigla=item.get("partido", {}).get("sigla", "") if isinstance(item.get("partido"), dict) else "",
                    partido_nome=item.get("partido", {}).get("nome", "") if isinstance(item.get("partido"), dict) else "",
                    uf=item.get("ufSigla", "") or item.get("uf", ""),
                    municipio=item.get("municipio", {}).get("nome", "") if isinstance(item.get("municipio"), dict) else "",
                    situacao=item.get("situacao", ""),
                    total_bens=float(item.get("totalDeBens", 0) or 0)
                ))
            
            return candidatos
            
    except Exception as e:
        print(f"TSE search error: {e}")
        return []


async def obter_candidato_detalhes(
    id_candidato: int,
    ano: int = 2024,
    eleicao_id: Optional[int] = None
) -> Optional[Candidato]:
    """Get detailed candidate information including assets"""
    try:
        if not eleicao_id:
            eleicao_id = {2024: 546, 2022: 544, 2020: 426, 2018: 295}.get(ano, 546)
        
        async with httpx.AsyncClient(timeout=30.0) as client:
            # Get candidate details
            response = await client.get(
                f"{TSE_DIVULGACAND_URL}/candidatura/buscar/{ano}/{eleicao_id}/candidato/{id_candidato}"
            )
            
            if response.status_code != 200:
                return None
            
            item = response.json()
            
            candidato = Candidato(
                id=item.get("id", 0),
                nome=item.get("nomeCompleto", ""),
                nome_urna=item.get("nomeUrna", ""),
                numero=str(item.get("numero", "")),
                cargo=item.get("cargo", {}).get("nome", "") if isinstance(item.get("cargo"), dict) else "",
                partido_sigla=item.get("partido", {}).get("sigla", "") if isinstance(item.get("partido"), dict) else "",
                partido_nome=item.get("partido", {}).get("nome", "") if isinstance(item.get("partido"), dict) else "",
                uf=item.get("ufSigla", ""),
                municipio=item.get("localCandidatura", ""),
                situacao=item.get("situacao", ""),
                data_nascimento=item.get("dataNascimento", ""),
                genero=item.get("genero", ""),
                grau_instrucao=item.get("grauInstrucao", ""),
                ocupacao=item.get("ocupacao", ""),
                total_bens=float(item.get("totalDeBens", 0) or 0)
            )
            
            # Try to get assets (bens)
            try:
                bens_response = await client.get(
                    f"{TSE_DIVULGACAND_URL}/candidatura/buscar/{ano}/{eleicao_id}/candidato/{id_candidato}/bens"
                )
                if bens_response.status_code == 200:
                    bens_data = bens_response.json()
                    candidato.bens = [
                        {
                            "tipo": b.get("tipoBem", ""),
                            "descricao": b.get("descricao", ""),
                            "valor": float(b.get("valor", 0) or 0)
                        }
                        for b in bens_data
                    ]
            except:
                pass
            
            return candidato
            
    except Exception as e:
        print(f"TSE details error: {e}")
        return None


async def buscar_politico(nome: str) -> Dict[str, Any]:
    """
    Search for a politician across multiple elections.
    Returns consolidated information.
    """
    resultado = {
        "nome": nome,
        "encontrado": False,
        "candidaturas": [],
        "ultimo_cargo": None,
        "total_patrimonio": 0.0,
        "partidos": set(),
        "ufs": set()
    }
    
    # Search in recent elections - continue through ALL years
    for ano in [2024, 2022, 2020, 2018]:
        try:
            candidatos = await buscar_candidatos(nome, ano=ano)
            print(f"TSE: Buscando '{nome}' em {ano} - encontrados: {len(candidatos)}")
            
            for c in candidatos:
                # Match if nome is in the candidate's full name
                if nome.lower() in c.nome.lower() or nome.lower() in c.nome_urna.lower():
                    resultado["encontrado"] = True
                    resultado["candidaturas"].append({
                        "ano": ano,
                        "cargo": c.cargo,
                        "partido": c.partido_sigla,
                        "uf": c.uf,
                        "situacao": c.situacao,
                        "patrimonio": c.total_bens
                    })
                    
                    if c.partido_sigla:
                        resultado["partidos"].add(c.partido_sigla)
                    if c.uf:
                        resultado["ufs"].add(c.uf)
                    
                    if c.total_bens > resultado["total_patrimonio"]:
                        resultado["total_patrimonio"] = c.total_bens
                    
                    if not resultado["ultimo_cargo"]:
                        resultado["ultimo_cargo"] = f"{c.cargo} ({ano})"
        except Exception as e:
            print(f"TSE search {ano} error: {e}")
            continue
    
    # Convert sets to lists for JSON
    resultado["partidos"] = list(resultado["partidos"])
    resultado["ufs"] = list(resultado["ufs"])
    
    print(f"TSE resultado para '{nome}': encontrado={resultado['encontrado']}, candidaturas={len(resultado['candidaturas'])}")
    
    return resultado