#!/usr/bin/env python3
"""
subidor-paralelo.py -- substitui o subidor.py para drenar a fila de ~/Downloads.

Por que existe (medido 2026-09-12): o subidor.py sobe um arquivo por vez e, a
cada arquivo, espera 1,5 s de "estabilidade" e dispara tres `aws s3 cp`
(arquivo, metadata.json, manifest). Cada um paga a partida do aws-cli: ~5 s por
arquivo, seja thumbnail de 74 KB ou preservacao de 40 MB. Resultado: ~1,3 GB/h
num link de 69 Mbps (~31 GB/h). O gargalo era o processo, nao o link.

O que muda:
  - upload em paralelo via boto3 (um cliente, N threads, sem partida de CLI),
    conferindo o tamanho no bucket antes de apagar o local
  - metadata.json + manifest-sha256.txt uma vez por item ao fim de cada lote
  - copia "(N)" do Chrome identica (mesmo handle, tipo, sha256 e nome-base de
    algo ja enviado) NAO sobe (achado 2.20)
  - arquivo que nao e imagem vai pra pipeline/rejeitados/ e ganha linha na
    tabela `rejeitados`, em vez de ser apagado sem rastro (achado 2.19)
  - manifest do item deixa de fora `derivado-thumb256`: esses objetos moram em
    derivados/thumb256/, nao dentro de itens/<handle>/

O que NAO muda: chaves no bucket, ledger `enviados`, validacao por magic bytes,
conteudo de metadata.json.

Uso:
    set -a; . .s3env; set +a
    python3 subidor-paralelo.py --catalogo catalogo.jsonl              # drena a fila
    python3 subidor-paralelo.py --limpar-duplicatas                    # apaga "(N)" identicas (bucket + ledger)
    python3 subidor-paralelo.py --catalogo catalogo.jsonl --so-manifestos  # regrava metadata+manifest de todo item
"""
import argparse, hashlib, json, mimetypes, os, re, shutil, sqlite3, sys, time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import boto3
from boto3.s3.transfer import TransferConfig
from botocore.config import Config

S3_ENDPOINT = os.environ.get("S3_ENDPOINT", "https://hel1.your-objectstorage.com")
BUCKET = os.environ.get("ARCHIVE_BUCKET", "brasiliana")
PREFIXO = "BRAS__"
ROOT = Path(__file__).resolve().parent
DB = ROOT / "subidor.sqlite"
REJ_DIR = ROOT / "rejeitados"
SUFIXO_N = re.compile(r"^(.*) \((\d+)\)(\.[^.]+)$")
TIPOS_DO_ITEM = ("preservacao", "thumbnail", "legado")  # o que mora em itens/<handle>/
TRANSF = TransferConfig(multipart_threshold=64 * 2**20, multipart_chunksize=16 * 2**20,
                        max_concurrency=4)


def agora():
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def log(msg):
    print(f"{time.strftime('%H:%M:%S')} {msg}", flush=True)


def base(nome):
    m = SUFIXO_N.match(nome)
    return m.group(1) + m.group(3) if m else nome


def cliente(threads):
    return boto3.client(
        "s3", endpoint_url=S3_ENDPOINT,
        config=Config(max_pool_connections=threads * 5 + 8,
                      retries={"max_attempts": 8, "mode": "adaptive"},
                      request_checksum_calculation="when_required",
                      response_checksum_validation="when_required"))


def db():
    con = sqlite3.connect(DB, timeout=60)
    con.executescript("""
    CREATE TABLE IF NOT EXISTS descartados (
      handle TEXT, tipo TEXT, nome TEXT, bytes INTEGER, sha256 TEXT,
      chave_no_bucket TEXT, quando TEXT, PRIMARY KEY (handle, tipo, nome, sha256));
    CREATE TABLE IF NOT EXISTS rejeitados (
      handle TEXT, tipo TEXT, nome TEXT, bytes INTEGER, sha256 TEXT,
      cabecalho TEXT, quando TEXT, PRIMARY KEY (handle, tipo, nome, sha256));""")
    return con


def carrega_catalogo(caminho):
    d = {}
    if caminho:
        for linha in open(caminho, encoding="utf-8"):
            if linha.strip():
                r = json.loads(linha)
                d[r["handle"]] = r
    return d


def interpreta(p):
    """BRAS__20.500.12156.1_10377__[THUMB__]nome -> (handle, pasta, tipo, nome)"""
    resto = p.name[len(PREFIXO):]
    if "__" not in resto:
        return None
    pasta, nome = resto.split("__", 1)
    m = re.fullmatch(r"(.+)_(\d+)", pasta)
    if not m:
        return None
    tipo = "preservacao"
    if nome.startswith("THUMB__"):
        nome, tipo = nome[len("THUMB__"):], "thumbnail"
    return f"{m.group(1)}/{m.group(2)}", pasta, tipo, nome


