#!/usr/bin/env python3
"""Build open-wikitable `compare/` shards — side-by-side per-qid responses.

For each question in the 500-qid wiki_opentable eval subset, joins the
six canonical c1-c6 responses, c7 naive-search, the existing E2E trajectory,
E2E v3, and E2E-v3 + rawtext. Each answer is parsed via the
canonical semicolon-Exact-Answer rule, and scores against the gold
``answer_list`` with the canonical ``wiki_opentable_adapter``-style
set-based F1 / EM / P / R.

The ten runs are:

  - closed-book        (cell1_closed_book)        — single-shot, no context
  - with-docs          (cell2_with_docs)          — single-shot, gold docs
  - with-structures    (cell3_with_structures)    — single-shot, eval structures
  - structure per q    (cell4_agentic_a)          — formerly Baseline A · per-qid scaffolds
  - structure per ds   (cell5_agentic_b_corpus)   — formerly Baseline B · flat structure corpus
  - DCI                (cell6_agentic_c_rawtext)  — formerly Baseline C · rawtext corpus
  - naive-search       — c7 dense-retrieval agent
  - e2e (legacy)       — existing 2026-06-22 trajectory shards
  - e2e v3             — full E2E-v3 pipeline run
  - e2e v3 + rawtext   — full E2E-v3 overlay with raw documents as fallback

The viewer intentionally computes one deterministic set-F1/EM metric for
every run so the per-qid cards remain directly comparable. Canonical
LLM-judge summaries exist for some full runs but are not mixed into this grid.

Output layout (mirrors monaco-benchmark-viewer/responses/):
    compare/index.json
    compare/records/<qid>.json

Usage:
    python scripts/build_compare.py
    python scripts/build_compare.py --scaffolds-root /path/to/information-scaffolds
"""
from __future__ import annotations

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

DATASET = "wiki_opentable"
HERE = Path(__file__).resolve().parent
REPO = HERE.parent
DEFAULT_SCAFFOLDS = Path(os.environ.get(
    "SCAFFOLDS_ROOT", "/home/azureuser/projects/information-scaffolds"
))
DEFAULT_GOLD = Path(os.environ.get(
    "DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data"
)) / "eval" / DATASET / "raw" / "test_with_chunks.jsonl"
DEFAULT_OUT = REPO / "compare"
DEFAULT_E2E_DIR = REPO / "trajectories_e2e"
DEFAULT_NAIVE = Path.home() / "run_logs/nsa_full_run/outputs/wiki_opentable/named-outputs/response/response"
DEFAULT_E2E_V3 = Path.home() / "run_logs/e2e-wiki_opentable-v3/outputs/v3/named-outputs/predictions/predictions"
DEFAULT_E2E_V3_RAWTEXT = Path.home() / "run_logs/e2e-wiki_opentable-v3-merged-rawtext/outputs/named-outputs/predictions/predictions"

# (label, source, shape)
# source: "cell:<basename>" loads from outputs/agentic_wiki_opentable/<basename>;
#         "e2e_shards"      loads from trajectories_e2e/records/*.json
# shape:  "single" / "agentic" (drives which extra metadata is kept).
# The label is what surfaces in the UI grid; ordering controls left→right cell order.
CELL_CONFIGS: List[Tuple[str, str, str]] = [
    ("closed-book",      "cell:cell1_closed_book.response.jsonl",       "single"),
    ("with-docs",        "cell:cell2_with_docs.response.jsonl",         "single"),
    ("with-structures",  "cell:cell3_with_structures.response.jsonl",   "single"),
    ("structure per q",  "cell:cell4_agentic_a.response.jsonl",         "agentic"),
    ("structure per ds", "cell:cell5_agentic_b_corpus.response.jsonl",  "agentic"),
    ("DCI",              "cell:cell6_agentic_c_rawtext.response.jsonl", "agentic"),
]


# ─── Inlined: semicolon Exact-Answer parser (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_items(model_answer: str) -> List[str]:
    if not model_answer:
        return []
    m = _EXACT_ANSWER_RE.search(model_answer)
    payload = m.group(1).strip() if m else model_answer.strip()
    if not payload:
        return []
    stripped = payload.lstrip()
    if stripped.startswith("["):
        end = stripped.rfind("]")
        if end > 0:
            try:
                parsed = json.loads(stripped[: end + 1])
                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()
    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 []


