"""Turn what our runs actually produce into the viewer's contract.

The stack already writes two things we do not control the shape of: `eval_pass_at_k.py` dumps a JSON of
per-harness summaries plus raw rows, and AsyncGRPO logs training metrics to a trackio sqlite. Rather
than change either — a viewer should not dictate how a trainer logs — this converts them into
`runs/<run_id>/` as CONTRACT.md describes.

    # an eval sweep
    python tools/ingest.py eval  --json logs/eval_6harness.json --run-id 2b-6harness --model Qwen/Qwen3.5-2B

    # a training run, straight from the trackio db
    python tools/ingest.py train --trackio ~/runs/agrpo_harbor/trackio/<project>.db \
                                 --run-name Qwen3.5-2B-mini-swe-agent-20steps-46552 --run-id 2b-mini-20
"""

from __future__ import annotations

import argparse
import json
import sqlite3
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parents[1]


def _now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _write(run_dir: Path, rel: str, payload) -> None:
    p = run_dir / rel
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(payload, indent=2, default=str) if not isinstance(payload, str) else payload)
    print(f"  wrote {p.relative_to(HERE)}")


def _task_metadata(split: str) -> dict[int, dict]:
    """index -> {id, question, answer, difficulty_level} from the Harbor suite itself."""
    import sys as _sys

    _sys.path.insert(0, str(HERE.parents[0] / "async_grpo_harbor_data_agent" / "src"))
    from harbor_tasks import _read_meta, download_suite  # type: ignore

    root = download_suite(split)
    out = {}
    for i, toml in enumerate(sorted(root.glob("tasks/*/task.toml"))):
        m = dict(_read_meta(toml.parent).get("metadata") or {})
        instr = toml.parent / "instruction.md"
        # `meta` is whatever the suite records, carried verbatim. Promoting a fixed set of keys to
        # columns means the viewer breaks the moment a suite adds or renames one — and difficulty,
        # package_tier and reward_mode are exactly the kind of field that changes between suites. The
        # viewer discovers the keys instead and offers them as filters.
        out[i] = {
            "id": toml.parent.name,
            "answer": m.pop("gold_answer", None),
            "question": (instr.read_text()[:4000] if instr.exists() else None),
            "meta": m,
        }
    return out


def ds_dir_for(args):
    return HERE / "data" / "projects" / args.project / "datasets" / args.dataset_id


