"""Recover raw per-round tables from orx run logs.

This project runs in orx LOCAL mode, where `orx artifacts` does not exist and the run
log is the only evidence channel. The expensive stages therefore print every per-round
metric to stdout, and this script parses those lines back into CSV so the numbers can
be published as downloadable raw data next to the claims that rest on them.

Theory-stage artifacts are NOT recovered this way: those stages are deterministic and
run in seconds, so they are regenerated by rerunning them, which is a stronger
guarantee than parsing text.

Usage:
    python tools/extract_evidence.py <out_dir> <label>=<run_id> [...]
"""

from __future__ import annotations

import csv
import hashlib
import json
import re
import subprocess
import sys
from pathlib import Path

# Numbers may be negative ("H=-0.0000" for a fully collapsed round), so every
# numeric group allows a leading minus -- without it those rows are silently dropped,
# which loses exactly the collapsed rounds the comparison depends on.
# "    round  7                                       H(L)=1.9312  mean_len= 18.4  distinct= 31  nearA=0.14  nearB=0.09  loss=3.201  74s"
TEXT_ROUND = re.compile(
    r"round\s+(\d+)\s+H\(L\)=(-?[\d.]+)\s+mean_len=\s*(-?[\d.]+)\s+distinct=\s*(\d+)\s+"
    r"nearA=(-?[\d.nan]+)\s+nearB=(-?[\d.nan]+)\s+loss=(-?[\d.]+)\s+(\d+)s"
)
# "    round  3                                       H=1.9312  KL=0.0821  featVar=0.412  intraVar=0.377  classes=10  91s"
FLOW_ROUND = re.compile(
    r"round\s+(\d+)\s+H=(-?[\d.]+)\s+KL=(-?[\d.]+)\s+featVar=(-?[\d.]+)\s+"
    r"intraVar=(-?[\d.]+)\s+classes=(\d+)\s+(\d+)s"
)
KV = re.compile(r"^\s{4}(\S.*?)\s{2,}(.+?)\s*$")

EMIT_BEGIN = re.compile(
    r"^<<<ORX-ARTIFACT path=(\S+) bytes=(\d+) sha256=([0-9a-f]{64}) "
    r"encoding=(\S+) status=(\S+)"
)


def recover_artifacts(log: str, dest: Path) -> list[dict]:
    """Rebuild the artifact files a job wrote, from the framed blocks in its log.

    Each block declares its own length and SHA-256, so a truncated or interleaved
    log is rejected rather than silently yielding a corrupt file.
    """
    import base64

    recovered, lines, i = [], log.splitlines(), 0
    while i < len(lines):
        m = EMIT_BEGIN.match(lines[i])
        if not m:
            i += 1
            continue
        rel, size, digest, encoding, status = m.groups()
        i += 1
        payload = []
        while i < len(lines) and not lines[i].startswith(">>>ORX-ARTIFACT-END"):
            payload.append(lines[i].strip())
            i += 1
        i += 1
        rec = {"path": rel, "declared_bytes": int(size), "declared_sha256": digest,
               "status": status, "ok": False}
        if encoding == "base64" and status == "OK":
            try:
                raw = base64.b64decode("".join(payload), validate=True)
            except Exception as exc:  # noqa: BLE001 - report, never guess
                rec["error"] = f"base64 decode failed: {exc}"
            else:
                actual = hashlib.sha256(raw).hexdigest()
                rec["actual_bytes"], rec["actual_sha256"] = len(raw), actual
                if actual == digest and len(raw) == int(size):
                    out = dest / rel
                    out.parent.mkdir(parents=True, exist_ok=True)
                    out.write_bytes(raw)
                    rec["ok"] = True
                else:
                    rec["error"] = "hash/length mismatch — log is truncated or interleaved"
        else:
            rec["error"] = f"not recoverable (encoding={encoding} status={status})"
        recovered.append(rec)
    return recovered


def fetch(run_id: str) -> str:
    """Fetch the WHOLE log.

    `orx logs` tails by default, so a byte budget smaller than the log silently
    drops the beginning — which, with artifact frames spread through the run,
    means silently losing evidence. Ask for far more than any run produces.
    """
    return subprocess.run(
        ["orx", "logs", run_id, "--bytes", "100000000"],
        capture_output=True, text=True, check=True,
    ).stdout


def parse(log: str) -> dict:
    text_rows, flow_rows = [], []
    for m in TEXT_ROUND.finditer(log):
        text_rows.append({
            "round": int(m.group(1)), "H_L": float(m.group(2)),
            "mean_len": float(m.group(3)), "n_distinct_lengths": int(m.group(4)),
            "frac_near_T_A": float(m.group(5)), "frac_near_T_B": float(m.group(6)),
            "train_loss": float(m.group(7)), "seconds": int(m.group(8)),
        })
    for m in FLOW_ROUND.finditer(log):
        flow_rows.append({
            "round": int(m.group(1)), "class_entropy": float(m.group(2)),
            "kl_to_uniform": float(m.group(3)), "feature_variance": float(m.group(4)),
            "intra_class_variance": float(m.group(5)),
            "n_classes_present": int(m.group(6)), "seconds": int(m.group(7)),
        })
    meta = {}
    for line in log.splitlines():
        m = KV.match(line)
        if m and len(m.group(1)) < 60:
            meta.setdefault(m.group(1).strip(), m.group(2).strip())
    verdicts = re.findall(r"VERDICT (\S+): (VERIFIED|FALSIFIED|BLOCKED)", log)
    gates = re.findall(r"GATE: (PASS|FAIL)", log)
    return {
        "text_rounds": text_rows, "flow_rounds": flow_rows, "meta": meta,
        "verdicts": [{"claim": c, "status": s} for c, s in verdicts],
        "gates": gates,
        "all_gates_pass": bool(gates) and all(g == "PASS" for g in gates),
        "exit_ok": "OK: every recorded verdict is supported by its own evidence." in log,
    }


def main() -> int:
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    out = Path(sys.argv[1])
    out.mkdir(parents=True, exist_ok=True)
    index = {}
    for spec in sys.argv[2:]:
        label, run_id = spec.split("=", 1)
        log = fetch(run_id)
        (out / f"log_{label}.txt").write_text(log)
        d = parse(log)
        d["run_id"] = run_id
        d["artifacts"] = recover_artifacts(log, out / "artifacts" / label)
        bad = [a for a in d["artifacts"] if not a["ok"]]
        print(f"{label}: recovered {len(d['artifacts']) - len(bad)}/{len(d['artifacts'])} artifacts"
              + (f"  ({len(bad)} FAILED)" if bad else ""))
        for a in bad:
            print(f"  !! {a['path']}: {a.get('error')}")
        for key, rows in (("text_rounds", d["text_rounds"]), ("flow_rounds", d["flow_rounds"])):
            if rows:
                p = out / f"rounds_{label}.csv"
                with p.open("w", newline="") as fh:
                    w = csv.DictWriter(fh, fieldnames=list(rows[0]))
                    w.writeheader()
                    w.writerows(rows)
                print(f"{label}: {len(rows)} rounds -> {p}")
        index[label] = d
        (out / f"meta_{label}.json").write_text(json.dumps(d, indent=2))
    (out / "index.json").write_text(json.dumps(index, indent=2))
    print(f"wrote {out/'index.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())