#!/usr/bin/env python3
"""Build monaco_dev standalone-baseline trajectory shards for the viewer.

Reads a `response` JSONL (each row carries the full agent loop in `events`
for agentic cells, or a single `answer` string for non-agentic cells 1-3),
joins against the matching AML LLM-judge output (with
`--deterministic-extract` = the canonical MoNaCo scoring protocol), and
writes one JSON shard per qid plus a wrapped `index.json` with summary
stats.

Divergences from the wiki_opentable variant:
- Scoring column is `judge_score` (LLM judge output), not deterministic F1.
- We surface `deterministic_exact_answer.items` from the judge as
  `pred_items` — this is the same list-extraction protocol the judge uses
  when computing precision/recall.
- Accepts both `dataset == "monaco"` and `dataset == "monaco_dev"` on
  predictions (submitters stamp both flavors).

Per-event tool_result content is truncated at 8 KB to keep each shard
browser-friendly.

Usage:

    # Baseline A / cell 4 → "structure per q"
    python scripts/build_trajectories.py \\
        --predictions /home/azureuser/projects/information-scaffolds/outputs/monaco_dev/baselines/cell4_agentic_a/named-outputs/response/response \\
        --judged      /home/azureuser/projects/information-scaffolds/outputs/monaco_dev/judges/c4_agentic_a/named-outputs/judged/judged \\
        --out         trajectories \\
        --label       "structure per q (monaco_dev · Baseline A · cell4_agentic_a)"

    # Baseline B / cell 5 → "structure per ds"
    python scripts/build_trajectories.py \\
        --predictions /home/azureuser/projects/information-scaffolds/outputs/monaco_dev/baselines/cell5_agentic_b/named-outputs/response/response \\
        --judged      /home/azureuser/projects/information-scaffolds/outputs/monaco_dev/judges/c5_agentic_b/named-outputs/judged/judged \\
        --out         trajectories_corpus \\
        --label       "structure per ds (monaco_dev · Baseline B · cell5_agentic_b)"

    # Baseline C / cell 6 → "DCI"  (use the good v2 rerun, not the broken v1)
    python scripts/build_trajectories.py \\
        --predictions /home/azureuser/projects/information-scaffolds/outputs/cell6_v2_reruns/monaco_dev_cell6_v2/named-outputs/response/response \\
        --judged      /home/azureuser/projects/information-scaffolds/outputs/monaco_dev/judges/c6_agentic_c_v2/named-outputs/judged/judged \\
        --out         trajectories_rawtext \\
        --label       "DCI (monaco_dev · Baseline C · cell6_agentic_c v2)"

    # Cells 1-3 (single-turn, no `events`) also work — no trajectory section
    # is rendered but scores + full model answer show up.
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path
from typing import Any, Dict, Iterable, List


_CONTENT_CAP = 8000


def _truncate_content(s: str):
    if s is None:
        return "", False
    if len(s) <= _CONTENT_CAP:
        return s, False
    head = s[: _CONTENT_CAP // 2]
    tail = s[-_CONTENT_CAP // 2 :]
    return f"{head}\n\n… [truncated {len(s) - _CONTENT_CAP} chars] …\n\n{tail}", True


def project_event(ev: Dict[str, Any]) -> Dict[str, Any]:
    content, trunc = _truncate_content(ev.get("content"))
    out = {
        "type": ev.get("type"),
        "name": ev.get("name"),
        "input": ev.get("input"),
        "content": content,
    }
    if trunc:
        out["truncated"] = True
    return out


# ─── I/O ──────────────────────────────────────────────────────────────────────


def load_predictions(path: Path) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            d = json.loads(line)
            if d.get("dataset") not in ("monaco", "monaco_dev"):
                continue
            out[str(d["qid"])] = d
    return out


def load_judged(path: Path) -> Dict[str, Dict[str, Any]]:
    """Read judged JSONL: one row per qid with `parsed.judge_score` etc."""
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            d = json.loads(line)
            if d.get("dataset") not in ("monaco", "monaco_dev"):
                continue
            out[str(d["qid"])] = d
    return out


def load_gold(path: Path) -> Dict[str, Dict[str, Any]]:
    """Load gold from unified JSONL or unified/records/*.json."""
    out: Dict[str, Dict[str, Any]] = {}
    if path.is_dir():
        for shard in sorted(path.glob("*.json")):
            with shard.open() as f:
                d = json.load(f)
            qid = str(d.get("qid") or shard.stem)
            out[qid] = d
        return out
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            out[str(d["qid"])] = d
    return out


# ─── Build ────────────────────────────────────────────────────────────────────


def build_record(
    qid: str,
    pred: Dict[str, Any],
    judged: Dict[str, Any] | None,
    gold_row: Dict[str, Any],
) -> Dict[str, Any]:
    gold_items = list(gold_row.get("answers") or gold_row.get("answer_list") or [])

    j = (judged or {}).get("parsed") or {}
    det = (judged or {}).get("deterministic_exact_answer") or {}
    pred_items = list(det.get("items") or [])
    judge_score = float(j.get("judge_score") or 0.0)
    correct = 1 if judge_score >= 0.99 else 0

    events = [project_event(ev) for ev in (pred.get("events") or [])]

    # cells 1-3 are single-turn — no events, no stop_reason, no turns/tool_calls.
    n_turns = pred.get("turns")
    stop_reason = pred.get("stop_reason") or (
        pred.get("finish_reason") if not events else None
    )

    return {
        "qid": qid,
        "dataset": pred.get("dataset") or "monaco_dev",
        "dataset_origin": pred.get("dataset"),
        "question": pred.get("question") or gold_row.get("question"),
        "gold_answers": gold_items,
        "model": pred.get("model"),
        "mode": pred.get("mode"),
        "system_prompt_file": pred.get("system_prompt_file"),
        "system_prompt": pred.get("system_prompt"),
        "user_prompt": pred.get("user_prompt"),
        "max_turns": pred.get("max_turns"),
        "max_completion_tokens": pred.get("max_completion_tokens"),
        "reasoning_effort": pred.get("reasoning_effort"),
        "stop_reason": stop_reason,
        "finish_reason": pred.get("finish_reason"),
        "error": pred.get("error"),
        "finish_reasons": pred.get("finish_reasons"),
        "n_turns": n_turns,
        "tool_call_counts": pred.get("tool_call_counts"),
        "tokens": pred.get("tokens") or pred.get("usage"),
        "latency_ms": pred.get("latency_ms"),
        "attempts": pred.get("attempts"),
        "timeout_retries": pred.get("timeout_retries"),
        "model_answer": pred.get("answer"),
        "pred_items": pred_items,
        "metrics": {
            "judge_score": round(judge_score, 4),
            "correct": correct,
            "precision": round(float(j.get("precision") or 0.0), 4),
            "recall": round(float(j.get("recall") or 0.0), 4),
            "extracted_final_answer": j.get("extracted_final_answer"),
        },
        "judge_model": (judged or {}).get("judge_model"),
        "events": events,
    }


def build_index_row(rec: Dict[str, Any]) -> Dict[str, Any]:
    m = rec["metrics"]
    return {
        "qid": rec["qid"],
        "question": rec["question"],
        "n_turns": rec["n_turns"],
        "stop_reason": rec["stop_reason"],
        "has_error": bool(rec["error"]) or (m["judge_score"] == 0 and not rec["pred_items"] and not rec["model_answer"]),
        "judge_score": m["judge_score"],
        "correct": m["correct"],
        "precision": m["precision"],
        "recall": m["recall"],
        "gold_answers_length": len(rec["gold_answers"]),
    }


def _counter(it: Iterable[Any]) -> Dict[str, int]:
    out: Dict[str, int] = {}
    for x in it:
        k = str(x)
        out[k] = out.get(k, 0) + 1
    return dict(sorted(out.items(), key=lambda kv: -kv[1]))


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--predictions", required=True, type=Path,
                    help="Raw response JSONL from the AML baseline job")
    ap.add_argument("--judged", required=True, type=Path,
                    help="Judged JSONL from submit_judge_responses_aml.py "
                         "(with --deterministic-extract, auto-on for monaco_dev)")
    ap.add_argument(
        "--gold",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "unified" / "records",
        help="Gold source. Defaults to the local unified/records/ dir built by "
             "scripts/build_unified.py; also accepts the raw unified jsonl.",
    )
    ap.add_argument("--out", required=True, type=Path,
                    help="Output dir (e.g. trajectories/ or trajectories_corpus/)")
    ap.add_argument("--label", default="standalone agentic run",
                    help="Human label for this run; written to index meta")
    args = ap.parse_args()

    preds = load_predictions(args.predictions)
    judged = load_judged(args.judged)
    gold = load_gold(args.gold)

    common = sorted(set(preds) & set(gold))
    only_judged_missing = [q for q in common if q not in judged]
    print(f"qids: predictions={len(preds)}, judged={len(judged)}, gold={len(gold)}, common={len(common)}, missing_judge={len(only_judged_missing)}")
    if not common:
        print("ERROR: no overlap between predictions and gold", file=sys.stderr)
        return 1

    out_dir = args.out
    out_dir.mkdir(parents=True, exist_ok=True)
    rec_dir = out_dir / "records"
    if rec_dir.exists():
        shutil.rmtree(rec_dir)
    rec_dir.mkdir(parents=True)

    index_rows: List[Dict[str, Any]] = []
    sum_j = sum_correct = sum_p = sum_r = 0.0
    for qid in common:
        rec = build_record(qid, preds[qid], judged.get(qid), gold[qid])
        (rec_dir / f"{qid}.json").write_text(json.dumps(rec, ensure_ascii=False))
        index_rows.append(build_index_row(rec))
        sum_j += rec["metrics"]["judge_score"]
        sum_correct += rec["metrics"]["correct"]
        sum_p += rec["metrics"]["precision"]
        sum_r += rec["metrics"]["recall"]

    n = len(common)
    summary = {
        "label": args.label,
        "n": n,
        "mean_judge_score": round(sum_j / n, 4),
        "mean_correct": round(sum_correct / n, 4),
        "mean_precision": round(sum_p / n, 4),
        "mean_recall": round(sum_r / n, 4),
        "fraction_error": round(
            sum(1 for r in index_rows if r["has_error"]) / n, 4
        ),
        "stop_reason_counts": _counter(r["stop_reason"] for r in index_rows),
    }
    index_payload = {"meta": summary, "rows": index_rows}
    (out_dir / "index.json").write_text(json.dumps(index_payload, ensure_ascii=False))

    print(f"\n✓ Wrote {out_dir}/index.json + {n} record shards")
    print(f"   mean judge_score = {summary['mean_judge_score']*100:.2f}")
    print(f"   mean correct     = {summary['mean_correct']*100:.2f}")
    print(f"   mean precision   = {summary['mean_precision']*100:.2f}")
    print(f"   mean recall      = {summary['mean_recall']*100:.2f}")
    print(f"   error rate       = {summary['fraction_error']*100:.2f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())