#!/usr/bin/env python3
"""
Enriquecimento de dados via Apollo.io
Lê empresas do JSON mais recente, busca contatos reais na Apollo
Atualiza email, telefone, LinkedIn, cargo, tech stack, etc.
"""

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

# ============ CONFIG ============
APOLLO_API_KEY = "btmgozMBQeNVzgE5WumGjQ"
APOLLO_BASE = "https://api.apollo.io/api/v1"

# Arquivos de entrada/saída
DATA_DIR = Path("/root/lelyon_prospecting/data")
INPUT_JSON = DATA_DIR / "lelyon_prospects_latest.json"
OUTPUT_JSON = DATA_DIR / f"lelyon_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
OUTPUT_CSV = DATA_DIR / f"lelyon_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
OUTPUT_EXCEL = DATA_DIR / f"lelyon_enriched_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
PROGRESS_FILE = DATA_DIR / "enrichment_apollo_progress.json"

# Rate limiting Apollo (plano básico: ~100 req/min)
DELAY_BETWEEN_REQUESTS = 1.0  # segundos
MAX_RETRIES = 3
TIMEOUT = 30

# Headers Apollo
headers = {
    "Authorization": f"Bearer {APOLLO_API_KEY}",
    "Content-Type": "application/json",
    "Cache-Control": "no-cache"
}

session = requests.Session()
session.headers.update(headers)


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


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


def load_companies() -> List[Dict]:
    """Carrega empresas do JSON mais recente."""
    if INPUT_JSON.exists():
        with open(INPUT_JSON, "r", encoding="utf-8") as f:
            data = json.load(f)
            # Compatibilidade: pode ser dict com b2b/b2c ou lista direta
            if isinstance(data, dict):
                companies = data.get("b2b", []) + data.get("b2c", [])
            else:
                companies = data
            print(f"✅ Carregadas {len(companies)} empresas de {INPUT_JSON}")
            return companies

    raise FileNotFoundError(f"Arquivo de entrada não encontrado: {INPUT_JSON}")


def clean_domain(url: str) -> Optional[str]:
    """Extrai domínio limpo de URL."""
    if not url:
        return None
    # Remove protocolo
    url = re.sub(r'^https?://', '', url)
    # Remove www
    url = re.sub(r'^www\.', '', url)
    # Pega apenas o domínio (antes do primeiro /)
    domain = url.split('/')[0].split(':')[0]
    # Valida formato básico
    if re.match(r'^[a-z0-9.-]+\.[a-z]{2,}$', domain):
        return domain.lower()
    return None


def guess_domain(company: Dict) -> Optional[str]:
    """Tenta adivinhar domínio da empresa baseado no nome."""
    razao = company.get("razao_social", "").lower()
    fantasia = company.get("nome_fantasia", "").lower()
    cnpj = company.get("cnpj", "").replace(".", "").replace("/", "").replace("-", "")
    
    nome_busca = fantasia or razao
    # Limpar: remove sufixos jurídicos, pontuação
    nome_busca = re.sub(r'\b(ltda|me|epp|sa|eireli|mei|s\.a\.)\b', '', nome_busca)
    nome_busca = re.sub(r'[^\w\s]', '', nome_busca)
    nome_busca = '_'.join(nome_busca.split()[:3])  # primeiras 3 palavras
    
    if len(nome_busca) < 4:
        return None
    
    # Tenta padrões comuns
    for tld in ['.com.br', '.com', '.net.br']:
        yield f"{nome_busca}{tld}"


def apollo_search_organization(domain: str) -> Optional[Dict]:
    """Busca organização na Apollo por domínio."""
    payload = {
        "q_organization_domains": domain,
        "page": 1,
        "per_page": 1
    }
    
    for attempt in range(MAX_RETRIES):
        try:
            resp = session.post(f"{APOLLO_BASE}/organizations/search", json=payload, timeout=TIMEOUT)
            
            if resp.status_code == 200:
                data = resp.json()
                orgs = data.get("organizations", [])
                if orgs:
                    return orgs[0]
                return None
            elif resp.status_code == 429:
                wait = 2 ** attempt * 10
                print(f"   ⏳ Rate limited (org search), aguardando {wait}s...")
                time.sleep(wait)
                continue
            elif resp.status_code == 401:
                print(f"   🔐 Erro de autenticação Apollo - verifique API key")
                return None
            elif resp.status_code == 402:
                print(f"   💳 Limite de créditos Apollo atingido")
                return None
            else:
                print(f"   ⚠️ Erro org search {resp.status_code}: {resp.text[:200]}")
                
        except requests.Timeout:
            print(f"   ⏱️ Timeout org search (tentativa {attempt+1})")
        except Exception as e:
            print(f"   ❌ Erro org search: {e}")
        
        if attempt < MAX_RETRIES - 1:
            time.sleep(2 ** attempt)
    
    return None


