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

For each of the 100-qid wiki_opentable_smoke subset, joins the model's
``answer`` from all six smoke ``outputs/wiki_opentable_smoke/baselines/cell*``
configurations PLUS the e2e v1 trajectory (read from the local
``trajectories_e2e/records/<qid>.json`` shards), parses each via the
canonical semicolon-Exact-Answer rule, and scores against the gold
``answer_list`` (from the local ``unified/records/`` smoke shards) with
the canonical ``wiki_opentable_adapter``-style set-based F1 / EM / P / R.

The seven configs:

  - closed-book        (cell1_closedbook)         — single-shot, no context
  - with-docs          (cell2_withdocs)           — single-shot, gold docs
  - with-structures    (cell3_withstructures)     — single-shot, eval structures
  - structure per q    (cell4_agentic_a)          — formerly Baseline A · per-qid scaffolds
  - structure per ds   (cell5_agentic_b)          — formerly Baseline B · flat structure corpus
  - DCI                (cell6_agentic_c)          — formerly Baseline C · rawtext corpus
  - e2e v1             (trajectories_e2e shards)  — e2e_pipeline run "slug-desc-dev-v0"
                                                    (on-the-fly per-shape extraction; framework
                                                    substring-eval reported 97/100)

Output layout (mirrors the parent open-wikitable-viewer/compare/):
    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_smoke"
HERE = Path(__file__).resolve().parent
REPO = HERE.parent
DEFAULT_SCAFFOLDS = Path(os.environ.get(
    "SCAFFOLDS_ROOT", "/home/azureuser/projects/information-scaffolds"
))
# Gold defaults to the local per-qid unified shards (committed in this repo).
DEFAULT_GOLD = REPO / "unified" / "records"
DEFAULT_OUT = REPO / "compare"
DEFAULT_E2E_DIR = REPO / "trajectories_e2e"
DEFAULT_E2E_V2_DIR = REPO / "trajectories_e2e_v2"

# (label, source, shape)
# source: "cell:<dir>" loads from outputs/wiki_opentable_smoke/baselines/<dir>/named-outputs/response/response;
#         "e2e_shards:<dir>" loads from <dir>/records/*.json (defaults below if no path given).
# shape:  "single" / "agentic" (drives which extra metadata is kept).
CONFIGS: List[Tuple[str, str, str]] = [
    ("closed-book",      "cell:cell1_closedbook",                "single"),
    ("with-docs",        "cell:cell2_withdocs",                  "single"),
    ("with-structures",  "cell:cell3_withstructures",            "single"),
    ("structure per q",  "cell:cell4_agentic_a",                 "agentic"),
    ("structure per ds", "cell:cell5_agentic_b",                 "agentic"),
    ("DCI",              "cell:cell6_agentic_c",                 "agentic"),
    ("e2e v1",           "e2e_shards:trajectories_e2e",          "agentic"),
    ("e2e v2",           "e2e_shards:trajectories_e2e_v2",       "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]]:
    """Load a JSONL keyed by qid.

    ``dataset_filter`` is treated as a hint: rows whose ``dataset`` doesn't
    match are skipped, but ``"wiki_opentable_smoke"`` / ``"wiki_opentable_dev"``
    / ``"wiki_opentable"`` are all considered the same canonical set (different
    stamps appear depending on which AML job produced the row).
    """
    accepted = None
    if dataset_filter is not None:
        accepted = {dataset_filter, "wiki_opentable", "wiki_opentable_smoke", "wiki_opentable_dev"}
    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 accepted is not None and d.get("dataset") not in accepted:
                continue
            out[str(d["qid"])] = d
    return out


def load_gold(path: Path) -> Dict[str, Dict[str, Any]]:
    """Load gold answers, accepting either a per-qid shard dir or a JSONL.

    Synthesises ``answer_list`` / ``question_text`` from the smoke unified
    shape (``answers``, ``question``) so downstream scoring/projection don't
    need to special-case it.
    """
    def _norm(d: Dict[str, Any]) -> Dict[str, Any]:
        if "answer_list" not in d and "answers" in d:
            d = dict(d)
            d["answer_list"] = list(d.get("answers") or [])
        if "question_text" not in d and "question" in d:
            d.setdefault("question_text", d["question"])
        return d

    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] = _norm(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"])] = _norm(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("--e2e-records-dir", type=Path, default=DEFAULT_E2E_DIR / "records",
                    help="Per-qid e2e v1 shards (defaults to ./trajectories_e2e/records/).")
    ap.add_argument("--e2e-v2-records-dir", type=Path, default=DEFAULT_E2E_V2_DIR / "records",
                    help="Per-qid e2e v2 shards (defaults to ./trajectories_e2e_v2/records/).")
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    args = ap.parse_args()

    cells_root = args.scaffolds_root / "outputs" / "wiki_opentable_smoke" / "baselines"
    if not cells_root.exists():
        print(f"error: smoke baselines dir not found: {cells_root}", file=sys.stderr)
        return 2

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

    # Load all configs.
    # Map source-dir-stem → records dir override (lets CLI flags retarget shards).
    e2e_dir_overrides = {
        "trajectories_e2e":    args.e2e_records_dir,
        "trajectories_e2e_v2": args.e2e_v2_records_dir,
    }
    cells: Dict[str, Tuple[str, Dict[str, Dict[str, Any]]]] = {}
    for label, source, shape in CONFIGS:
        if source.startswith("cell:"):
            cell_dir = source[len("cell:"):]
            path = cells_root / cell_dir / "named-outputs" / "response" / "response"
            if not path.exists():
                print(f"warning: cell response missing, skipping: {path}", file=sys.stderr)
                continue
            d = load_jsonl(path, dataset_filter=DATASET)
            label_path = f"{cell_dir}/named-outputs/response/response"
        elif source.startswith("e2e_shards:"):
            dir_stem = source[len("e2e_shards:"):]
            records_dir = e2e_dir_overrides.get(dir_stem, REPO / dir_stem / "records")
            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.relative_to(REPO))
        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())