Enriquecer uma planilha de CNPJs gastando o mínimo de créditos

"""Enriquece uma planilha de CNPJs gastando o mínimo de créditos.

    CNPJIA_KEY=... python enriquecer_csv.py                    # lê empresas.csv
    CNPJIA_KEY=... python enriquecer_csv.py minha_lista.csv --saida resultado.csv

O fluxo:
  1. valida o formato e o dígito verificador de cada CNPJ aqui mesmo: CNPJ
     inválido não vira chamada;
  2. consulta o perfil basic (cadastro, CNAE, endereço e os sinais has_phone,
     has_email...);
  3. consulta o perfil full só de quem tem has_phone ou has_email: é o full que
     traz telefones e e-mail.

Quando compensa: basic primeiro sai mais barato quando boa parte da lista não
tem contato; se quase todas têm, full direto evita a chamada dupla. O custo
vigente de cada perfil está em https://cnpj.ia.br/precos; o total impresso no
fim é a soma de meta.credits_charged das respostas, o que a API de fato cobrou.

A planilha precisa de uma coluna "cnpj". Exige chave (a demonstração sem chave
só aceita os CNPJs da allowlist dela).
"""

import argparse
import csv
import os
import re
import sys
import time
from pathlib import Path

import requests

API = "https://api.cnpj.ia.br"
ESPERA_MAXIMA = 60
PARAR_EM = {"invalid_api_key", "key_expired", "quota_exceeded", "payment_required", "insufficient_plan"}
PESOS_DV = [2, 3, 4, 5, 6, 7, 8, 9]
COLUNAS = ["cnpj", "razao_social", "situacao", "cnae_fiscal", "uf", "telefone", "email", "perfil", "creditos", "erro"]


def cnpj_valido(cnpj):
    """Formato e dígitos verificadores, inclusive do CNPJ alfanumérico.

    Cada caractere vale (código ASCII - 48): dígitos valem 0 a 9, letras de 17
    (A) a 42 (Z). Módulo 11 com pesos 2 a 9 da direita para a esquerda; resto
    menor que 2 dá 0, senão 11 - resto. Para CNPJ só com dígitos, é o cálculo
    de sempre.
    """
    if not re.fullmatch(r"[A-Z0-9]{12}[0-9]{2}", cnpj) or len(set(cnpj)) == 1:
        return False

    def dv(base):
        soma = sum((ord(c) - 48) * PESOS_DV[i % 8] for i, c in enumerate(reversed(base)))
        resto = soma % 11
        return 0 if resto < 2 else 11 - resto

    d1 = dv(cnpj[:12])
    d2 = dv(cnpj[:12] + str(d1))
    return cnpj[12:] == f"{d1}{d2}"


def get(sessao, url, params):
    """GET com timeout; em 429/503 espera o Retry-After e repete uma vez."""
    for tentativa in (1, 2):
        resp = sessao.get(url, params=params, timeout=(5, 30))
        if tentativa == 1 and resp.status_code in (429, 503):
            espera = resp.headers.get("Retry-After", "")
            if espera.isdigit() and int(espera) <= ESPERA_MAXIMA:
                print(f"  HTTP {resp.status_code}: aguardando {espera}s (Retry-After)", file=sys.stderr)
                time.sleep(int(espera))
                continue
        return resp


def consultar(sessao, cnpj, perfil):
    """Devolve (data, créditos cobrados, código de erro)."""
    try:
        resp = get(sessao, f"{API}/v1/cnpjs/{cnpj}", {"profile": perfil})
        corpo = resp.json()
    except (requests.RequestException, ValueError):
        return None, 0, "falha_de_rede"
    if resp.ok:
        return corpo["data"], corpo["meta"]["credits_charged"], None
    return None, 0, corpo.get("error", {}).get("code", f"http_{resp.status_code}")


def main():
    pasta = Path(__file__).parent
    parser = argparse.ArgumentParser(description="Enriquece uma planilha de CNPJs.")
    parser.add_argument("entrada", nargs="?", default=pasta / "empresas.csv")
    parser.add_argument("--saida", default="empresas_enriquecidas.csv")
    args = parser.parse_args()

    chave = os.environ.get("CNPJIA_KEY")
    if not chave:
        print("Defina CNPJIA_KEY.", file=sys.stderr)
        sys.exit(2)

    sessao = requests.Session()
    sessao.headers["Authorization"] = f"Bearer {chave}"

    with open(args.entrada, newline="", encoding="utf-8") as f:
        cnpjs = [linha["cnpj"] for linha in csv.DictReader(f)]

    linhas = []
    n_invalidos = n_basic = n_full = creditos = 0
    parada = None

    for bruto in cnpjs:
        # Tira a pontuação: a barra de "00.000.000/0001-91" quebra a rota da URL.
        cnpj = re.sub(r"[./-]", "", bruto.strip()).upper()
        linha = {"cnpj": cnpj}
        linhas.append(linha)

        if not cnpj_valido(cnpj):
            n_invalidos += 1
            linha["erro"] = "invalid_cnpj (validado localmente, sem chamada)"
            print(f"{cnpj}: inválido, sem chamada")
            continue

        basic, custo, erro = consultar(sessao, cnpj, "basic")
        n_basic += 1
        creditos += custo
        if erro:
            linha["erro"] = erro
            print(f"{cnpj}: {erro}")
            if erro in PARAR_EM:
                parada = erro
                break
            continue

        linha.update(
            razao_social=basic["razao_social"],
            situacao=basic["situacao_cadastral"]["descricao"],
            cnae_fiscal=basic["cnae_fiscal"]["codigo"],
            uf=basic["endereco"]["uf"],
            perfil="basic",
            creditos=custo,
        )

        if not (basic.get("has_phone") or basic.get("has_email")):
            print(f"{cnpj}: {basic['razao_social']} (basic, sem contato)")
            continue

        full, custo_full, erro = consultar(sessao, cnpj, "full")
        n_full += 1
        creditos += custo_full
        if erro:
            linha["erro"] = erro
            print(f"{cnpj}: {basic['razao_social']} (basic; full falhou: {erro})")
            if erro in PARAR_EM:
                parada = erro
                break
            continue

        # Com meta.suppressed, contatos podem vir omitidos: nada aqui supõe o campo.
        telefones = full.get("telefones") or []
        if telefones:
            linha["telefone"] = f"({telefones[0]['ddd']}) {telefones[0]['numero']}"
        linha["email"] = full.get("email") or ""
        linha["perfil"] = "basic+full"
        linha["creditos"] = custo + custo_full
        print(f"{cnpj}: {basic['razao_social']} (basic + full)")

    with open(args.saida, "w", newline="", encoding="utf-8") as f:
        escritor = csv.DictWriter(f, fieldnames=COLUNAS)
        escritor.writeheader()
        escritor.writerows(linhas)

    print(
        f"\n{len(cnpjs)} linhas: {n_invalidos} inválida(s) sem chamada, "
        f"{n_basic} consulta(s) basic, {n_full} consulta(s) full"
    )
    print(f"Créditos cobrados: {creditos} (soma de meta.credits_charged)")
    print(f"Saída: {args.saida}")

    if parada:
        print(f"Parou em {parada}: veja https://cnpj.ia.br/docs/erros-e-retentativas#{parada}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()