def ingest_eval(args) -> None:
    raw = json.loads(Path(args.json).read_text())
    summary_in = raw.get("summary", {})
    per_harness_in = summary_in.get("harnesses") or {}
    rows_in = raw.get("rows") or {}

    # Two shapes exist in the wild and both are real output from this stack: the multi-harness sweep
    # writes {"rows": {harness: [...]}} while the earlier single-harness runs wrote
    # {"per_task": {index: [...]}} with no harness on the rows. Normalising here rather than rejecting
    # the older one keeps already-collected results usable — they are the only baseline we have.
    if not rows_in and raw.get("per_task"):
        harness = summary_in.get("harness") or args.harness_fallback
        rows_in = {harness: [r for rows in raw["per_task"].values() for r in rows]}
        if not per_harness_in:
            per_harness_in = {harness: {k: v for k, v in summary_in.items() if k.startswith("pass@") or k == "mean_turns"}}

    # pass@k over TASKS, pass@1 over SAMPLES, and `reward is None` excluded rather than scored 0 —
    # the same rule the eval tool applies, restated here so the published numbers cannot drift from it.
    by_harness, tasks_acc = [], defaultdict(dict)
    task_meta: dict[int, dict] = {}
    if args.tasks_from:
        # Question, gold answer and difficulty come from the suite, not from the eval output. Without
        # them a row is an opaque index and a cell cannot be judged by eye.
        task_meta = _task_metadata(args.tasks_from)
    totals = {"tasks_total": 0, "tasks_any_pass": 0, "attempts_total": 0, "attempts_passed": 0, "n_all_infra": 0}
    k = summary_in.get("k", args.k)

    for harness, rows in rows_in.items():
        by_task = defaultdict(list)
        for r in rows:
            by_task[r["index"]].append(r)
        measured = solved = infra = 0
        samples, turns = [], []
        for index, rs in by_task.items():
            graded = [r for r in rs if r.get("reward") is not None]
            key = f"{args.model}|{harness}"
            attempts = []
            for i, r in enumerate(sorted(rs, key=lambda r: r.get("rep", 0))):
                att = {"attempt": i + 1, "reward": r.get("reward"), "n_turns": r.get("n_turns")}
                # A trace is written only when the row carries one. Referencing a file that does not
                # exist would give the viewer an "open" button that always 404s.
                att["elapsed_sec"] = r.get("wall_s")
                if r.get("messages"):
                    tid = f"{index}-{harness}-{i + 1}.json"
                    _write(ds_dir_for(args), f"traces/{tid}", {
                        "task_index": index, "task_id": r.get("task_id"),
                        "model": args.model, "harness": harness, "attempt": i + 1,
                        "reward": r.get("reward"), "rewards": r.get("rewards") or {},
                        "n_turns": r.get("n_turns"), "elapsed_sec": r.get("wall_s"),
                        "rollout_type": r.get("rollout_type"),
                        "n_trainable_tokens": r.get("n_trainable_tokens"),
                        "trial_name": r.get("trial_name"),
                        "messages": r["messages"],
                    })
                    att["trace"] = f"traces/{tid}"
                attempts.append(att)
            first_pass = next((a["attempt"] for a in attempts if (a["reward"] or 0) > 0), None)
            tasks_acc[index][key] = {"passed_at": first_pass, "attempts": attempts}
            if not graded:
                infra += 1
                continue
            measured += 1
            if any((r["reward"] or 0) > 0 for r in graded):
                solved += 1
            samples.extend(r["reward"] for r in graded)
            turns.extend(r.get("n_turns") or 0 for r in graded)
        m = per_harness_in.get(harness, {})
        by_harness.append({
            "harness": harness,
            f"pass@{k}": m.get(f"pass@{k}", round(solved / measured, 4) if measured else None),
            "pass@1": m.get("pass@1", round(sum(samples) / len(samples), 4) if samples else None),
            "mean_turns": m.get("mean_turns", round(sum(turns) / len(turns), 2) if turns else None),
            "n_measured": measured, "cells": len(by_task), "n_all_infra": infra,
        })
        totals["attempts_total"] += sum(len(v) for v in by_task.values())
        totals["attempts_passed"] += sum(1 for s in samples if s > 0)
        totals["n_all_infra"] += infra

    totals["tasks_total"] = len(tasks_acc)
    totals["tasks_any_pass"] = sum(
        1 for cells in tasks_acc.values() if any(c["passed_at"] for c in cells.values())
    )

    proj = HERE / "data" / "projects" / args.project
    ds = proj / "datasets" / args.dataset_id

    # by_model mirrors by_harness so the viewer's pivot has aggregates on both axes. With one model in a
    # sweep it is a single row, which is honest rather than empty.
    model_samples = [a["reward"] for cells in tasks_acc.values() for c in cells.values()
                     for a in c["attempts"] if a["reward"] is not None]
    by_model = [{
        "model": args.model,
        "cells": totals["attempts_total"],
        f"pass@{k}": round(totals["tasks_any_pass"] / totals["tasks_total"], 4) if totals["tasks_total"] else None,
        "pass@1": round(sum(1 for r in model_samples if r > 0) / len(model_samples), 4) if model_samples else None,
        "n_measured": totals["tasks_total"] - totals["n_all_infra"],
        "mean_turns": None,
    }]

    _write(ds, "summary.json", {
        "k_max": k,
        "models": [args.model],
        "harnesses": sorted(rows_in),
        "summary": totals,
        "by_harness": sorted(by_harness, key=lambda h: -(h.get(f"pass@{k}") or 0)),
        "by_model": by_model,
        "tasks": [
            {"id": task_meta.get(i, {}).get("id", str(i)), "index": i,
             "question": task_meta.get(i, {}).get("question"),
             "answer": task_meta.get(i, {}).get("answer"),
             "meta": task_meta.get(i, {}).get("meta", {}),
             "cells": cells}
            for i, cells in sorted(tasks_acc.items())
        ],
    })
    _write(ds, "dataset.json", {
        "dataset_id": args.dataset_id, "label": args.dataset_label or args.dataset_id,
        "split": "eval", "source": args.dataset, "k": k,
        "created_at": _now(), "notes": args.notes,
    })
    if not (proj / "project.json").exists():
        _write(proj, "project.json", {
            "project_id": args.project, "label": args.project_label or args.project,
            "description": args.project_description,
            "source": {"hf_dataset": args.dataset},
            "support": {},
        })