# ─── Inlined: 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]) -> Dict[str, float]:
    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 {"precision": 1.0, "recall": 1.0, "f1": 1.0, "em": 1}
    if not P:
        return {"precision": 0.0, "recall": 0.0, "f1": 0.0, "em": 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 {
        "precision": round(prec, 4),
        "recall": round(rec, 4),
        "f1": round(f1, 4),
        "em": em,
    }


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


def load_jsonl(path: Path, dataset_filter: str | None = None) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            if dataset_filter is not None and d.get("dataset") != dataset_filter:
                continue
            out[str(d["qid"])] = d
    return out


def load_e2e_shards(records_dir: Path) -> Dict[str, Dict[str, Any]]:
    """Load all per-qid e2e shards from trajectories_e2e/records/.

    Returns a dict shaped roughly like a response.jsonl row, so it slots
    into ``project_config`` without a separate code path:
        {qid, dataset, answer (= model_answer), model, mode, stop_reason,
         turns (= n_turns), max_turns, tool_call_counts, tokens, latency_ms,
         system_prompt_file, max_completion_tokens, finish_reasons}
    """
    out: Dict[str, Dict[str, Any]] = {}
    if not records_dir.exists():
        return out
    for fp in sorted(records_dir.glob("*.json")):
        try:
            r = json.loads(fp.read_text())
        except Exception:
            continue
        qid = str(r.get("qid") or fp.stem)
        out[qid] = {
            "qid": qid,
            "dataset": r.get("dataset") or DATASET,
            "question": r.get("question"),
            # Map e2e-shard field names to response.jsonl-style for project_config:
            "answer": r.get("model_answer") or "",
            "model": r.get("model"),
            "mode": r.get("mode"),
            "system_prompt_file": r.get("system_prompt_file"),
            "max_completion_tokens": r.get("max_completion_tokens"),
            "latency_ms": r.get("latency_ms"),
            "stop_reason": r.get("stop_reason"),
            "turns": r.get("n_turns"),
            "max_turns": r.get("max_turns"),
            "tool_call_counts": r.get("tool_call_counts") or {},
            "tokens": r.get("tokens") or {},
            "finish_reasons": r.get("finish_reasons"),
            "timeout_retries": 0,
        }
    return out


def project_config(pred: Dict[str, Any], gold_items: List[str], shape: str) -> Dict[str, Any]:
    """Per-config payload: light-weight; events live in the trajectory shards."""
    pred_items = extract_answer_items(pred.get("answer") or "")
    metrics = score_one(pred_items, gold_items)
    base: Dict[str, Any] = {
        "answer": pred.get("answer") or "",
        "pred_items": pred_items,
        "metrics": metrics,
        "model": pred.get("model"),
        "mode": pred.get("mode"),
        "system_prompt_file": pred.get("system_prompt_file"),
        "max_completion_tokens": pred.get("max_completion_tokens"),
        "latency_ms": pred.get("latency_ms"),
    }
    if shape == "single":
        base["finish_reason"] = pred.get("finish_reason")
        base["usage"] = pred.get("usage")
    else:  # agentic
        base["stop_reason"] = pred.get("stop_reason")
        base["n_turns"] = pred.get("turns")
        base["max_turns"] = pred.get("max_turns")
        base["tool_call_counts"] = pred.get("tool_call_counts") or {}
        base["tokens"] = pred.get("tokens") or {}
        base["finish_reasons"] = pred.get("finish_reasons")
        base["timeout_retries"] = pred.get("timeout_retries") or 0
    return base


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--scaffolds-root", type=Path, default=DEFAULT_SCAFFOLDS)
    ap.add_argument("--gold", type=Path, default=DEFAULT_GOLD)
    ap.add_argument(
        "--legacy-e2e-records-dir",
        "--e2e-records-dir",
        dest="legacy_e2e_records_dir",
        type=Path,
        default=DEFAULT_E2E_DIR / "records",
        help="Per-qid shards for the existing 2026-06-22 E2E run.",
    )
    ap.add_argument("--naive-predictions", type=Path, default=DEFAULT_NAIVE)
    ap.add_argument("--e2e-v3-predictions", type=Path, default=DEFAULT_E2E_V3)
    ap.add_argument("--e2e-v3-rawtext-predictions", type=Path, default=DEFAULT_E2E_V3_RAWTEXT)
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    args = ap.parse_args()

    cells_dir = args.scaffolds_root / "outputs" / "agentic_wiki_opentable"
    if not cells_dir.exists():
        print(f"error: cells dir not found: {cells_dir}", file=sys.stderr)
        return 2

    gold = load_jsonl(args.gold)
    print(f"gold: {len(gold)} qids from {args.gold}", file=sys.stderr)

    configs = [
        *CELL_CONFIGS,
        ("naive-search", f"file:{args.naive_predictions}", "agentic"),
        ("e2e (legacy)", f"shards:{args.legacy_e2e_records_dir}", "agentic"),
        ("e2e v3", f"file:{args.e2e_v3_predictions}", "agentic"),
        ("e2e v3 + rawtext", f"file:{args.e2e_v3_rawtext_predictions}", "agentic"),
    ]

    # Load all runs.
    cells: Dict[str, Tuple[str, Dict[str, Dict[str, Any]]]] = {}
    for label, source, shape in configs:
        if source.startswith("cell:"):
            basename = source[len("cell:"):]
            path = cells_dir / basename
            if not path.exists():
                print(f"warning: cell file missing, skipping: {path}", file=sys.stderr)
                continue
            d = load_jsonl(path, dataset_filter=DATASET)
            label_path = basename
        elif source.startswith("shards:"):
            records_dir = Path(source[len("shards:"):])
            d = load_e2e_shards(records_dir)
            if not d:
                print(f"warning: e2e shards dir empty/missing, skipping: {records_dir}", file=sys.stderr)
                continue
            label_path = str(records_dir)
        elif source.startswith("file:"):
            path = Path(source[len("file:"):])
            if not path.exists():
                print(f"warning: prediction file missing, skipping: {path}", file=sys.stderr)
                continue
            d = load_jsonl(path, dataset_filter=DATASET)
            label_path = str(path)
        else:
            print(f"warning: unknown source spec '{source}', skipping {label!r}", file=sys.stderr)
            continue
        cells[label] = (shape, d)
        print(f"  {label:18s} {len(d):4d} rows  ← {label_path}", file=sys.stderr)

    # Use the intersection of all cells × gold so every record has all configs.
    qids = set(gold)
    for label, (_, d) in cells.items():
        qids &= set(d)
    qids = sorted(qids)
    print(f"common qids: {len(qids)}", file=sys.stderr)
    if not qids:
        print("ERROR: no overlap across cells × gold", file=sys.stderr)
        return 1

    # Write per-qid shards + index.
    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_metrics: Dict[str, Dict[str, float]] = {label: {"f1": 0.0, "em": 0.0} for label, _ in cells.items()}
    for qid in qids:
        gold_row = gold[qid]
        gold_items = list(gold_row.get("answer_list") or [])
        configs_payload: Dict[str, Dict[str, Any]] = {}
        per_cell_brief: List[Dict[str, Any]] = []
        for label, (shape, d) in cells.items():
            pred = d[qid]
            cfg = project_config(pred, gold_items, shape)
            configs_payload[label] = cfg
            per_cell_brief.append({
                "label": label,
                "f1": cfg["metrics"]["f1"],
                "em": cfg["metrics"]["em"],
            })
            sum_metrics[label]["f1"] += cfg["metrics"]["f1"]
            sum_metrics[label]["em"] += cfg["metrics"]["em"]

        record = {
            "qid": qid,
            "question": gold_row.get("question_text") or (
                next((d[qid].get("question") for label, (_, d) in cells.items() if d[qid].get("question")), "")
            ),
            "dataset_origin": gold_row.get("dataset_origin"),
            "original_table_id": gold_row.get("original_table_id"),
            "sql": gold_row.get("sql"),
            "gold_answers": gold_items,
            "configs_order": [label for label, _ in cells.items()],
            "configs": configs_payload,
        }
        (rec_dir / f"{qid}.json").write_text(json.dumps(record, ensure_ascii=False))
        index_rows.append({
            "qid": qid,
            "question": record["question"],
            "dataset_origin": record["dataset_origin"],
            "n_gold": len(gold_items),
            "cells": per_cell_brief,
        })

    n = len(qids)
    summary = {
        "n": n,
        "configs": [label for label, _ in cells.items()],
        "mean_metrics": {
            label: {
                "mean_f1": round(s["f1"] / n, 4),
                "mean_em": round(s["em"] / n, 4),
            }
            for label, s in sum_metrics.items()
        },
    }
    (out_dir / "index.json").write_text(json.dumps({"meta": summary, "rows": index_rows}, ensure_ascii=False))

    print(f"\n✓ Wrote {out_dir}/index.json + {n} record shards", file=sys.stderr)
    for label in summary["mean_metrics"]:
        mm = summary["mean_metrics"][label]
        print(f"   {label:18s} mean F1 = {mm['mean_f1']*100:6.2f}   mean EM = {mm['mean_em']*100:6.2f}", file=sys.stderr)
    return 0


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