# assinaturas de imagem aceitas. Achado 2026-09-12: parte das "preservacao" com
# extensao .jpg e TIFF de verdade (II*\0 / MM\0*, ~8 MB, faixa 145xx); a regra
# antiga (so JPEG/PNG) do subidor.py apagava esses arquivos como se fossem o
# bitstream de licenca. O que se quer barrar e texto, nao formato de imagem.
ASSINATURAS = (
    (b"\xff\xd8", "jpeg"), (b"\x89PNG\r\n\x1a\n", "png"),
    (b"II*\x00", "tiff"), (b"MM\x00*", "tiff"), (b"GIF87a", "gif"), (b"GIF89a", "gif"),
    (b"\x00\x00\x00\x0cjP  ", "jp2"), (b"\xff\x4f\xff\x51", "j2k"), (b"BM", "bmp"),
)


def formato(cab):
    for sig, nome in ASSINATURAS:
        if cab.startswith(sig):
            return nome
    if cab[:4] == b"RIFF" and cab[8:12] == b"WEBP":
        return "webp"
    return None


def examina(p):
    h, cab = hashlib.sha256(), None
    with p.open("rb") as fh:
        for b in iter(lambda: fh.read(1 << 20), b""):
            if cab is None:
                cab = b[:20]
            h.update(b)
    cab = cab or b""
    return h.hexdigest(), p.stat().st_size, formato(cab) is not None, cab


MIME = {"jpeg": "image/jpeg", "png": "image/png", "tiff": "image/tiff", "gif": "image/gif",
        "jp2": "image/jp2", "j2k": "image/jp2", "bmp": "image/bmp", "webp": "image/webp"}


def tipo_mime(p):
    """pelo conteudo, nao pela extensao: TIFF chega com nome .jpg"""
    with open(p, "rb") as fh:
        f = formato(fh.read(20))
    return MIME.get(f) or mimetypes.guess_type(str(p))[0] or "binary/octet-stream"


def sobe_arquivo(s3, p, key, tam):
    s3.upload_file(str(p), BUCKET, key, Config=TRANSF,
                   ExtraArgs={"ContentType": tipo_mime(p)})
    got = s3.head_object(Bucket=BUCKET, Key=key)["ContentLength"]
    if got != tam:
        raise RuntimeError(f"tamanho no bucket {got} != local {tam}")


def confere_existente(s3, key, tam):
    """True se o objeto ja esta no bucket com o mesmo tamanho."""
    try:
        return s3.head_object(Bucket=BUCKET, Key=key)["ContentLength"] == tam
    except Exception:
        return False


def sobe_bytes(s3, dados, key, ctype):
    s3.put_object(Bucket=BUCKET, Key=key, Body=dados, ContentType=ctype)


def refaz_metadados(con, s3, cat, handles, threads):
    """metadata.json + manifest-sha256.txt, mesmo formato do subidor.py."""
    trabalhos = []
    for h in sorted(handles):
        rec = cat.get(h)
        if not rec:
            continue
        mj = json.dumps(rec, ensure_ascii=False, indent=2).encode("utf-8")
        linhas = [f"{sha}  {tipo}/{nome}" for nome, sha, tipo in con.execute(
            f"SELECT nome, sha256, tipo FROM enviados WHERE handle=? AND tipo IN "
            f"({','.join('?' * len(TIPOS_DO_ITEM))}) ORDER BY tipo, nome", (h, *TIPOS_DO_ITEM))]
        linhas.append(f"{hashlib.sha256(mj).hexdigest()}  metadata.json")
        mf = ("\n".join(linhas) + "\n").encode("utf-8")
        pasta = h.replace("/", "_")
        trabalhos.append((mj, f"itens/{pasta}/metadata.json", "application/json"))
        trabalhos.append((mf, f"itens/{pasta}/manifest-sha256.txt", "text/plain"))
    falhas = 0
    with ThreadPoolExecutor(threads) as ex:
        for f in as_completed([ex.submit(sobe_bytes, s3, *t) for t in trabalhos]):
            try:
                f.result()
            except Exception as e:
                falhas += 1
                log(f"FALHA metadado: {e}")
    return len(trabalhos) // 2, falhas


