"""Preenche `url_arquivada` com cópias do Wayback Machine (Internet Archive).

Uso:
    python3 scripts/arquivar.py            # só procura cópias que já existem; não envia nada
    python3 scripts/arquivar.py --salvar   # pede ao Wayback para arquivar os links sem cópia

Lê eventos/*.csv (coluna fonte_url) e fontes/registro/*.csv (coluna url). Páginas de
governo mudam ou somem; a cópia arquivada garante que o link continue verificável.
Dados de séries não passam por aqui: o arquivo bruto com hash, em dados/brutos/, já é
o arquivo.
"""

import argparse
import csv
import json
import re
import time
import urllib.error
import urllib.parse
import urllib.request

from comum import RAIZ, USER_AGENT

ARQUIVOS = [("eventos", "fonte_url"), ("fontes/registro", "url")]


class Pausa(Exception):
    """O Internet Archive respondeu 429: insistir só prolonga o bloqueio."""


def pedir(url, timeout=60, tentativas=3):
    """Abre a URL. Falha passageira (5xx) tenta de novo; 429 vira Pausa, sem novas tentativas."""
    for i in range(tentativas):
        req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        try:
            return urllib.request.urlopen(req, timeout=timeout)
        except urllib.error.HTTPError as e:
            if e.code == 429:
                raise Pausa(url) from e
            if e.code not in (500, 502, 503, 504) or i == tentativas - 1:
                raise
            time.sleep(15 * (i + 1))


def copia_no_indice(url):
    """Última cópia com status 200 no índice CDX do Wayback, que acerta mais que a API
    `available` (o índice ignora maiúsculas, http/https e www). None se não houver; ValueError
    se o índice estiver fora do ar (página HTML no lugar do JSON)."""
    consulta = ("https://web.archive.org/cdx/search/cdx?output=json&limit=-1&filter=statuscode:200"
                "&fl=timestamp,original&url=" + urllib.parse.quote(url, safe=""))
    for i in range(3):
        with pedir(consulta) as r:
            corpo = r.read().decode("utf-8", "replace")
        if corpo.lstrip().startswith("["):
            linhas = json.loads(corpo)
            if len(linhas) < 2:
                return None
            data, original = linhas[-1]
            return f"https://web.archive.org/web/{data}/{original}"
        time.sleep(20 * (i + 1))  # "Temporarily Offline"
    raise ValueError("índice CDX fora do ar")


def copia_existente(url):
    try:
        return copia_no_indice(url)
    except ValueError:
        pass
    consulta = "https://archive.org/wayback/available?url=" + urllib.parse.quote(url, safe="")
    with pedir(consulta) as r:
        dados = json.loads(r.read().decode("utf-8"))
    mais_proxima = dados.get("archived_snapshots", {}).get("closest") or {}
    if mais_proxima.get("available"):
        return mais_proxima["url"].replace("http://", "https://", 1)
    return None


def salvar_agora(url):
    """Save Page Now anônimo. Devolve a URL da cópia ou None."""
    with pedir("https://web.archive.org/save/" + url, timeout=180) as r:
        final = r.geturl()
        local = r.headers.get("Content-Location")
    if local:
        return "https://web.archive.org" + local
    if "/web/" in final:
        return final
    return None


def separar_copia(linha, coluna):
    """Se a fonte foi registrada direto pela cópia do Wayback, devolve o link original
    à coluna da fonte e guarda a cópia em url_arquivada. A URL de uma cópia tem a forma
    https://web.archive.org/web/<data>/<url original>."""
    url = (linha.get(coluna) or "").strip()
    m = re.match(r"https?://web\.archive\.org/web/\d+[a-z_]*/(https?://.+)", url)
    if m and not linha.get("url_arquivada"):
        linha["url_arquivada"] = url.replace("http://", "https://", 1)
        linha[coluna] = m.group(1)
        return True
    return False


def gravar(caminho, campos, linhas):
    with open(caminho, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=campos)
        w.writeheader()
        w.writerows(linhas)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--salvar", action="store_true", help="arquivar no Wayback os links sem cópia")
    ap.add_argument("--so", help="caminhos relativos separados por vírgula, ex.: eventos/economia.csv")
    ap.add_argument("--pausa", type=int, default=600, help="segundos de espera depois de um 429 (padrão 600)")
    args = ap.parse_args()
    somente = {RAIZ / p.strip() for p in args.so.split(",")} if args.so else None

    cache, contagem = {}, {"ja_tinha": 0, "separada": 0, "achada": 0, "salva": 0, "sem_copia": 0, "erro": 0}
    pausas, parar = 0, False
    for pasta, coluna in ARQUIVOS:
        for caminho in sorted((RAIZ / pasta).glob("*.csv")):
            if somente and caminho not in somente:
                continue
            with open(caminho, encoding="utf-8", newline="") as f:
                leitor = csv.DictReader(f)
                campos, linhas = leitor.fieldnames, list(leitor)
            if not campos or coluna not in campos or "url_arquivada" not in campos:
                continue
            mudou, pendentes = False, 0
            for linha in linhas:
                if parar:
                    break
                if separar_copia(linha, coluna):
                    contagem["separada"] += 1
                    mudou = True
                url = (linha.get(coluna) or "").strip()
                if not url.startswith("http"):
                    continue
                if linha.get("url_arquivada"):
                    contagem["ja_tinha"] += 1
                    continue
                if url not in cache:
                    try:
                        copia = copia_existente(url)
                        if copia:
                            contagem["achada"] += 1
                        elif args.salvar:
                            copia = salvar_agora(url)
                            contagem["salva" if copia else "sem_copia"] += 1
                            time.sleep(8)  # o Save Page Now anônimo limita a frequência
                        else:
                            contagem["sem_copia"] += 1
                        pausas = 0
                    except Pausa:
                        # 429: uma pausa longa; se voltar logo, grava o que tem e para
                        pausas += 1
                        if pausas >= 2:
                            print("  O Internet Archive pediu pausa de novo (429). Progresso gravado; rode mais tarde.", flush=True)
                            parar = True
                            break
                        print(f"  429 em {url}: pausa de {args.pausa} s", flush=True)
                        time.sleep(args.pausa)
                        continue
                    except (urllib.error.URLError, TimeoutError, ValueError) as e:
                        print(f"  ! {url}: {e}", flush=True)
                        contagem["erro"] += 1
                        copia = None
                    cache[url] = copia
                    time.sleep(2)
                if cache[url]:
                    linha["url_arquivada"] = cache[url]
                    mudou = True
                    pendentes += 1
                    if pendentes >= 20:  # grava o progresso aos poucos
                        gravar(caminho, campos, linhas)
                        pendentes = 0
            if mudou:
                gravar(caminho, campos, linhas)
                print(f"{caminho.relative_to(RAIZ)} atualizado", flush=True)
            if parar:
                break
        if parar:
            break
    print(", ".join(f"{k}: {v}" for k, v in contagem.items()))
    if contagem["sem_copia"] and not args.salvar:
        print("Para arquivar os links sem cópia: python3 scripts/arquivar.py --salvar")


if __name__ == "__main__":
    main()