#!/usr/bin/env python3
"""Comércio do Brasil com os membros do BRICS, EUA, União Europeia e mundo (Comex Stat/MDIC).

Coleta na API oficial do Comex Stat e grava:

- dados/brutos/comexstat/<consulta>_<AAAA-MM-DD>.<ext>: resposta intocada, com linha no
  manifesto.csv da pasta (arquivo, consulta, data e hora, SHA-256, bytes);
- dados/tratados/politica_externa_comercio.csv: formato longo `indicador,data,valor,unidade,fonte`;
- o bloco "Extração" das fichas em fontes/fichas/politica_externa/.

No fim, imprime a participação média de China, BRICS-5 e EUA por mandato.

Uso:
    python3 scripts/coletar_comex.py              # coleta e trata
    python3 scripts/coletar_comex.py --so-tratar  # trata os brutos mais recentes, sem baixar

Só biblioteca padrão (decisão D-004). Um bruto nunca é sobrescrito: se a consulta já foi
baixada no mesmo dia, o arquivo do dia é reaproveitado e a API não é chamada de novo.

Comportamento da API, testado em 2026-10-05 (a documentação no próprio host responde 403):
- POST /general com corpo JSON. O GET com o filtro na URL é bloqueado (403).
- "period" funciona como "anos de A a B, meses de M1 a M2 em cada ano": {"from": "2025-01",
  "to": "2026-08"} devolve janeiro a agosto de 2025 e janeiro a agosto de 2026. Para anos
  completos, o período vai de janeiro do primeiro ano a dezembro do último.
- Há limite de requisições (HTTP 429, "tente novamente em 10 segundos"). O script espera
  INTERVALO segundos entre chamadas e repete com espera crescente.
- A composição da UE vem da própria API (bloco 22 aberto por país). A tabela PAIS_BLOCO do
  download em massa (balanca.mdic.gov.br) traz a mesma lista, mas o servidor não envia o
  certificado intermediário e o Python recusa a conexão; por isso ela não é baixada aqui.
"""

import argparse
import csv
import datetime as dt
import hashlib
import json
import re
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

RAIZ = Path(__file__).resolve().parent.parent
API = "https://api-comexstat.mdic.gov.br"
USER_AGENT = "Mozilla/5.0 (projeto comparacao-governos; coleta de dados publicos)"
FONTE = "MDIC/Comex Stat"
ANO_INICIAL = 2002  # D-002
BLOCO_UE = 22  # "União Europeia - UE" na tabela de blocos do Comex Stat
INTERVALO = 15  # segundos entre chamadas à API
MESES = "jan fev mar abr mai jun jul ago set out nov dez".split()

# Sufixo do id (ISO 3166 alfa-3, minúsculo) e código do país no Comex Stat.
PAISES = [
    ("chn", 160), ("rus", 676), ("ind", 361), ("zaf", 756),
    ("egy", 240), ("eth", 253), ("irn", 372), ("are", 244),
    ("sau", 53), ("idn", 365), ("usa", 249),
]
# Composição fixa em todos os anos (ver fichas e notas).
AGREGADOS = {
    "brics5": ["chn", "rus", "ind", "zaf"],
    "brics11": ["chn", "rus", "ind", "zaf", "egy", "eth", "irn", "are", "sau", "idn"],
}
PARCEIROS = [s for s, _ in PAISES] + ["ue"] + list(AGREGADOS) + ["mundo"]
US = "US$ (FOB, correntes)"
MEDIDAS = {
    "exportacoes_usd": US,
    "importacoes_usd": US,
    "corrente_usd": US,
    "saldo_usd": US,
    "participacao_corrente": "% da corrente de comércio total do Brasil",
}
RESUMO = [("chn", "China"), ("brics5", "BRICS-5"), ("usa", "EUA")]


def configurar(raiz):
    global BRUTOS, TRATADO, FICHAS, MANDATOS
    BRUTOS = raiz / "dados" / "brutos" / "comexstat"
    TRATADO = raiz / "dados" / "tratados" / "politica_externa_comercio.csv"
    FICHAS = raiz / "fontes" / "fichas" / "politica_externa"
    MANDATOS = raiz / "metodologia" / "mandatos.csv"


# ---------------------------------------------------------------- brutos