def apollo_search_people_by_organization(org_id: str, per_page: int = 10) -> List[Dict]:
    """Busca pessoas/contatos de uma organização na Apollo."""
    payload = {
        "organization_ids": [org_id],
        "page": 1,
        "per_page": per_page,
        "sort_by_field": "department",
        "sort_ascending": True
    }
    
    for attempt in range(MAX_RETRIES):
        try:
            resp = session.post(f"{APOLLO_BASE}/mixed_people/search", json=payload, timeout=TIMEOUT)
            
            if resp.status_code == 200:
                data = resp.json()
                return data.get("people", [])
            elif resp.status_code == 429:
                wait = 2 ** attempt * 10
                print(f"   ⏳ Rate limited (people search), aguardando {wait}s...")
                time.sleep(wait)
                continue
            elif resp.status_code == 401:
                print(f"   🔐 Erro de autenticação Apollo")
                return []
            elif resp.status_code == 402:
                print(f"   💳 Limite de créditos Apollo atingido")
                return []
            else:
                print(f"   ⚠️ Erro people search {resp.status_code}: {resp.text[:200]}")
                
        except requests.Timeout:
            print(f"   ⏱️ Timeout people search (tentativa {attempt+1})")
        except Exception as e:
            print(f"   ❌ Erro people search: {e}")
        
        if attempt < MAX_RETRIES - 1:
            time.sleep(2 ** attempt)
    
    return []


def apollo_enrich_contact(email: str) -> Optional[Dict]:
    """Enriquece um contato específico por email (revela telefone se disponível)."""
    payload = {
        "email": email,
        "reveal_personal_emails": True,
        "reveal_phone_number": True
    }
    
    try:
        resp = session.post(f"{APOLLO_BASE}/contacts/enrich", json=payload, timeout=TIMEOUT)
        
        if resp.status_code == 200:
            data = resp.json()
            if data.get("contact"):
                return data["contact"]
        elif resp.status_code == 429:
            time.sleep(10)
    except Exception:
        pass
    return None


def select_best_contact(people: List[Dict], company: Dict) -> Optional[Dict]:
    """Seleciona o melhor contato para prospecção B2B."""
    if not people:
        return None
    
    # Prioridade de cargos para decisão de compra
    priority_titles = [
        "ceo", "cto", "cio", "founder", "co-founder", "owner", "director",
        "gerente", "manager", "head", "lead", "coordenador", "supervisor",
        "sócio", "partner", "presidente", "vp", "vice president"
    ]
    
    # Departamentos relevantes
    priority_depts = ["engineering", "technology", "it", "marketing", "sales", 
                      "business development", "operations", "founder", "executive"]
    
    scored = []
    for p in people:
        score = 0
        title = (p.get("title") or "").lower()
        dept = (p.get("department") or "").lower()
        seniority = (p.get("seniority") or "").lower()
        
        # Título
        for i, pt in enumerate(priority_titles):
            if pt in title:
                score += (len(priority_titles) - i) * 10
                break
        
        # Departamento
        for i, pd in enumerate(priority_depts):
            if pd in dept:
                score += (len(priority_depts) - i) * 5
                break
        
        # Senioridade
        if seniority in ["c_level", "vp", "director", "head"]:
            score += 50
        elif seniority in ["manager", "senior"]:
            score += 20
        
        # Tem email corporativo?
        email = p.get("email", "")
        if email and not any(g in email.lower() for g in ["gmail", "hotmail", "outlook", "yahoo"]):
            score += 30
        
        # Tem telefone?
        if p.get("phone_numbers"):
            score += 20
        
        # LinkedIn?
        if p.get("linkedin_url"):
            score += 10
        
        scored.append((score, p))
    
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[0][1] if scored else None


