"""
Baixa e descompacta os arquivos mensais de Dados Abertos do CNPJ (RFB).

Fonte primária : WebDAV SERPRO+ da Receita Federal
                 https://arquivos.receitafederal.gov.br  (token público)
Fallback        : mirror Casa dos Dados
                 https://dados-abertos-rf-cnpj.casadosdados.com.br

Uso:
    python -m cnpj_monitor.download_rfb 2026-07
"""
import re
import sys
import time
import zipfile
from pathlib import Path

import requests
from requests.exceptions import RequestException

from .config import RFB_WEBDAV_URL, RFB_WEBDAV_TOKEN, RFB_BASE_URL, N_PARTES

CONNECT_TIMEOUT = 30    # segundos para estabelecer conexão
READ_TIMEOUT    = 300   # segundos entre chunks (arquivos grandes)
MAX_TENTATIVAS  = 4
BACKOFF_BASE    = 15    # segundos — dobra a cada tentativa (15, 30, 60, 120)

REPO_ROOT = Path(__file__).resolve().parents[2]
RAW_DIR = REPO_ROOT / "data" / "raw"
EXTRACTED_DIR = REPO_ROOT / "data" / "extracted"


def descobrir_pasta_fallback(mes_ano: str) -> str:
    """Descobre a pasta AAAA-MM-DD no mirror Casa dos Dados para o mês informado."""
    print(f"[fallback] buscando pasta para {mes_ano} no mirror...")
    r = requests.get(RFB_BASE_URL + "/", timeout=(CONNECT_TIMEOUT, READ_TIMEOUT))
    r.raise_for_status()
    candidatos = sorted(set(re.findall(rf'{re.escape(mes_ano)}-\d{{2}}', r.text)))
    if not candidatos:
        raise ValueError(
            f"Nenhuma pasta encontrada para {mes_ano} no mirror.\n"
            f"Verifique os meses disponíveis em {RFB_BASE_URL}/"
        )
    pasta = candidatos[-1]
    print(f"[fallback] pasta encontrada: {pasta}")
    return pasta


def _baixar(url: str, destino: Path, auth=None) -> None:
    """Download com retry/backoff exponencial. Lança RequestException após esgotar tentativas."""
    destino.parent.mkdir(parents=True, exist_ok=True)
    for tentativa in range(1, MAX_TENTATIVAS + 1):
        try:
            print(f"[download] {url} (tentativa {tentativa}/{MAX_TENTATIVAS})")
            with requests.get(
                url, stream=True,
                timeout=(CONNECT_TIMEOUT, READ_TIMEOUT),
                auth=auth,
            ) as r:
                r.raise_for_status()
                with open(destino, "wb") as f:
                    for chunk in r.iter_content(chunk_size=1024 * 1024):
                        f.write(chunk)
            return
        except RequestException as e:
            destino.unlink(missing_ok=True)
            if tentativa == MAX_TENTATIVAS:
                raise
            espera = BACKOFF_BASE * (2 ** (tentativa - 1))
            print(f"[retry] falha: {e}. Aguardando {espera}s...")
            time.sleep(espera)


def baixar_arquivo(nome: str, mes_ano: str, destino: Path, pasta_fallback: str) -> None:
    if destino.exists():
        print(f"[skip] já existe: {destino.name}")
        return

    # --- Primário: WebDAV RFB ---
    url_rfb = f"{RFB_WEBDAV_URL}/{mes_ano}/{nome}"
    try:
        _baixar(url_rfb, destino, auth=(RFB_WEBDAV_TOKEN, ""))
        return
    except RequestException as e:
        print(f"[fallback] RFB indisponível ({type(e).__name__}: {e}), tentando mirror...")

    # --- Fallback: mirror Casa dos Dados ---
    url_mirror = f"{RFB_BASE_URL}/{pasta_fallback}/{nome}"
    _baixar(url_mirror, destino)


def extrair(zip_path: Path, destino_dir: Path) -> None:
    print(f"[unzip] {zip_path.name}")
    destino_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path) as z:
        z.extractall(destino_dir)


def main(mes_ano: str):
    """mes_ano no formato AAAA-MM (ex: 2026-07)."""
    raw_mes = RAW_DIR / mes_ano
    extracted_mes = EXTRACTED_DIR / mes_ano

    arquivos = (
        [f"Empresas{i}.zip"         for i in range(N_PARTES)] +
        [f"Estabelecimentos{i}.zip" for i in range(N_PARTES)] +
        [f"Socios{i}.zip"           for i in range(N_PARTES)]
    )
    arquivos += ["Municipios.zip", "Cnaes.zip", "Simples.zip"]

    # Pasta do fallback descoberta uma vez só (evita N requisições ao mirror)
    pasta_fallback = None

    erros = []
    for nome in arquivos:
        destino = raw_mes / nome
        try:
            if pasta_fallback is None:
                try:
                    pasta_fallback = descobrir_pasta_fallback(mes_ano)
                except RequestException:
                    pasta_fallback = mes_ano  # melhor esforço se mirror também falhar
            baixar_arquivo(nome, mes_ano, destino, pasta_fallback)
            extrair(destino, extracted_mes)
        except RequestException as e:
            print(f"[erro] {nome}: {e}")
            erros.append(nome)

    print(f"\nArquivos extraídos em: {extracted_mes.resolve()}")
    if erros:
        print(f"[aviso] {len(erros)} arquivo(s) falharam: {', '.join(erros)}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Uso: python -m cnpj_monitor.download_rfb AAAA-MM   (ex: 2026-07)")
        sys.exit(1)
    main(sys.argv[1])