def pedir(url, corpo=None, tentativas=8):
    dados = json.dumps(corpo).encode("utf-8") if corpo is not None else None
    cabecalhos = {"User-Agent": USER_AGENT}
    if dados is not None:
        cabecalhos["Content-Type"] = "application/json"
    erro = None
    for i in range(tentativas):
        req = urllib.request.Request(url, data=dados, headers=cabecalhos,
                                     method="POST" if dados is not None else "GET")
        try:
            with urllib.request.urlopen(req, timeout=180) as r:
                return r.read()
        except urllib.error.HTTPError as e:
            if e.code not in (429, 500, 502, 503, 504):
                raise
            erro = e
        except (urllib.error.URLError, TimeoutError) as e:
            erro = e
        espera = min(15 * (i + 1), 90)
        print(f"    {erro}; nova tentativa em {espera} s", file=sys.stderr)
        time.sleep(espera)
    raise RuntimeError(f"falha em {url}: {erro}")


def brutos_de(nome, data=None):
    """Arquivos `<nome>_<AAAA-MM-DD>[-n].<ext>`, do mais antigo ao mais recente."""
    achados = []
    for p in BRUTOS.glob(f"{nome}_*"):
        m = re.fullmatch(re.escape(nome) + r"_(\d{4}-\d{2}-\d{2})(?:-(\d+))?\.\w+", p.name)
        if m and (data is None or m.group(1) == data):
            achados.append(((m.group(1), int(m.group(2) or 1)), p))
    return [p for _, p in sorted(achados)]


def bruto_mais_recente(nome):
    achados = brutos_de(nome)
    if not achados:
        raise SystemExit(f"Sem arquivo bruto de '{nome}' em {BRUTOS}. Rode sem --so-tratar.")
    return achados[-1]