def enrich_company(company: Dict) -> Dict:
    """Enriquece uma única empresa via Apollo."""
    cnpj = company.get("cnpj", "")
    nome = company.get("nome_fantasia") or company.get("razao_social", "")
    website = company.get("website", "")
    
    print(f"\n🔍 [{cnpj}] {nome}")
    
    enriched = company.copy()
    enriched.update({
        "fonte_enriquecimento": "apollo",
        "data_enriquecimento": datetime.now().isoformat(),
        "apollo_org_id": "",
        "apollo_contacts_found": 0,
    })
    
    # 1. Tentar buscar por domínio do website conhecido
    domain = clean_domain(website)
    org = None
    
    if domain and domain != "cnpj.biz":
        print(f"   🌐 Buscando por domínio: {domain}")
        org = apollo_search_organization(domain)
        time.sleep(DELAY_BETWEEN_REQUESTS)
    
    # 2. Se não achou, tentar domínios baseados no nome
    if not org:
        print(f"   🔎 Tentando domínios baseados no nome...")
        for guessed_domain in guess_domain(company):
            print(f"      Tentando: {guessed_domain}")
            org = apollo_search_organization(guessed_domain)
            time.sleep(DELAY_BETWEEN_REQUESTS)
            if org:
                domain = guessed_domain
                break
    
    # 3. Se ainda não achou, buscar por nome da empresa (mais lento, menos preciso)
    if not org:
        print(f"   📛 Buscando por nome da empresa...")
        payload = {
            "q_organization_name": nome[:100],
            "page": 1,
            "per_page": 1
        }
        try:
            resp = session.post(f"{APOLLO_BASE}/organizations/search", json=payload, timeout=TIMEOUT)
            if resp.status_code == 200:
                orgs = resp.json().get("organizations", [])
                if orgs:
                    org = orgs[0]
        except Exception as e:
            print(f"      ❌ Erro busca por nome: {e}")
        time.sleep(DELAY_BETWEEN_REQUESTS)
    
    if not org:
        print(f"   🚫 Organização não encontrada na Apollo")
        enriched["observacoes"] = "Apollo: Org não encontrada"
        return enriched
    
    # Organização encontrada
    org_id = org.get("id", "")
    enriched["apollo_org_id"] = org_id
    enriched["website"] = org.get("website_url") or enriched.get("website", "")
    enriched["website_status"] = "apollo_verified"
    
    print(f"   ✅ Org encontrada: {org.get('name')} (ID: {org_id})")
    print(f"      Domínio: {org.get('website_url', 'N/A')}")
    print(f"      Tecnologias: {', '.join(org.get('technologies', [])[:5]) or 'N/A'}")
    print(f"      Funcionários: {org.get('estimated_num_employees', 'N/A')}")
    print(f"      Receita est.: {org.get('annual_revenue', 'N/A')}")
    
    # Buscar contatos da organização
    people = apollo_search_people_by_organization(org_id, per_page=15)
    enriched["apollo_contacts_found"] = len(people)
    
    if not people:
        print(f"   👥 Nenhum contato encontrado")
        enriched["observacoes"] = f"Apollo: Org OK, 0 contatos"
        return enriched
    
    print(f"   👥 {len(people)} contatos encontrados")
    
    # Selecionar melhor contato
    best = select_best_contact(people, company)
    
    if best:
        print(f"   🎯 Melhor contato: {best.get('name')} - {best.get('title')} ({best.get('department')})")
        print(f"      Email: {best.get('email', 'N/A')}")
        print(f"      Telefone: {best.get('phone_numbers', ['N/A'])[0] if best.get('phone_numbers') else 'N/A'}")
        print(f"      LinkedIn: {best.get('linkedin_url', 'N/A')}")
        
        # Tentar enriquecer telefone se tiver email
        if best.get("email") and not best.get("phone_numbers"):
            print(f"      🔍 Enriquecendo contato para revelar telefone...")
            enriched_contact = apollo_enrich_contact(best["email"])
            if enriched_contact:
                best = enriched_contact
                print(f"      📞 Telefone revelado: {best.get('phone_numbers', ['N/A'])[0] if best.get('phone_numbers') else 'N/A'}")
        
        # Atualizar empresa com melhor contato
        enriched.update({
            "email": best.get("email", ""),
            "telefone": best.get("phone_numbers", [""])[0] if best.get("phone_numbers") else "",
            "whatsapp": "",  # Apollo não dá WhatsApp direto
            "contato_nome": best.get("name", ""),
            "contato_cargo": best.get("title", ""),
            "linkedin": best.get("linkedin_url", ""),
            "instagram": "",
            "facebook": "",
            "apollo_contact_id": best.get("id", ""),
            "apollo_department": best.get("department", ""),
            "apollo_seniority": best.get("seniority", ""),
        })
        
        # Adicionar tecnologias da empresa
        if org.get("technologies"):
            enriched["tech_stack"] = org["technologies"][:20]
        
        # Adicionar info da empresa
        enriched["apollo_employee_count"] = org.get("estimated_num_employees", 0)
        enriched["apollo_revenue"] = org.get("annual_revenue", 0)
        enriched["apollo_industry"] = org.get("industry", "")
        
        enriched["observacoes"] = f"Apollo: {len(people)} contatos, melhor: {best.get('title', 'N/A')}"
    else:
        enriched["observacoes"] = f"Apollo: {len(people)} contatos, nenhum adequado"
    
    return enriched