def drena(con, s3, cat, pasta, threads, lote, max_lotes=0):
    REJ_DIR.mkdir(exist_ok=True)
    env = dup = rej = 0
    bytes_env, t0 = 0, time.time()
    falhou = set()
    n_lotes = 0
    while True:
        if max_lotes and n_lotes >= max_lotes:
            log(f"parando depois de {n_lotes} lote(s) (--max-lotes)")
            break
        n_lotes += 1
        corte = time.time() - 10
        fila = []
        for p in pasta.glob(PREFIXO + "*"):
            if (not p.is_file() or p.suffix == ".crdownload" or p.name in falhou
                    or p.stat().st_mtime > corte or Path(str(p) + ".crdownload").exists()):
                continue
            info = interpreta(p)
            if info:
                fila.append((p, *info))
            else:
                log(f"ignorado (nome inesperado): {p.name}")
                falhou.add(p.name)
        if not fila:
            break
        # nome sem "(N)" primeiro: entre duas copias identicas, fica o nome original
        fila.sort(key=lambda a: (SUFIXO_N.match(a[4]) is not None, a[0].name))
        pedaco = fila[:lote]

        with ThreadPoolExecutor(threads) as ex:
            exames = list(zip(pedaco, ex.map(lambda a: examina(a[0]), pedaco)))

        no_ledger = defaultdict(dict)  # (handle, tipo, sha) -> {nome-base: nome ja enviado}
        for h in {a[1] for a, _ in exames}:
            for tipo, sha, nome in con.execute(
                    "SELECT tipo, sha256, nome FROM enviados WHERE handle=?", (h,)):
                no_ledger[(h, tipo, sha)].setdefault(base(nome), nome)
        neste_lote = set()             # (handle, tipo, sha, nome-base) que ainda vai subir agora

        subir, conferir = [], []
        for (p, handle, pasta_h, tipo, nome), (sha, tam, eh_img, cab) in exames:
            if not eh_img:
                shutil.move(str(p), REJ_DIR / p.name)
                con.execute("INSERT OR REPLACE INTO rejeitados VALUES (?,?,?,?,?,?,?)",
                            (handle, tipo, nome, tam, sha, cab.hex(), agora()))
                rej += 1
                log(f"REJEITADO (nao e imagem, {tam} bytes, inicio {cab[:12]!r}): {p.name} -> rejeitados/")
                continue
            k = (handle, tipo, sha)
            existente = no_ledger[k].get(base(nome))
            if existente:
                # so apaga o local depois do HEAD confirmar o objeto no bucket
                conferir.append((p, handle, pasta_h, tipo, nome, sha, tam, existente))
                continue
            if k + (base(nome),) in neste_lote:
                continue            # copia identica sobe neste lote; esta fica pro proximo
            neste_lote.add(k + (base(nome),))
            subir.append((p, handle, pasta_h, tipo, nome, sha, tam))
        con.commit()

        tocados = set()
        with ThreadPoolExecutor(threads) as ex:
            futs = {ex.submit(sobe_arquivo, s3, a[0], f"itens/{a[2]}/{a[3]}/{a[4]}", a[6]): a
                    for a in subir}
            for f in as_completed(futs):
                p, handle, pasta_h, tipo, nome, sha, tam = futs[f]
                try:
                    f.result()
                except Exception as e:
                    falhou.add(p.name)
                    log(f"FALHA no upload (mantido em disco): {p.name}: {e}")
                    continue
                con.execute("INSERT OR REPLACE INTO enviados (handle,nome,bytes,sha256,quando,tipo) "
                            "VALUES (?,?,?,?,?,?)", (handle, nome, tam, sha, agora(), tipo))
                con.commit()
                p.unlink()
                tocados.add(handle)
                env += 1
                bytes_env += tam
        with ThreadPoolExecutor(threads) as ex:
            futs = {ex.submit(confere_existente, s3, f"itens/{a[2]}/{a[3]}/{a[7]}", a[6]): a
                    for a in conferir}
            for f in as_completed(futs):
                p, handle, pasta_h, tipo, nome, sha, tam, existente = futs[f]
                chave = f"itens/{pasta_h}/{tipo}/{existente}"
                if f.result():
                    con.execute("INSERT OR REPLACE INTO descartados VALUES (?,?,?,?,?,?,?)",
                                (handle, tipo, nome, tam, sha, chave, agora()))
                    con.commit()
                    p.unlink()
                    dup += 1
                else:
                    # ledger dizia que estava la, o bucket discorda: sobe de novo
                    log(f"ledger x bucket divergem, reenviando: {chave}")
                    try:
                        sobe_arquivo(s3, p, chave, tam)
                    except Exception as e:
                        falhou.add(p.name)
                        log(f"FALHA no reenvio (mantido em disco): {p.name}: {e}")
                        continue
                    con.execute("INSERT OR REPLACE INTO enviados (handle,nome,bytes,sha256,quando,tipo) "
                                "VALUES (?,?,?,?,?,?)", (handle, existente, tam, sha, agora(), tipo))
                    con.commit()
                    p.unlink()
                    tocados.add(handle)
                    env += 1
                    bytes_env += tam
        n_meta, f_meta = refaz_metadados(con, s3, cat, tocados, threads)

        horas = (time.time() - t0) / 3600
        log(f"lote ok: {len(tocados)} itens | total {env} arq, {bytes_env/2**30:.1f} GB "
            f"({bytes_env/2**30/horas:.1f} GB/h) | dup {dup} | rej {rej} | falha {len(falhou)} "
            f"| meta {n_meta} (falha {f_meta}) | fila ~{len(fila) - len(pedaco)}")
    log(f"FIM da fila: {env} enviados, {bytes_env/2**30:.1f} GB, {dup} duplicatas descartadas, "
        f"{rej} rejeitados, {len(falhou)} falhas mantidas em disco")
    for n in sorted(falhou):
        log(f"  pendente: {n}")


