"""Baixa as séries listadas nas fichas de um tema e salva os arquivos brutos.

Uso:
    python3 scripts/coletar.py --tema economia
    python3 scripts/coletar.py --tema economia --so selic_meta,ipca_12m

Cada ficha diz de onde vem o dado (campo `api`) e qual é a consulta (campo `consulta`).
O arquivo bruto vai para dados/brutos/<fonte>/ com a data do dia no nome, e o download
é registrado no manifesto.csv da pasta.
"""

import argparse
import csv
import datetime as dt
import json
import sys
import time
import urllib.error

from comum import HOJE, METODOLOGIA, baixar, consultas, ler_fichas, salvar_bruto

INICIO_SGS = dt.date(2000, 1, 1)
SDMX_IMF = "https://api.imf.org/external/sdmx/3.0/data/dataflow/IMF.RES/{fluxo}/+/{chave}"
SDMX_PARAMS = "?dimensionAtObservation=TIME_PERIOD&attributes=dsd&measures=all"


def br(data):
    return data.strftime("%d/%m/%Y")


def baixar_json(url, tentativas=6, **kw):
    """Baixa e confere que a resposta é JSON. O firewall do Banco Central às vezes devolve,
    com status 200, uma página HTML de "Requisição inválida"; nesse caso, espera e repete."""
    for i in range(tentativas):
        conteudo = baixar(url, **kw)
        try:
            json.loads(conteudo)
            return conteudo
        except ValueError:
            time.sleep(5 * (i + 1))
    raise RuntimeError(f"resposta não é JSON depois de {tentativas} tentativas: {url}")


def coletar_sgs(codigo):
    """SGS do Banco Central. Séries diárias aceitam no máximo 10 anos por consulta (erro 406):
    nesse caso a consulta é fatiada em blocos de 5 anos e cada fatia vira um arquivo bruto."""
    base = f"https://api.bcb.gov.br/dados/serie/bcdata.sgs.{codigo}/dados?formato=json"
    url = f"{base}&dataInicial={br(INICIO_SGS)}&dataFinal={br(HOJE)}"
    try:
        conteudo = baixar_json(url, aceitar_erro=(406,))
        return [salvar_bruto("bcb_sgs", f"sgs{codigo}", conteudo, url, "json")]
    except urllib.error.HTTPError as e:
        if e.code != 406:
            raise
    arquivos = []
    inicio = INICIO_SGS
    while inicio <= HOJE:
        fim = min(dt.date(inicio.year + 4, 12, 31), HOJE)
        url = f"{base}&dataInicial={br(inicio)}&dataFinal={br(fim)}"
        conteudo = baixar_json(url)
        fatia = f"sgs{codigo}_p{inicio.year}-{fim.year}"
        arquivos.append(salvar_bruto("bcb_sgs", fatia, conteudo, url, "json"))
        inicio = dt.date(fim.year + 1, 1, 1)
        time.sleep(1)
    return arquivos


def nome_sidra(caminho):
    """/t/6784/n1/all/v/9814,93/p/all -> t6784_v9814-93"""
    partes = caminho.strip("/").split("/")
    tabela = partes[partes.index("t") + 1]
    variaveis = partes[partes.index("v") + 1].replace(",", "-") if "v" in partes else "all"
    return f"t{tabela}_v{variaveis}"


def coletar_sidra(caminho):
    url = "https://apisidra.ibge.gov.br/values" + caminho
    return [salvar_bruto("ibge_sidra", nome_sidra(caminho), baixar_json(url), url, "json")]


def paises_pares():
    with open(METODOLOGIA / "pares.csv", encoding="utf-8") as f:
        return sorted({linha["iso3"] for linha in csv.DictReader(f)})


def coletar_weo(indicador, agregados):
    paises = ["BRA"] + paises_pares() + [a for a in agregados if a]
    chave = f"{'+'.join(paises)}.{indicador}.A"
    url = SDMX_IMF.format(fluxo="WEO", chave=chave) + SDMX_PARAMS
    return [salvar_bruto("fmi_weo", f"weo_{indicador}", baixar_json(url), url, "json")]


def coletar_pcps(chave):
    url = SDMX_IMF.format(fluxo="PCPS", chave=chave) + SDMX_PARAMS
    return [salvar_bruto("fmi_pcps", "pcps_" + chave.replace(".", "-"), baixar_json(url), url, "json")]


def coletar_fred(serie):
    url = f"https://fred.stlouisfed.org/graph/fredgraph.csv?id={serie}"
    return [salvar_bruto("fred", f"fred_{serie}", baixar(url), url, "csv")]


def coletar_ficha(ficha):
    api = ficha.get("api")
    arquivos = []
    for consulta in consultas(ficha):
        if api == "bcb_sgs":
            arquivos += coletar_sgs(consulta)
        elif api == "sidra":
            arquivos += coletar_sidra(consulta)
        elif api == "fmi_weo":
            arquivos += coletar_weo(consulta, consultas(ficha, "agregados"))
        elif api == "fmi_pcps":
            arquivos += coletar_pcps(consulta)
        elif api == "fred":
            arquivos += coletar_fred(consulta)
        elif api in ("derivado", "manual"):
            pass
        elif api == "comexstat":
            print("  coletado por scripts/coletar_comex.py")
            break
        else:
            print(f"  ! api desconhecida '{api}' em {ficha['id']}", file=sys.stderr)
    for codigo in consultas(ficha, "consulta_extra_bcb_sgs"):
        arquivos += coletar_sgs(codigo)
    return arquivos


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tema", required=True)
    ap.add_argument("--so", help="ids de fichas separados por vírgula")
    args = ap.parse_args()

    fichas = ler_fichas(args.tema)
    if args.so:
        ids = set(args.so.split(","))
        fichas = [f for f in fichas if f["id"] in ids]

    falhas = []
    for ficha in fichas:
        print(f"{ficha['id']} ({ficha.get('api')})")
        try:
            for arquivo in coletar_ficha(ficha):
                print(f"  {arquivo.relative_to(arquivo.parents[2])}")
        except Exception as e:  # segue para a próxima ficha e reporta no fim
            print(f"  ! {e}", file=sys.stderr)
            falhas.append(ficha["id"])
    if falhas:
        print(f"\nFalharam: {', '.join(falhas)}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()