def ingest_train(args) -> None:
    """Read trackio's sqlite directly: it is the source of truth for a finished run, and re-deriving
    metrics from stdout would invent numbers the trainer never logged."""
    con = sqlite3.connect(args.trackio)
    rows = con.execute(
        "select step, metrics from metrics where run_name like ? and length(metrics) > 4 order by step",
        (f"%{args.run_name}%",),
    ).fetchall()
    if not rows:
        raise SystemExit(f"no metric rows matching {args.run_name!r} in {args.trackio}")

    lines = []
    for step, blob in rows:
        d = json.loads(blob if isinstance(blob, (str, bytes)) else str(blob))
        lines.append(json.dumps({
            "step": step,
            "reward": d.get("train/reward"), "reward_std": d.get("train/reward_std"),
            "loss": d.get("train/loss"), "ratio": d.get("train/ratio"),
            "kl": d.get("train/kl"), "entropy": d.get("train/entropy"),
            "learning_rate": d.get("train/learning_rate"),
        }))

    run_dir = HERE / "data" / "projects" / args.project / "runs" / args.run_id
    _write(run_dir, "train/metrics.jsonl", "\n".join(lines) + "\n")
    with_grad = sum(1 for line in lines if (json.loads(line).get("reward_std") or 0) > 0)
    _write(run_dir, "run.json", {
        "run_id": args.run_id, "kind": "train", "created_at": _now(), "updated_at": _now(),
        "model": args.model, "harnesses": [args.harness] if args.harness else [],
        "sandbox": args.sandbox, "dataset": args.dataset, "split": "train",
        "notes": args.notes or f"{len(lines)} steps, {with_grad} with a non-zero gradient",
        "config": {"trackio_run": args.run_name},
    })
    print(f"  {len(lines)} steps, {with_grad} with reward_std > 0")


def main() -> int:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    e = sub.add_parser("eval")
    e.add_argument("--json", required=True, help="output of tools/eval_pass_at_k.py")
    e.add_argument("--model", default="Qwen/Qwen3.5-2B")
    e.add_argument("--sandbox", default="e2b")
    e.add_argument("--dataset", default="AdithyaSK/data_agent_rl_environment_eval")
    e.add_argument("--k", type=int, default=4)
    e.add_argument("--notes", default="")
    e.add_argument("--project", default="data-agent")
    e.add_argument("--project-label", default="Data-Agent Bench")
    e.add_argument("--project-description", default="")
    e.add_argument("--dataset-id", required=True, help="a variation within the project, e.g. eval-easy50")
    e.add_argument("--dataset-label", default="")
    e.add_argument("--tasks-from", default="", help="Harbor suite to pull question/answer/difficulty from")
    e.add_argument("--harness-fallback", default="mini-swe-agent",
                   help="harness name for older single-harness JSON that does not record one")
    e.set_defaults(fn=ingest_eval)

    t = sub.add_parser("train")
    t.add_argument("--trackio", required=True, help="path to the trackio sqlite db")
    t.add_argument("--run-name", required=True, help="substring of the trackio run name")
    t.add_argument("--run-id", required=True)
    t.add_argument("--model", default="Qwen/Qwen3.5-2B")
    t.add_argument("--harness", default="mini-swe-agent")
    t.add_argument("--sandbox", default="e2b")
    t.add_argument("--dataset", default="AdithyaSK/data_agent_rl_environment_train")
    t.add_argument("--notes", default="")
    t.add_argument("--project", default="data-agent")
    t.set_defaults(fn=ingest_train)

    args = ap.parse_args()
    where = (f"projects/{args.project}/datasets/{args.dataset_id}" if args.cmd == "eval"
             else f"projects/{args.project}/runs/{args.run_id}")
    print(f"ingesting {args.cmd} -> data/{where}/")
    args.fn(args)
    return 0


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