def limpa_duplicatas(con, s3):
    grupos = defaultdict(list)
    for h, tipo, nome, sha in con.execute(
            "SELECT handle, tipo, nome, sha256 FROM enviados WHERE tipo IN ('preservacao','thumbnail')"):
        grupos[(h, tipo, sha)].append(nome)
    apagar = []
    for (h, tipo, sha), nomes in grupos.items():
        if len(nomes) < 2:
            continue
        for b, iguais in _por_base(nomes).items():
            if len(iguais) < 2:
                continue
            fica = sorted(iguais, key=lambda n: (SUFIXO_N.match(n) is not None, n))[0]
            apagar += [(h, tipo, n, fica) for n in iguais if n != fica]
    log(f"duplicatas identicas no ledger: {len(apagar)}")
    feitos, tocados = 0, set()

    def um(a):
        h, tipo, n, fica = a
        pasta = h.replace("/", "_")
        k_dup, k_fica = f"itens/{pasta}/{tipo}/{n}", f"itens/{pasta}/{tipo}/{fica}"
        t1 = s3.head_object(Bucket=BUCKET, Key=k_dup)["ContentLength"]
        t2 = s3.head_object(Bucket=BUCKET, Key=k_fica)["ContentLength"]
        if t1 != t2:
            raise RuntimeError(f"tamanhos divergem no bucket: {k_dup}")
        s3.delete_object(Bucket=BUCKET, Key=k_dup)
        return a

    with ThreadPoolExecutor(16) as ex:
        for f in as_completed([ex.submit(um, a) for a in apagar]):
            try:
                h, tipo, n, _ = f.result()
            except Exception as e:
                log(f"  nao apagado: {e}")
                continue
            con.execute("DELETE FROM enviados WHERE handle=? AND tipo=? AND nome=?", (h, tipo, n))
            tocados.add(h)
            feitos += 1
    con.commit()
    log(f"duplicatas apagadas: {feitos}")
    return tocados


def _por_base(nomes):
    d = defaultdict(list)
    for n in nomes:
        d[base(n)].append(n)
    return d


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--downloads", default=str(Path.home() / "Downloads"))
    ap.add_argument("--catalogo")
    # 6 e nao 16: com muitas threads o OpenSSL do macOS estoura na verificacao
    # de certificado (SIGSEGV em X509_STORE_CTX_get1_issuer, 13/09) -- corrida
    # dentro do libcrypto, nao do nosso codigo. 6 threads ja dao ~19 GB/h.
    ap.add_argument("--threads", type=int, default=6)
    ap.add_argument("--lote", type=int, default=400)
    ap.add_argument("--max-lotes", type=int, default=0, help="0 = ate esvaziar a fila")
    ap.add_argument("--limpar-duplicatas", action="store_true")
    ap.add_argument("--so-manifestos", action="store_true")
    args = ap.parse_args()
    if not os.environ.get("AWS_ACCESS_KEY_ID"):
        sys.exit("credenciais S3 ausentes. Rode:  set -a; . .s3env; set +a")

    con, s3 = db(), cliente(args.threads)
    cat = carrega_catalogo(args.catalogo)
    if args.limpar_duplicatas:
        tocados = limpa_duplicatas(con, s3)
        if cat and tocados:
            log("regravando manifest dos itens tocados: %d (falha %d)" %
                refaz_metadados(con, s3, cat, tocados, args.threads))
    elif args.so_manifestos:
        if not cat:
            sys.exit("--so-manifestos precisa de --catalogo")
        hs = {r[0] for r in con.execute("SELECT DISTINCT handle FROM enviados WHERE tipo IN (?,?,?)",
                                        TIPOS_DO_ITEM)}
        log(f"regravando metadata+manifest de {len(hs)} itens")
        log("feito: %d itens (falha %d)" % refaz_metadados(con, s3, cat, hs, args.threads))
    else:
        log(f"drenando {args.downloads} -> s3://{BUCKET}/ ({args.threads} threads, lotes de {args.lote})")
        drena(con, s3, cat, Path(args.downloads), args.threads, args.lote, args.max_lotes)


if __name__ == "__main__":
    main()