"""Etapa 1: baixa os dados públicos usados na Fase 1.

- SIM/DATASUS: declarações de óbito de 2020 a 2024, por UF
- IBGE/SIDRA: população do Censo 2022 por município, sexo e faixa etária (tabela 9514)
- IBGE: lista de municípios e malha municipal (para o mapa)

Arquivos já baixados são reaproveitados. Uso: python baixar.py
"""
import json
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

import pandas as pd
import requests

from config import ANOS, BRUTOS, CENSO_IDADE, COD_UF, FTP_SIM, UFS


def baixar_url(url, destino, tentativas=4):
    if destino.exists() and destino.stat().st_size > 0:
        return "já existe"
    tmp = destino.with_suffix(destino.suffix + ".part")
    for i in range(tentativas):
        try:
            if url.startswith("ftp://"):
                urllib.request.urlretrieve(url, tmp)
            else:
                r = requests.get(url, timeout=300)
                r.raise_for_status()
                tmp.write_bytes(r.content)
            tmp.rename(destino)
            return "ok"
        except Exception as e:  # rede do DATASUS oscila; tenta de novo
            if i == tentativas - 1:
                raise RuntimeError(f"falhou {url}: {e}")
            time.sleep(5 * (i + 1))


def baixar_sim():
    pasta = BRUTOS / "sim"
    pasta.mkdir(parents=True, exist_ok=True)
    tarefas = [(uf, ano) for ano in ANOS for uf in UFS]
    with ThreadPoolExecutor(max_workers=4) as ex:
        futuros = {
            ex.submit(baixar_url, FTP_SIM.format(uf=uf, ano=ano), pasta / f"DO{uf}{ano}.dbc"): (uf, ano)
            for uf, ano in tarefas
        }
        for n, f in enumerate(as_completed(futuros), 1):
            uf, ano = futuros[f]
            print(f"  SIM {uf} {ano}: {f.result()}  ({n}/{len(tarefas)})", flush=True)


def baixar_populacao():
    destino = BRUTOS / "populacao_censo2022.csv"
    if destino.exists():
        print("  População: já existe")
        return
    idades = ",".join(str(c) for c in CENSO_IDADE)
    linhas = []
    for uf, cod in COD_UF.items():
        url = (f"https://apisidra.ibge.gov.br/values/t/9514/n6/in%20n3%20{cod}/v/93/p/2022"
               f"/c2/4,5/c287/{idades}/c286/113635/f/c")
        for i in range(4):
            try:
                r = requests.get(url, timeout=300)
                r.raise_for_status()
                dados = r.json()
                break
            except Exception:
                if i == 3:
                    raise
                time.sleep(10 * (i + 1))
        cab = dados[0]
        for d in dados[1:]:
            linhas.append({
                "cod_ibge": int(d["D1C"]),
                "sexo": 1 if d["D4C"] == "4" else 2,
                "idade_censo": int(d["D5C"]),
                "pop": int(d["V"]) if d["V"].isdigit() else 0,
            })
        print(f"  População {uf}: {len(dados) - 1} valores", flush=True)
    df = pd.DataFrame(linhas)
    df["faixa"] = df["idade_censo"].map(CENSO_IDADE)
    df = df.groupby(["cod_ibge", "sexo", "faixa"], as_index=False)["pop"].sum()
    df.to_csv(destino, index=False)
    print(f"  População: {df['cod_ibge'].nunique()} municípios, {df['pop'].sum():,} pessoas")


def baixar_municipios():
    destino = BRUTOS / "municipios.json"
    baixar_url("https://servicodados.ibge.gov.br/api/v1/localidades/municipios", destino)
    print("  Lista de municípios: ok")
    malha = BRUTOS / "malha_municipios.geojson"
    baixar_url("https://servicodados.ibge.gov.br/api/v3/malhas/paises/BR"
               "?formato=application/vnd.geo+json&qualidade=minima&intrarregiao=municipio", malha)
    malha_uf = BRUTOS / "malha_ufs.geojson"
    baixar_url("https://servicodados.ibge.gov.br/api/v3/malhas/paises/BR"
               "?formato=application/vnd.geo+json&qualidade=minima&intrarregiao=UF", malha_uf)
    print("  Malhas: ok")


if __name__ == "__main__":
    BRUTOS.mkdir(parents=True, exist_ok=True)
    etapas = sys.argv[1:] or ["municipios", "populacao", "sim"]
    if "municipios" in etapas:
        print("Municípios e malha (IBGE)")
        baixar_municipios()
    if "populacao" in etapas:
        print("População do Censo 2022 (SIDRA)")
        baixar_populacao()
    if "sim" in etapas:
        print("Óbitos (SIM/DATASUS)")
        baixar_sim()
    print("Pronto.")