def save_results(companies: List[Dict]):
    """Salva resultados em JSON, CSV e Excel."""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    # JSON
    json_file = DATA_DIR / f"lelyon_enriched_{timestamp}.json"
    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(companies, f, ensure_ascii=False, indent=2)
    
    # CSV
    csv_file = DATA_DIR / f"lelyon_enriched_{timestamp}.csv"
    fieldnames = [
        "cnpj", "razao_social", "nome_fantasia", "cidade", "bairro", 
        "cnae", "cnae_desc", "situacao", "porte", "capital_social",
        "data_abertura", "tipo_prospeccao", "b2c_categoria",
        "best_fit_service", "fit_score", "lelyon_fit",
        "email", "telefone", "whatsapp", "website", "website_status",
        "contato_nome", "contato_cargo", "linkedin", "instagram", "facebook",
        "fonte_enriquecimento", "data_enriquecimento", "observacoes",
        # Apollo specific
        "apollo_org_id", "apollo_contact_id", "apollo_contacts_found",
        "apollo_department", "apollo_seniority", "apollo_employee_count",
        "apollo_revenue", "apollo_industry", "tech_stack"
    ]
    with open(csv_file, "w", encoding="utf-8-sig", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for c in companies:
            row = {k: c.get(k, "") for k in fieldnames}
            row["lelyon_fit"] = json.dumps(c.get("lelyon_fit", {}), ensure_ascii=False)
            row["tech_stack"] = json.dumps(c.get("tech_stack", []), ensure_ascii=False)
            writer.writerow(row)
    
    # Excel com abas
    excel_file = DATA_DIR / f"lelyon_enriched_{timestamp}.xlsx"
    with pd.ExcelWriter(excel_file, engine="openpyxl") as writer:
        df_all = pd.DataFrame(companies)
        # Ensure tech_stack column exists
        if "tech_stack" not in df_all.columns:
            df_all["tech_stack"] = [[] for _ in range(len(df_all))]
        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_all["tech_stack_str"] = df_all["tech_stack"].apply(lambda x: json.dumps(x, ensure_ascii=False) if isinstance(x, list) else str(x))
        
        df_b2b = df_all[df_all["tipo_prospeccao"] == "B2B"]
        df_b2c = df_all[df_all["tipo_prospeccao"] == "B2C"]
        
        # Write sheets only if they have data, otherwise write a placeholder row
        df_all.to_excel(writer, sheet_name="Todos_Enriquecidos", index=False)
        
        if len(df_b2b) > 0:
            df_b2b.to_excel(writer, sheet_name="B2B_Enriquecidos", index=False)
        else:
            pd.DataFrame([{"Aviso": "Nenhuma empresa B2B encontrada"}]).to_excel(writer, sheet_name="B2B_Enriquecidos", index=False)
            
        if len(df_b2c) > 0:
            df_b2c.to_excel(writer, sheet_name="B2C_Enriquecidos", index=False)
        else:
            pd.DataFrame([{"Aviso": "Nenhuma empresa B2C encontrada"}]).to_excel(writer, sheet_name="B2C_Enriquecidos", index=False)
        
        # Resumo enriquecimento
        enriched_count = len([c for c in companies if c.get("email") or c.get("telefone") or c.get("website")])
        stats = [
            {"Métrica": "Total Empresas", "Valor": len(companies)},
            {"Métrica": "Com Email", "Valor": len([c for c in companies if c.get("email")])},
            {"Métrica": "Com Telefone", "Valor": len([c for c in companies if c.get("telefone")])},
            {"Métrica": "Com WhatsApp", "Valor": len([c for c in companies if c.get("whatsapp")])},
            {"Métrica": "Com Website", "Valor": len([c for c in companies if c.get("website")])},
            {"Métrica": "Total Enriquecidas", "Valor": enriched_count},
            {"Métrica": "Taxa Sucesso", "Valor": f"{enriched_count/len(companies)*100:.1f}%"},
        ]
        pd.DataFrame(stats).to_excel(writer, sheet_name="Resumo_Enriquecimento", index=False)
    
    print(f"\n💾 Resultados salvos:")
    print(f"   📄 JSON: {json_file}")
    print(f"   📊 CSV:  {csv_file}")
    print(f"   📈 Excel: {excel_file}")
    
    # Atualizar symlink latest
    for src, dst_name in [
        (json_file, "lelyon_enriched_latest.json"),
        (csv_file, "lelyon_enriched_latest.csv"),
        (excel_file, "lelyon_enriched_latest.xlsx"),
    ]:
        dst = DATA_DIR / dst_name
        if dst.exists():
            dst.unlink()
        dst.symlink_to(src.name)


# ============ MAIN ============
def main():
    print("=" * 60)
    print("APOLLO.IO ENRIQUECIMENTO - LELYON Prospects")
    print("=" * 60)
    print(f"API Key: {APOLLO_API_KEY[:12]}...")
    print(f"Delay: {DELAY_BETWEEN_REQUESTS}s | Timeout: {TIMEOUT}s")
    print()
    
    # Carregar empresas
    companies = load_companies()
    if not companies:
        print("❌ Nenhuma empresa para enriquecer")
        return
    
    # Carregar progresso
    progress = load_progress()
    processed_cnpjs = set(progress.get("processed", []))
    enriched_data = progress.get("enriched", {})
    
    # Filtrar apenas não processadas
    pending = [c for c in companies if c.get("cnpj") not in processed_cnpjs]
    print(f"📋 Total: {len(companies)} | Já processadas: {len(processed_cnpjs)} | Pendentes: {len(pending)}")
    
    if not pending:
        print("✅ Todas já processadas!")
        # Reconstruir lista completa com dados enriquecidos
        final_companies = []
        for c in companies:
            cnpj = c.get("cnpj")
            if cnpj in enriched_data:
                final_companies.append(enriched_data[cnpj])
            else:
                final_companies.append(c)
        save_results(final_companies)
        return
    
    # Processar
    print(f"\n🚀 Iniciando enriquecimento Apollo de {len(pending)} empresas...")
    print("   (Ctrl+C para pausar - progresso é salvo automaticamente)\n")
    
    enriched_count = 0
    failed_count = 0
    
    try:
        for idx, company in enumerate(pending):
            cnpj = company.get("cnpj", "")
            
            # Enriquecer
            enriched = enrich_company(company)
            
            # Salvar progresso
            enriched_data[cnpj] = enriched
            processed_cnpjs.add(cnpj)
            progress["processed"] = list(processed_cnpjs)
            progress["enriched"] = enriched_data
            progress["last_index"] = idx
            save_progress(progress)
            
            if enriched.get("email") or enriched.get("telefone") or enriched.get("website"):
                enriched_count += 1
            else:
                failed_count += 1
            
            # Rate limiting
            if idx < len(pending) - 1:
                time.sleep(DELAY_BETWEEN_REQUESTS)
                
    except KeyboardInterrupt:
        print("\n\n⏸️ Interrompido pelo usuário. Progresso salvo.")
    
    # Reconstruir lista final completa
    print("\n📦 Consolidando resultados...")
    final_companies = []
    for c in companies:
        cnpj = c.get("cnpj")
        if cnpj in enriched_data:
            final_companies.append(enriched_data[cnpj])
        else:
            final_companies.append(c)
    
    save_results(final_companies)
    
    # Stats finais
    total_enriched = len([c for c in final_companies if c.get("email") or c.get("telefone") or c.get("website")])
    print(f"\n✅ CONCLUÍDO!")
    print(f"   Total processadas nesta sessão: {len(pending)}")
    print(f"   Enriquecidas com sucesso: {enriched_count}")
    print(f"   Sem dados encontrados: {failed_count}")
    print(f"   Total enriquecidas (geral): {total_enriched}/{len(final_companies)} ({total_enriched/len(final_companies)*100:.1f}%)")
    
    # Limpar progresso se completou tudo
    if len(processed_cnpjs) >= len(companies):
        PROGRESS_FILE.unlink(missing_ok=True)
        print("   🧹 Progresso limpo (todas processadas)")


if __name__ == "__main__":
    main()