def salvar_bruto(nome, ext, conteudo, consulta, hoje):
    BRUTOS.mkdir(parents=True, exist_ok=True)
    sha = hashlib.sha256(conteudo).hexdigest()
    caminho = BRUTOS / f"{nome}_{hoje}.{ext}"
    n = 1
    while caminho.exists():
        if hashlib.sha256(caminho.read_bytes()).hexdigest() == sha:
            return caminho
        n += 1
        caminho = BRUTOS / f"{nome}_{hoje}-{n}.{ext}"
    caminho.write_bytes(conteudo)
    manifesto = BRUTOS / "manifesto.csv"
    novo = not manifesto.exists()
    with open(manifesto, "a", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        if novo:
            w.writerow(["arquivo", "consulta_url", "data_extracao", "sha256", "bytes"])
        w.writerow([caminho.name, consulta, dt.datetime.now().isoformat(timespec="seconds"),
                    sha, len(conteudo)])
    return caminho


def obter(nome, ext, url, corpo, hoje, validar):
    """Baixa e grava um bruto, ou reaproveita o do dia. Devolve (caminho, chamou_api)."""
    do_dia = brutos_de(nome, hoje)
    if do_dia:
        print(f"  {nome}: já baixado hoje ({do_dia[-1].name})")
        return do_dia[-1], False
    conteudo = pedir(url, corpo)
    validar(conteudo)
    consulta = url if corpo is None else (
        f"POST {url} " + json.dumps(corpo, ensure_ascii=False, separators=(",", ":")))
    caminho = salvar_bruto(nome, ext, conteudo, consulta, hoje)
    print(f"  {nome}: {caminho.name} ({len(conteudo)} bytes)")
    return caminho, True


def validar_json(conteudo):
    d = json.loads(conteudo)
    if not isinstance(d, dict) or d.get("success") is False or "data" not in d:
        raise RuntimeError(f"resposta inesperada da API: {conteudo[:300]!r}")


def ler_lista(caminho):
    d = json.loads(caminho.read_text(encoding="utf-8"))["data"]
    return d["list"] if isinstance(d, dict) and "list" in d else d


def periodos_disponiveis(ano, mes):
    """Anos completos e, se o ano corrente não fechou, o recorte parcial jan–mes."""
    ano_completo = ano if mes == 12 else ano - 1
    periodos = {"anual": (f"{ANO_INICIAL}-01", f"{ano_completo}-12",
                          list(range(ANO_INICIAL, ano_completo + 1)))}
    if mes < 12:
        # Mesmo recorte de meses no ano corrente e no anterior, para comparar igual com igual.
        periodos["parcial"] = (f"{ano - 1}-01", f"{ano}-{mes:02d}", [ano - 1, ano])
    return periodos


def consultas_api(periodos):
    """Lista de (nome, sigla do fluxo, recorte, período, corpo da consulta)."""
    # Os códigos vão como texto, com três dígitos: a API ignora o inteiro 53 e aceita "053".
    ue = [{"filter": "economicBlock", "values": [str(BLOCO_UE)]}]
    recortes = {
        "parceiros": ([{"filter": "country", "values": [f"{c:03d}" for _, c in PAISES]}],
                      ["country"]),
        "ue": (ue, ["economicBlock"]),
        "total": ([], []),
        "ue_paises": (ue, ["country"]),  # só para documentar e conferir a composição da UE
    }
    lista = []
    for fluxo, sigla in (("export", "exp"), ("import", "imp")):
        for recorte, (filtros, detalhes) in recortes.items():
            for periodo, (de, ate, _) in periodos.items():
                if recorte == "ue_paises" and periodo != "anual":
                    continue
                corpo = {"flow": fluxo, "monthDetail": False, "period": {"from": de, "to": ate},
                         "filters": filtros, "details": detalhes, "metrics": ["metricFOB"]}
                lista.append((f"{sigla}_{recorte}_{periodo}", sigla, recorte, periodo, corpo))
    return lista


def coletar(hoje):
    print("Coleta no Comex Stat")
    chamou = False

    def api(nome, url, corpo=None):
        nonlocal chamou
        if chamou and not brutos_de(nome, hoje):
            time.sleep(INTERVALO)
        caminho, usou = obter(nome, "json", url, corpo, hoje, validar_json)
        chamou = chamou or usou
        return caminho

    att = json.loads(api("atualizacao", f"{API}/general/dates/updated").read_text("utf-8"))["data"]
    api("tabela_paises", f"{API}/tables/countries?language=pt")
    api("tabela_blocos", f"{API}/tables/economic-blocks?language=pt")
    periodos = periodos_disponiveis(int(att["year"]), int(att["monthNumber"]))
    for nome, _, _, _, corpo in consultas_api(periodos):
        api(nome, f"{API}/general?language=pt", corpo)


# ---------------------------------------------------------------- tratamento

def ler_manifesto():
    caminho = BRUTOS / "manifesto.csv"
    if not caminho.exists():
        return {}
    with open(caminho, encoding="utf-8") as f:
        return {linha["arquivo"]: linha for linha in csv.DictReader(f)}


def tratar():
    att = json.loads(bruto_mais_recente("atualizacao").read_text("utf-8"))["data"]
    ano, mes = int(att["year"]), int(att["monthNumber"])
    periodos = periodos_disponiveis(ano, mes)
    blocos = {int(x["id"]): x["text"] for x in ler_lista(bruto_mais_recente("tabela_blocos"))}
    if "União Europeia" not in blocos.get(BLOCO_UE, ""):
        raise SystemExit(f"Bloco {BLOCO_UE} não é a União Europeia: {blocos.get(BLOCO_UE)!r}")
    paises = {int(x["id"]): x["text"] for x in ler_lista(bruto_mais_recente("tabela_paises"))}
    sufixo_do_nome = {paises[c]: s for s, c in PAISES}

    valores, usados, faltas = {}, {}, []
    membros_ue, soma_ue = set(), {}
    for nome, sigla, recorte, periodo, _ in consultas_api(periodos):
        caminho = bruto_mais_recente(nome)
        usados[nome] = caminho
        for linha in ler_lista(caminho):
            if recorte == "ue_paises":
                if linha["country"] in sufixo_do_nome:
                    raise SystemExit(f"{linha['country']} aparece no bloco da UE")
                membros_ue.add(linha["country"])
                chave = (sigla, int(linha["year"]))
                soma_ue[chave] = soma_ue.get(chave, 0) + int(linha["metricFOB"])
                continue
            if recorte == "parceiros":
                parceiro = sufixo_do_nome[linha["country"]]
            elif recorte == "ue":
                if "União Europeia" not in linha.get("economicBlock", ""):
                    raise SystemExit(f"{caminho.name}: linha fora da UE: {linha}")
                parceiro = "ue"
            else:
                parceiro = "mundo"
            chave = (sigla, periodo, parceiro, int(linha["year"]))
            if chave in valores:
                raise SystemExit(f"{caminho.name}: valor repetido para {chave}")
            valores[chave] = int(linha["metricFOB"])
    for (sigla, a), soma in soma_ue.items():  # soma dos países do bloco = total do bloco
        if soma != valores.get((sigla, "anual", "ue", a)):
            raise SystemExit(f"UE {sigla} {a}: soma dos países {soma} difere do total do bloco")
    for periodo in periodos:  # parceiro ausente em todos os anos indica filtro errado
        for sigla in ("exp", "imp"):
            for p in [s for s, _ in PAISES] + ["ue"]:
                if not any(k[:3] == (sigla, periodo, p) for k in valores):
                    raise SystemExit(f"Nenhum dado de {p} em {sigla}_{periodo}: confira o filtro")

    linhas, participacao = [], {}
    base = {}  # (periodo, parceiro, ano) -> (exportações, importações)
    for periodo, (_, _, anos) in periodos.items():
        for a in anos:
            for p in [s for s, _ in PAISES] + ["ue", "mundo"]:
                par = []
                for sigla in ("exp", "imp"):
                    v = valores.get((sigla, periodo, p, a))
                    if v is None:
                        if p == "mundo":
                            raise SystemExit(f"Falta o total do Brasil: {sigla} {periodo} {a}")
                        faltas.append(f"{sigla}_{p}_{periodo}_{a}")
                        v = 0  # a API omite a linha quando não há registro de comércio
                    par.append(v)
                base[(periodo, p, a)] = tuple(par)
            for ag, membros in AGREGADOS.items():
                base[(periodo, ag, a)] = tuple(sum(base[(periodo, m, a)][i] for m in membros)
                                               for i in (0, 1))

    for medida, unidade in MEDIDAS.items():
        for periodo, (_, _, anos) in periodos.items():
            sufixo = "" if periodo == "anual" else "_parcial"
            for p in PARCEIROS:
                if medida == "participacao_corrente" and p == "mundo":
                    continue
                for a in anos:
                    exp, imp = base[(periodo, p, a)]
                    if medida == "exportacoes_usd":
                        valor = str(exp)
                    elif medida == "importacoes_usd":
                        valor = str(imp)
                    elif medida == "corrente_usd":
                        valor = str(exp + imp)
                    elif medida == "saldo_usd":
                        valor = str(exp - imp)
                    else:
                        x = 100 * (exp + imp) / sum(base[(periodo, "mundo", a)])
                        participacao[(periodo, p, a)] = x
                        valor = f"{x:.4f}"
                    linhas.append([f"{medida}{sufixo}__{p}", str(a), valor, unidade, FONTE])

    TRATADO.parent.mkdir(parents=True, exist_ok=True)
    with open(TRATADO, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["indicador", "data", "valor", "unidade", "fonte"])
        w.writerows(linhas)
    print(f"\n{TRATADO.relative_to(TRATADO.parents[2])}: {len(linhas)} linhas")
    if faltas:
        print("Sem registro na API (gravado como 0):", ", ".join(faltas))
    print(f"União Europeia (bloco {BLOCO_UE}): {len(membros_ue)} países e territórios com comércio "
          f"no período, soma conferida com o total do bloco:", "; ".join(sorted(membros_ue)))

    atualizar_fichas(usados, periodos, att)
    imprimir_resumo(participacao, periodos)


def rotulo(periodo, a, mes):
    return str(a) if periodo == "anual" else f"{a} (jan–{MESES[mes - 1]})"


def atualizar_fichas(usados, periodos, att):
    manifesto = ler_manifesto()
    mes = int(att["monthNumber"])
    for medida in MEDIDAS:
        ficha = FICHAS / f"{medida}.md"
        if not ficha.exists():
            continue
        siglas = {"exportacoes_usd": ["exp"], "importacoes_usd": ["imp"]}.get(medida, ["exp", "imp"])
        n_series = len(PARCEIROS) - (medida == "participacao_corrente")
        tabela = ["| Série | Arquivos brutos | Extraído em | Observações | Primeira | Última |",
                  "|---|---|---|---|---|---|"]
        for periodo, (_, _, anos) in periodos.items():
            sufixo = "" if periodo == "anual" else "_parcial"
            arquivos = [usados[f"{s}_{r}_{periodo}"] for s in siglas
                        for r in ("parceiros", "ue", "total")]
            quando = max((manifesto.get(a.name, {}).get("data_extracao", "") for a in arquivos),
                         default="").replace("T", " ")[:16]
            nomes = "<br>".join(f"`dados/brutos/comexstat/{a.name}`" for a in arquivos)
            tabela.append(f"| `{medida}{sufixo}__*` ({n_series} séries) | {nomes} | {quando} | "
                          f"{len(anos)} por série | {rotulo(periodo, anos[0], mes)} | "
                          f"{rotulo(periodo, anos[-1], mes)} |")
        nota = (f"Base do Comex Stat atualizada em {att['updated']}, com dados até "
                f"{MESES[mes - 1]}/{att['year']} "
                f"(`dados/brutos/comexstat/{bruto_mais_recente('atualizacao').name}`). "
                f"Composição da UE conferida em "
                f"`{usados['exp_ue_paises_anual'].name}` e `{usados['imp_ue_paises_anual'].name}`.")
        bloco = "\n".join(tabela) + "\n\n" + nota
        texto = ficha.read_text(encoding="utf-8")
        novo = re.sub(r"<!-- extracao:inicio -->.*?<!-- extracao:fim -->",
                      lambda m: f"<!-- extracao:inicio -->\n{bloco}\n<!-- extracao:fim -->",
                      texto, flags=re.S)
        if novo != texto:
            ficha.write_text(novo, encoding="utf-8")


def anos_do_mandato(m):
    """Anos civis inteiros dentro do mandato (D-006)."""
    inicio, fim = dt.date.fromisoformat(m["inicio"]), dt.date.fromisoformat(m["fim"])
    return [a for a in range(inicio.year, fim.year + 1)
            if inicio <= dt.date(a, 1, 1) and dt.date(a, 12, 31) <= fim]


def imprimir_resumo(participacao, periodos):
    with open(MANDATOS, encoding="utf-8") as f:
        mandatos = {m["id"]: m for m in csv.DictReader(f)}
    disponiveis = set(periodos["anual"][2])

    def num(x):
        return f"{x:.1f}".replace(".", ",")

    def faixa(anos):
        seguidos = anos == list(range(anos[0], anos[-1] + 1))
        return f"{anos[0]}–{anos[-1]}" if seguidos else ", ".join(map(str, anos))

    def linha(rotulo_, anos, valores):
        print(f"| {rotulo_} | {anos} | " + " | ".join(num(v) for v in valores) + " |")

    print("\nParticipação na corrente de comércio do Brasil, % (média aritmética dos anos)\n")
    print("| Mandato | Anos | " + " | ".join(n for _, n in RESUMO) + " |")
    print("|---|---|" + "---:|" * len(RESUMO))
    for mid in ("lula1", "lula2", "bolsonaro", "lula3"):
        m = mandatos[mid]
        anos = [a for a in anos_do_mandato(m) if a in disponiveis]
        variantes = [(m["nome"] + (" (parcial)" if anos[-1] < anos_do_mandato(m)[-1] else ""), anos)]
        if mid == "bolsonaro":
            variantes.append((f"{m['nome']}, sem 2020–21", [a for a in anos if a not in (2020, 2021)]))
        for nome, lista in variantes:
            linha(nome, faixa(lista), [sum(participacao[("anual", p, a)] for a in lista) / len(lista)
                                       for p, _ in RESUMO])
    if "parcial" in periodos:
        _, ate, anos = periodos["parcial"]
        for a in anos:
            linha(f"Acumulado {rotulo('parcial', a, int(ate[-2:]))}", a,
                  [participacao[("parcial", p, a)] for p, _ in RESUMO])


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--so-tratar", action="store_true",
                    help="não baixa nada; trata os brutos mais recentes")
    ap.add_argument("--raiz", type=Path, default=RAIZ,
                    help="raiz do repositório (padrão: a pasta acima de scripts/)")
    args = ap.parse_args()
    configurar(args.raiz.resolve())
    if not args.so_tratar:
        coletar(dt.date.today().isoformat())
    tratar()


if __name__ == "__main__":
    main()