#!/usr/bin/env python3
"""
Enriquece os dados existentes (lelyon_prospects_latest.json) com dados REAIS
da API CNPJ Finder (cnpjfinder.com.br) - grátis, sem auth, 50 req/s
Mantém estrutura B2B/B2C mas preenche email, telefone, endereço, sócios reais
"""

import json
import csv
import time
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional
import pandas as pd
import requests

# ============ CONFIG ============
CNPJ_FINDER_API = "https://www.cnpjfinder.com.br/api/v1/cnpj/"
DATA_DIR = Path("/root/lelyon_prospecting/data")
INPUT_JSON = DATA_DIR / "lelyon_real_minhareceita_latest.json"
OUTPUT_JSON = DATA_DIR / f"lelyon_prospects_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
OUTPUT_CSV = DATA_DIR / f"lelyon_prospects_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
OUTPUT_EXCEL = DATA_DIR / f"lelyon_prospects_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
PROGRESS_FILE = DATA_DIR / "enrichment_real_progress.json"

DELAY_BETWEEN_REQUESTS = 0.05  # 20 req/s (bem dentro do limite de 50/s)
MAX_RETRIES = 3
TIMEOUT = 30

session = requests.Session()
session.headers.update({"Accept": "application/json"})

# ============ HELPERS ============
def load_progress() -> Dict:
    if PROGRESS_FILE.exists():
        with open(PROGRESS_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return {"processed_cnpjs": [], "enriched": {}, "failed": []}

def save_progress(progress: Dict):
    with open(PROGRESS_FILE, "w", encoding="utf-8") as f:
        json.dump(progress, f, ensure_ascii=False, indent=2)

def load_companies() -> Dict:
    """Carrega empresas do JSON latest - retorna dict cnpj->company."""
    if INPUT_JSON.exists():
        with open(INPUT_JSON, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, dict):
                companies = {}
                for cat in ["b2b", "b2c"]:
                    for c in data.get(cat, []):
                        cnpj = c.get("cnpj", "").replace(".", "").replace("/", "").replace("-", "")
                        if cnpj:
                            companies[cnpj] = c
                print(f"✅ Carregadas {len(companies)} empresas de {INPUT_JSON}")
                return companies
    raise FileNotFoundError(f"Arquivo não encontrado: {INPUT_JSON}")

def cnpj_finder_lookup(cnpj: str) -> Optional[Dict]:
    """Consulta CNPJ na API CNPJ Finder."""
    url = f"{CNPJ_FINDER_API}{cnpj}"
    
    for attempt in range(MAX_RETRIES):
        try:
            resp = session.get(url, timeout=TIMEOUT)
            
            if resp.status_code == 200:
                data = resp.json()
                if data.get("success"):
                    return data.get("data")
                else:
                    return None
            elif resp.status_code == 429:
                wait = 2 ** attempt * 2
                print(f"   ⏳ Rate limited, aguardando {wait}s...")
                time.sleep(wait)
                continue
            elif resp.status_code == 404:
                return None
            else:
                print(f"   ⚠️ Erro {resp.status_code}: {resp.text[:100]}")
        except requests.Timeout:
            print(f"   ⏱️ Timeout (tentativa {attempt+1})")
        except Exception as e:
            print(f"   ❌ Erro: {e}")
        
        if attempt < MAX_RETRIES - 1:
            time.sleep(2 ** attempt)
    
    return None

def merge_company_data(original: Dict, real: Dict) -> Dict:
    """Mescla dados originais com dados reais da API."""
    merged = original.copy()
    
    # Campos que vêm da API real
    merged.update({
        "email": real.get("email", "") or "",
        "telefone": format_phone(real.get("telefones", [])),
        "website": real.get("website", "") or "",
        "endereco_completo": format_address(real),
        "contato_nome": get_first_socio(real),
        "contato_cargo": get_first_socio_role(real),
        "capital_social": parse_capital(real.get("capital_social")),
        "natureza_juridica": real.get("natureza_juridica", ""),
        "simples_nacional": "Sim" if real.get("opcao_pelo_simples") == "S" else "Não",
        "data_situacao": real.get("data_situacao_cadastral", ""),
        "motivo_situacao": real.get("descricao_motivo_situacao_cadastral", ""),
        "cnaes_secundarios": format_cnaes_secundarios(real.get("cnaes_secundarios", [])),
        "qsa": format_qsa(real.get("QSA", [])),
        "fonte_enriquecimento": "cnpj_finder",
        "data_enriquecimento": datetime.now().isoformat(),
        "observacoes": f"Dados reais CNPJ Finder | {original.get('observacoes', '')}",
    })
    
    # Manter campos que a API não tem
    for key in ["tipo_prospeccao", "b2c_categoria", "best_fit_service", 
                "fit_score", "lelyon_fit", "whatsapp", "linkedin", 
                "instagram", "facebook", "tipo_estabelecimento", "gestao",
                "latitude", "longitude", "especialidades", "leitos_total",
                "leitos_sus", "leitos_uti", "data_atualizacao_cnes"]:
        if key in original:
            merged[key] = original[key]
    
    return merged

def format_phone(telefones: List) -> str:
    if not telefones:
        return ""
    parts = []
    for t in telefones[:2]:
        ddd = t.get("ddd", "")
        num = t.get("numero", "")
        if ddd and num:
            clean = num.replace("-", "").replace(" ", "")
            parts.append(f"({ddd}) {clean[:4]}-{clean[4:]}" if len(clean) == 8 else f"({ddd}) {clean}")
    return "; ".join(parts)

def format_address(real: Dict) -> str:
    parts = []
    if real.get("tipo_logradouro") and real.get("logradouro"):
        parts.append(f"{real['tipo_logradouro']} {real['logradouro']}")
    if real.get("numero"):
        parts.append(str(real["numero"]))
    if real.get("complemento"):
        parts.append(real["complemento"])
    if real.get("bairro"):
        parts.append(real["bairro"])
    if real.get("municipio") and real.get("uf"):
        parts.append(f"{real['municipio']}/{real['uf']}")
    if real.get("cep"):
        cep = str(real["cep"]).zfill(8)
        parts.append(f"CEP {cep[:5]}-{cep[5:]}")
    return ", ".join(parts)

def get_first_socio(real: Dict) -> str:
    qsa = real.get("QSA", [])
    if qsa:
        return qsa[0].get("nome_socio", "")
    return ""

def get_first_socio_role(real: Dict) -> str:
    qsa = real.get("QSA", [])
    if qsa:
        return qsa[0].get("qualificacao_socio", "")
    return ""

def parse_capital(capital_str) -> float:
    if not capital_str:
        return 0.0
    try:
        return float(str(capital_str).replace(",", "."))
    except:
        return 0.0

def format_cnaes_secundarios(cnaes_sec) -> str:
    if isinstance(cnaes_sec, str):
        return cnaes_sec
    if not cnaes_sec:
        return ""
    if not isinstance(cnaes_sec, list):
        return str(cnaes_sec)
    # Handle both list of objects {'codigo': '...'} and list of strings ['6201501', ...]
    codes = []
    for c in cnaes_sec:
        if isinstance(c, dict) and c.get("codigo"):
            codes.append(str(c["codigo"]))
        elif isinstance(c, str):
            codes.append(c)
    return ", ".join(codes)

def format_qsa(qsa: List) -> str:
    if not qsa:
        return ""
    return "; ".join([f"{s.get('nome_socio','')} ({s.get('qualificacao_socio','')})" for s in qsa if s.get("nome_socio")])

def save_results(companies: Dict):
    """Salva resultados em JSON, CSV e Excel."""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    companies_list = list(companies.values())
    
    # Separar B2B/B2C
    b2b = [c for c in companies_list if c.get("tipo_prospeccao") == "B2B"]
    b2c = [c for c in companies_list if c.get("tipo_prospeccao") == "B2C"]
    
    # JSON
    json_file = DATA_DIR / f"lelyon_prospects_real_{timestamp}.json"
    combined = {"b2b": b2b, "b2c": b2c}
    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(combined, f, ensure_ascii=False, indent=2)
    
    # CSV
    csv_file = DATA_DIR / f"lelyon_prospects_real_{timestamp}.csv"
    if companies_list:
        fieldnames = list(companies_list[0].keys())
        with open(csv_file, "w", encoding="utf-8-sig", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            for c in companies_list:
                row = {k: c.get(k, "") for k in fieldnames}
                row["lelyon_fit"] = json.dumps(c.get("lelyon_fit", {}), ensure_ascii=False)
                writer.writerow(row)
    
    # Excel
    excel_file = DATA_DIR / f"lelyon_prospects_real_{timestamp}.xlsx"
    with pd.ExcelWriter(excel_file, engine="openpyxl") as writer:
        df_all = pd.DataFrame(companies_list)
        df_all["lelyon_fit_str"] = df_all["lelyon_fit"].apply(lambda x: json.dumps(x, ensure_ascii=False) if isinstance(x, dict) else str(x))
        
        df_b2b = df_all[df_all["tipo_prospeccao"] == "B2B"]
        df_b2c = df_all[df_all["tipo_prospeccao"] == "B2C"]
        
        df_all.to_excel(writer, sheet_name="Todos", index=False)
        df_b2b.to_excel(writer, sheet_name="B2B", index=False)
        df_b2c.to_excel(writer, sheet_name="B2C", index=False)
        
        # Resumo
        enriched = len([c for c in companies_list if c.get("email") or c.get("telefone")])
        stats = [
            {"Métrica": "Total", "Valor": len(companies_list)},
            {"Métrica": "B2B", "Valor": len(b2b)},
            {"Métrica": "B2C", "Valor": len(b2c)},
            {"Métrica": "Com Email", "Valor": len([c for c in companies_list if c.get("email")])},
            {"Métrica": "Com Telefone", "Valor": len([c for c in companies_list if c.get("telefone")])},
            {"Métrica": "Enriquecidas", "Valor": enriched},
            {"Métrica": "Taxa Sucesso", "Valor": f"{enriched/len(companies_list)*100:.1f}%"},
        ]
        pd.DataFrame(stats).to_excel(writer, sheet_name="Resumo", index=False)
    
    # Atualizar symlinks latest
    for src, dst_name in [
        (json_file, "lelyon_prospects_latest.json"),
        (csv_file, "lelyon_prospects_latest.csv"),
        (excel_file, "lelyon_prospects_latest.xlsx"),
    ]:
        dst = DATA_DIR / dst_name
        if dst.exists():
            dst.unlink()
        dst.symlink_to(src.name)
    
    print(f"\n💾 Resultados salvos:")
    print(f"   📄 JSON: {json_file}")
    print(f"   📊 CSV:  {csv_file}")
    print(f"   📈 Excel: {excel_file}")

# ============ MAIN ============
def main():
    print("=" * 60)
    print("ENRIQUECIMENTO COM DADOS REAIS - CNPJ Finder API")
    print("=" * 60)
    print(f"Delay: {DELAY_BETWEEN_REQUESTS}s ({1/DELAY_BETWEEN_REQUESTS:.0f} req/s)")
    print()
    
    # Carregar empresas
    companies = load_companies()
    if not companies:
        print("❌ Nenhuma empresa para enriquecer")
        return
    
    # Carregar progresso
    progress = load_progress()
    processed = set(progress.get("processed_cnpjs", []))
    enriched_data = progress.get("enriched", {})
    
    pending = [cnpj for cnpj in companies.keys() if cnpj not in processed]
    print(f"📋 Total: {len(companies)} | Já processados: {len(processed)} | Pendentes: {len(pending)}")
    
    if not pending:
        print("✅ Todos já processados!")
        # Reconstruir lista final
        final_companies = {}
        for cnpj, c in companies.items():
            if cnpj in enriched_data:
                final_companies[cnpj] = enriched_data[cnpj]
            else:
                final_companies[cnpj] = c
        save_results(final_companies)
        return
    
    print(f"\n🚀 Iniciando enriquecimento de {len(pending)} empresas...")
    print("   (Ctrl+C para pausar - progresso salvo automaticamente)\n")
    
    enriched_count = 0
    failed_count = 0
    
    try:
        for i, cnpj in enumerate(pending):
            original = companies[cnpj]
            print(f"\n[{i+1}/{len(pending)}] {cnpj} - {original.get('razao_social', '')[:50]}")
            
            # Consultar API
            real = cnpj_finder_lookup(cnpj)
            
            if real:
                enriched = merge_company_data(original, real)
                enriched_count += 1
                print(f"   ✅ Enriquecido: email={bool(enriched.get('email'))}, tel={bool(enriched.get('telefone'))}")
            else:
                # Manter original se não encontrou
                enriched = original
                enriched["fonte_enriquecimento"] = "cnpj_finder_not_found"
                enriched["data_enriquecimento"] = datetime.now().isoformat()
                failed_count += 1
                print(f"   ⭕ Não encontrado na API")
            
            enriched_data[cnpj] = enriched
            processed.add(cnpj)
            
            # Salvar progresso a cada 50
            if (i + 1) % 50 == 0:
                save_progress({
                    "processed_cnpjs": list(processed),
                    "enriched": enriched_data,
                    "last_update": datetime.now().isoformat()
                })
            
            time.sleep(DELAY_BETWEEN_REQUESTS)
        
    except KeyboardInterrupt:
        print("\n\n⏸️ Interrompido pelo usuário. Progresso salvo.")
        save_progress({
            "processed_cnpjs": list(processed),
            "enriched": enriched_data,
            "last_update": datetime.now().isoformat()
        })
        return
    
    # Salvar resultados finais
    print(f"\n📦 Consolidando {len(enriched_data)} empresas...")
    final_companies = {}
    for cnpj, c in companies.items():
        if cnpj in enriched_data:
            final_companies[cnpj] = enriched_data[cnpj]
        else:
            final_companies[cnpj] = c
    
    save_results(final_companies)
    
    # Stats
    total_enriched = len([c for c in final_companies.values() if c.get("email") or c.get("telefone")])
    print(f"\n✅ CONCLUÍDO!")
    print(f"   Total: {len(final_companies)}")
    print(f"   Enriquecidas com sucesso: {enriched_count}")
    print(f"   Não encontradas: {failed_count}")
    print(f"   Com dados de contato: {total_enriched} ({total_enriched/len(final_companies)*100:.1f}%)")
    
    # Limpar progresso
    PROGRESS_FILE.unlink(missing_ok=True)

if __name__ == "__main__":
    main()