#!/usr/bin/env python3
"""Build per-qid trajectory + eval shards for the E2E view of the Open-WikiTable viewer.

Joins three sources keyed by qid:
  1. The agent's response.jsonl (predictions + events) from an e2e pipeline run
  2. The wiki_opentable gold (`answer_list`) from test_with_chunks.jsonl
  3. The framework eval_results.json (`exact_string_match` per qid)

For every qid it
  - parses the prediction into a list via the canonical Exact-Answer rules
  - computes set-based P / R / F1 / strict-EM after normalization
  - emits records/<qid>.json with the full trajectory (events truncated per-content)
  - appends a summary row to index.json (qid, question, n_turns, stop_reason,
    f1, em, framework_passed, has_error)

The scoring + normalization logic is INLINED from
`information-scaffolds/evaluation/adapters/wiki_opentable_adapter.py` +
`_parse_exact_answer_semicolon.py` so this script has no dependency on the
information-scaffolds repo. The wiki_opentable prompt instructs the model
to emit `Exact Answer: a; b; c` (semicolon-separated), so this parser
splits on `;` — the comma-split variant systematically under-reports
F1 by ~3–5 pp (see `evaluation/wiki_opentable_rollout_2026-06-22.md` §9
for the bug history and before/after table).

Usage:
    python scripts/build_trajectories_e2e.py \\
        --predictions  /path/to/predictions \\
        --gold         /path/to/wiki_opentable/raw/test_with_chunks.jsonl \\
        --eval-results /path/to/eval_results.json \\
        --out-dir      trajectories_e2e \\
        --label        "wiki-opentable-fullcorpus-allshapes16k (20260622)"
"""

from __future__ import annotations

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


# ─── Parsing the model's Exact-Answer block (mirrors _parse_exact_answer_semicolon.py) ──

_EXACT_ANSWER_RE = re.compile(
    r"Exact\s*Answer\s*:\s*(.*?)(?:\n\s*Confidence\s*:|\Z)",
    re.IGNORECASE | re.DOTALL,
)


def _strip_uncertainty(s: str) -> str:
    s = s.strip()
    if len(s) >= 2 and s[0] in '"\u201c\u201d\'' and s[-1] in '"\u201c\u201d\'':
        s = s[1:-1].strip()
    while s.endswith("?"):
        s = s[:-1].rstrip()
    return s


def extract_answer_payload(model_answer: str) -> str:
    if not model_answer:
        return ""
    m = _EXACT_ANSWER_RE.search(model_answer)
    if not m:
        return model_answer.strip()
    return m.group(1).strip()


def extract_answer_items(model_answer: str) -> List[str]:
    payload = extract_answer_payload(model_answer)
    if not payload:
        return []
    stripped = payload.lstrip()
    if stripped.startswith("["):
        end = stripped.rfind("]")
        if end > 0:
            candidate = stripped[: end + 1]
            try:
                parsed = json.loads(candidate)
                if isinstance(parsed, list):
                    items = [
                        _strip_uncertainty(str(x))
                        for x in parsed
                        if x is not None and str(x).strip()
                    ]
                    return [x for x in items if x]
            except json.JSONDecodeError:
                pass
    one_line = payload.splitlines()[0].strip()
    # Semicolon-separated list (the canonical wiki_opentable prompt instructs `;`)
    if ";" in one_line:
        items = [_strip_uncertainty(p) for p in one_line.split(";")]
        items = [x for x in items if x]
        if len(items) >= 2:
            return items
    s = _strip_uncertainty(payload)
    return [s] if s else []


# ─── Set-based scoring (mirrors wiki_opentable_adapter.py) ────────────────────

_PUNCT_TRAIL = ".,;:!?"
_QUOTE_CHARS = "\"'\u201c\u201d\u2018\u2019"
_WS_RE = re.compile(r"\s+")
_NUM_RE = re.compile(r"^-?\d{1,3}(?:,\d{3})*(?:\.\d+)?$|^-?\d+(?:\.\d+)?$")


def _normalize(s: str) -> str:
    if s is None:
        return ""
    t = str(s).strip()
    if len(t) >= 2 and t[0] in _QUOTE_CHARS and t[-1] in _QUOTE_CHARS:
        t = t[1:-1].strip()
    while t and t[-1] in _PUNCT_TRAIL:
        t = t[:-1].rstrip()
    t = _WS_RE.sub(" ", t).strip().lower()
    if _NUM_RE.match(t):
        t2 = t.replace(",", "")
        try:
            f = float(t2)
            if f.is_integer():
                return str(int(f))
            return str(f)
        except ValueError:
            return t2
    return t


def score_one(pred_items: List[str], gold_items: List[str]) -> Tuple[float, float, float, int]:
    P = {_normalize(x) for x in pred_items if _normalize(x)}
    G = {_normalize(x) for x in gold_items if _normalize(x)}
    if not P and not G:
        return 1.0, 1.0, 1.0, 1
    if not P:
        return 0.0, 0.0, 0.0, 0
    hit = len(P & G)
    prec = hit / len(P)
    rec = hit / len(G) if G else 0.0
    f1 = (2 * prec * rec / (prec + rec)) if (prec + rec) else 0.0
    em = 1 if P == G else 0
    return prec, rec, f1, em


