"""Coleta os resumos de avaliação (dmb-resultados e dmb-resultados-gru) para reports/<destino>/.

Para cada (candidato, dispositivo, partição) usa a execução mais recente bem-sucedida com o
número de exemplos pedido. Gera os JSONs, as predições e um README.md com a tabela comparativa.

Uso: PYTHONPATH=infra/oci .venv/bin/python infra/ds/collect.py --split dev --examples 50 --dest reports/pilot
"""
import argparse
import json
from pathlib import Path

import oci

import provision
from session import load_config

ROOT = Path(__file__).resolve().parents[2]
SOURCES = [(provision.REGION, "dmb-resultados"), (provision.GPU_REGION, f"dmb-resultados{provision.GPU_SUFFIX}")]


def fmt(v, nd=1):
    return "—" if v is None else f"{v:.{nd}f}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", default="dev")
    ap.add_argument("--examples", type=int)
    ap.add_argument("--dest", required=True)
    ap.add_argument("--all-reps", action="store_true", help="guarda cada repetição (DMB_REP) em separado")
    args = ap.parse_args()
    dest = ROOT / args.dest
    dest.mkdir(parents=True, exist_ok=True)
    latest = {}
    for region, bucket in SOURCES:
        c = oci.object_storage.ObjectStorageClient(load_config(region=region))
        ns = c.get_namespace().data
        names = [o.name for o in oci.pagination.list_call_get_all_results(
            c.list_objects, ns, bucket, prefix="runs/").data.objects]
        for n in names:
            fname = n.rsplit("/", 1)[-1]
            if not (fname.startswith("summary-") and fname.endswith(f"-{args.split}.json")):
                continue
            s = json.loads(c.get_object(ns, bucket, n).data.content)
            if args.examples and s["examples"] != args.examples:
                continue
            if s.get("operational_env"):  # execuções exploratórias com parâmetros alterados ficam de fora
                continue
            key = (s["candidate"]["name"], s["device"]) + ((s.get("repetition", 1),) if args.all_reps else ())
            run_id = n.split("/")[1]
            if key not in latest or run_id > latest[key][0]:
                pred = n.replace("summary-", "predictions-").replace(".json", ".jsonl")
                latest[key] = (run_id, s, c, ns, bucket, pred, region)
    rows = []
    for key, (run_id, s, c, ns, bucket, pred, region) in sorted(latest.items()):
        cand, dev = key[0], key[1]
        suffix = f"-{dev}-r{key[2]}" if args.all_reps else f"-{dev}"
        s["oci"] = {"region": region, "run_id": run_id}
        (dest / f"summary-{cand}{suffix}.json").write_text(json.dumps(s, indent=2, ensure_ascii=False) + "\n")
        try:
            (dest / f"predictions-{cand}{suffix}.jsonl").write_bytes(c.get_object(ns, bucket, pred).data.content)
        except oci.exceptions.ServiceError:
            pass
        q, lat, cal = s["quality"], s["latency_ms"], s.get("calibration_raw", {})
        rows.append(f"| {cand} | {dev}{' r' + str(key[2]) if args.all_reps else ''} | {s.get('precision') or '—'} | {q.get('n', 0)} | "
                    f"{fmt(q.get('accuracy', 0) * 100)}% | {fmt(q.get('macro_f1', 0) * 100)}% | "
                    f"{fmt(cal.get('ece'), 3)} | {fmt(lat['p50'])} | {fmt(lat['p95'])} | {fmt(lat['p99'])} | "
                    f"{fmt(s['load_seconds'])} | {fmt(s.get('rss_peak_mib'), 0)} | {fmt(s.get('gpu_mem_peak_mib'), 0)} | "
                    f"{sum(s['errors'].values())} | {s['truncated']} |")
    table = ("| Candidato | Disp. | Precisão | n | Acurácia | Macro F1 | ECE (bruto) | p50 ms | p95 ms | p99 ms | "
             "Carga s | RAM pico MiB | GPU pico MiB | Erros | Truncados |\n"
             "|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n" + "\n".join(rows))
    start, end = "<!-- TABELA:INICIO -->", "<!-- TABELA:FIM -->"
    block = f"{start}\n{table}\n{end}"
    readme_path = dest / "README.md"
    if readme_path.exists() and start in readme_path.read_text():
        # README editado à mão: só a tabela entre os marcadores é regenerada.
        text = readme_path.read_text()
        readme = text[:text.index(start)] + block + text[text.index(end) + len(end):]
    else:
        readme = (f"# Resultados — partição `{args.split}`"
                  + (f", {args.examples} exemplos" if args.examples else "") + "\n\n"
                  "Gerado por `infra/ds/collect.py`. Lote 1; latências sem os exemplos de aquecimento; "
                  "probabilidades sem calibração.\n\n" + block + "\n")
    readme_path.write_text(readme)
    print(table)


if __name__ == "__main__":
    main()