# ─── Event truncation ─────────────────────────────────────────────────────────

# Per-event content cap. Tool results can be entire file dumps (10s of KB each);
# capping at 8 KB keeps the per-qid shard browser-friendly while still showing
# enough to follow the agent's reasoning. The full content lives in the source
# response.jsonl if a maintainer needs the unabridged version.
_CONTENT_CAP = 8000


def _truncate_content(s: str) -> Tuple[str, bool]:
    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]:
    """Strip events down to {type, name, content, input, truncated?}."""
    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") != "wiki_opentable":
                continue
            out[str(d["qid"])] = d
    return out


def load_gold(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)
            out[str(d["qid"])] = d
    return out


def load_framework_passed(path: Path) -> Dict[str, bool]:
    if not path or not path.exists():
        return {}
    with path.open() as f:
        data = json.load(f)
    rows = data.get("rows") or []
    return {str(r["qid"]): bool(r.get("passed")) for r in rows}


# ─── Main ─────────────────────────────────────────────────────────────────────


def build_record(
    qid: str,
    pred: Dict[str, Any],
    gold_row: Dict[str, Any],
    framework_passed: Dict[str, bool],
) -> Dict[str, Any]:
    gold_items = list(gold_row.get("answer_list", []))
    pred_items = extract_answer_items(pred.get("answer", ""))
    prec, rec, f1, em = score_one(pred_items, gold_items)

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

    return {
        "qid": qid,
        "dataset": "wiki_opentable",
        "dataset_origin": gold_row.get("dataset_origin"),
        "original_table_id": gold_row.get("original_table_id"),
        "question": pred.get("question") or gold_row.get("question_text"),
        "gold_answers": gold_items,
        "sql": gold_row.get("sql"),
        "model": pred.get("model"),
        "mode": pred.get("mode"),
        "system_prompt_file": pred.get("system_prompt_file"),
        "max_turns": pred.get("max_turns"),
        "max_completion_tokens": pred.get("max_completion_tokens"),
        "stop_reason": pred.get("stop_reason"),
        "error": pred.get("error"),
        "finish_reasons": pred.get("finish_reasons"),
        "n_turns": pred.get("turns"),
        "tool_call_counts": pred.get("tool_call_counts"),
        "tokens": pred.get("tokens"),
        "latency_ms": pred.get("latency_ms"),
        "attempts": pred.get("attempts"),
        "model_answer": pred.get("answer"),
        "pred_items": pred_items,
        "metrics": {
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "f1": round(f1, 4),
            "em": em,
            "framework_passed": framework_passed.get(qid),
        },
        "events": events,
    }


def build_index_row(rec: Dict[str, Any]) -> Dict[str, Any]:
    m = rec["metrics"]
    return {
        "qid": rec["qid"],
        "question": rec["question"],
        "dataset_origin": rec["dataset_origin"],
        "n_turns": rec["n_turns"],
        "stop_reason": rec["stop_reason"],
        "has_error": rec["error"] is not None and rec["error"] != "",
        "f1": m["f1"],
        "em": m["em"],
        "framework_passed": m["framework_passed"],
        "n_pred": len(rec["pred_items"]),
        "n_gold": len(rec["gold_answers"]),
    }


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--predictions", required=True, type=Path)
    ap.add_argument("--gold", required=True, type=Path)
    ap.add_argument("--eval-results", type=Path, default=None,
                    help="Optional eval_results.json (for framework_passed comparison)")
    ap.add_argument("--out-dir", required=True, type=Path)
    ap.add_argument("--label", default="e2e run",
                    help="Human label for this run; written to index meta")
    args = ap.parse_args()

    preds = load_predictions(args.predictions)
    gold = load_gold(args.gold)
    fpassed = load_framework_passed(args.eval_results) if args.eval_results else {}

    common = sorted(set(preds) & set(gold))
    print(f"qids: predictions={len(preds)}, gold={len(gold)}, common={len(common)}")
    if not common:
        print("ERROR: no overlap", file=sys.stderr)
        return 1

    out_dir = args.out_dir
    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_f1 = sum_em = 0.0
    for qid in common:
        rec = build_record(qid, preds[qid], gold[qid], fpassed)
        (rec_dir / f"{qid}.json").write_text(json.dumps(rec, ensure_ascii=False))
        index_rows.append(build_index_row(rec))
        sum_f1 += rec["metrics"]["f1"]
        sum_em += rec["metrics"]["em"]

    n = len(common)
    summary = {
        "label": args.label,
        "n": n,
        "mean_f1": round(sum_f1 / n, 4),
        "mean_em": round(sum_em / n, 4),
        "fraction_passed_framework": round(
            sum(1 for r in index_rows if r["framework_passed"]) / n, 4
        ) if fpassed else None,
        "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 F1 = {summary['mean_f1']*100:.2f}")
    print(f"   mean EM = {summary['mean_em']*100:.2f}")
    if summary["fraction_passed_framework"] is not None:
        print(f"   framework exact_string_match = {summary['fraction_passed_framework']*100:.2f}")
    print(f"   error rate = {summary['fraction_error']*100:.2f}")
    return 0


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]))


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