#!/usr/bin/env python3
"""Build full-run and comparison shards for the canonical PhantomWiki split.

Run from the viewer repository root:

    python3 scripts/build_runs.py

The response files may contain multiple datasets. They are streamed line by
line, only ``dataset == "phantom_wiki"`` records are retained, and c1/c2
recovery rows overlay their original full-run responses.
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
from pathlib import Path
from typing import Any, Iterator


ROOT = Path(__file__).resolve().parents[1]
DATASET = "phantom_wiki"
SCOPE = "size5000_seed1"
STRING_LIMIT = 8192
EVENTS_LIMIT_BYTES = 256 * 1024

DEFAULT_EVAL = ROOT / "eval_size5000_seed1.json"
DEFAULT_SOURCES = {
    "c1": {
        "label": "c1 Closed-book",
        "response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_full/full_closedbook/named-outputs/response/response"
        ),
        "recovery_response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_recover/cb_len4x/named-outputs/response/response"
        ),
        "judge": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_judge/cb_phantom_len4x/named-outputs/judged/judged"
        ),
    },
    "c2": {
        "label": "c2 With-docs",
        "response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_full/full_openbook/named-outputs/response/response"
        ),
        "recovery_response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_recover/ob_len4x/named-outputs/response/response"
        ),
        "judge": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_judge/ob_phantom_len4x/named-outputs/judged/judged"
        ),
    },
    "c6": {
        "label": "c6 Agentic-DCI",
        "response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_full/full_dci/named-outputs/response/response"
        ),
        "judge": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_judge/dci_phantom/named-outputs/judged/judged"
        ),
    },
    "naive": {
        "label": "Naive-search",
        "response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_full/full_naive_phantom/named-outputs/response/response"
        ),
        "judge": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/"
            "herb_phantom_judge/naive_phantom/named-outputs/judged/judged"
        ),
    },
    "e2e": {
        "label": "E2E v3",
        "response": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/e2e_runs/"
            "new-datasets-full-20260711/phantom_wiki/named-outputs/"
            "predictions/predictions"
        ),
        "judge": Path(
            "/home/azureuser/projects/information-scaffolds/outputs/e2e_runs/"
            "new-datasets-full-20260711/judges/phantom_wiki/named-outputs/"
            "judged/judged"
        ),
    },
    "e2e_rawtext": {
        "label": "E2E v3 + rawtext",
        "response": Path(
            "/tmp/viewer-overlay-phantom/predictions/named-outputs/"
            "predictions/predictions"
        ),
        "judge": Path(
            "/tmp/viewer-overlay-phantom/evaluated/named-outputs/"
            "canonical_evaluated/evaluated"
        ),
        "judge_format": "f1_csv",
        "score_label": "Mean answer F1",
    },
}


def stream_jsonl(path: Path) -> Iterator[dict[str, Any]]:
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            if not line.strip():
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: {exc}") from exc


def load_dataset_rows(path: Path) -> dict[str, dict[str, Any]]:
    rows: dict[str, dict[str, Any]] = {}
    for row in stream_jsonl(path):
        if row.get("dataset") != DATASET:
            continue
        qid = str(row.get("qid") or "")
        if not qid:
            raise ValueError(f"{path}: PhantomWiki row has no qid")
        if qid in rows:
            raise ValueError(f"{path}: duplicate PhantomWiki qid {qid}")
        rows[qid] = row
    return rows


def load_f1_rows(path: Path) -> dict[str, dict[str, Any]]:
    rows: dict[str, dict[str, Any]] = {}
    with path.open(encoding="utf-8", newline="") as handle:
        for row in csv.DictReader(handle):
            qid = str(row.get("qid") or "")
            if not qid:
                raise ValueError(f"{path}: F1 row has no qid")
            if qid in rows:
                raise ValueError(f"{path}: duplicate F1 qid {qid}")
            rows[qid] = {
                "qid": qid,
                "precision": float(row["precision"]),
                "recall": float(row["recall"]),
                "f1": float(row["f1"]),
                "n_pred": int(row["n_pred"]),
                "n_gold": int(row["n_gold"]),
                "n_hit": int(row["n_hit"]),
            }
    return rows


def cap_string(value: str, limit: int = STRING_LIMIT) -> str:
    if len(value) <= limit:
        return value
    marker = f"\n… [truncated from {len(value)} chars]"
    return value[: max(0, limit - len(marker))] + marker


def cap_value(value: Any) -> Any:
    if isinstance(value, str):
        return cap_string(value)
    if isinstance(value, list):
        return [cap_value(item) for item in value]
    if isinstance(value, tuple):
        return [cap_value(item) for item in value]
    if isinstance(value, dict):
        return {cap_string(str(key)): cap_value(item) for key, item in value.items()}
    return value


def compact_events(events: Any) -> list[dict[str, Any]]:
    if not isinstance(events, list):
        return []
    compacted: list[dict[str, Any]] = []
    used = 2
    original_bytes = len(
        json.dumps(events, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    )
    for index, raw in enumerate(events):
        if not isinstance(raw, dict):
            event = {"type": "event", "content": cap_value(raw)}
        else:
            event = {
                key: cap_value(raw[key])
                for key in ("type", "name", "input", "content")
                if raw.get(key) is not None
            }
        encoded = json.dumps(event, ensure_ascii=False, separators=(",", ":")).encode(
            "utf-8"
        )
        if used + len(encoded) + 1 > EVENTS_LIMIT_BYTES:
            compacted.append(
                {
                    "type": "truncated",
                    "content": (
                        f"{len(events) - index} events omitted; original events "
                        f"were {original_bytes} bytes and the compact limit is "
                        f"{EVENTS_LIMIT_BYTES} bytes."
                    ),
                }
            )
            break
        compacted.append(event)
        used += len(encoded) + 1
    return compacted


def nonempty_text(value: Any) -> str:
    return value.strip() if isinstance(value, str) else ""


def failure_reason(response: dict[str, Any] | None, answered: bool) -> str | None:
    if answered:
        return None
    if not response:
        return "missing_response"
    for key in ("stop_reason", "finish_reason"):
        if response.get(key):
            return str(response[key])
    for event in reversed(response.get("events") or []):
        if isinstance(event, dict) and event.get("type") == "error":
            return nonempty_text(event.get("content")) or "error"
    attempts = response.get("attempt_log") or []
    for attempt in reversed(attempts):
        if not isinstance(attempt, dict):
            continue
        if attempt.get("exc"):
            exc = str(attempt["exc"]).split(":", 1)[0]
            return f"exc:{exc}"
        if attempt.get("outcome") and attempt.get("outcome") != "success":
            return str(attempt["outcome"])
    return "unanswered"


def normalized_tokens(response: dict[str, Any] | None) -> dict[str, Any]:
    if not response:
        return {}
    tokens = response.get("tokens")
    if not isinstance(tokens, dict):
        tokens = response.get("usage")
    return cap_value(tokens) if isinstance(tokens, dict) else {}


def normalize_record(
    gold: dict[str, Any],
    response: dict[str, Any] | None,
    judge: dict[str, Any] | None,
) -> dict[str, Any]:
    prediction = nonempty_text((response or {}).get("answer"))
    answered = bool(prediction)
    is_f1 = isinstance((judge or {}).get("f1"), float)
    parsed = (judge or {}).get("parsed")
    parsed = parsed if isinstance(parsed, dict) else {}
    score = float(judge["f1"]) if is_f1 else None
    correct = bool(
        answered
        and (
            (score == 1.0 if score is not None else parsed.get("correct") is True)
        )
    )
    tool_counts = (response or {}).get("tool_call_counts")
    tool_counts = tool_counts if isinstance(tool_counts, dict) else {}
    metadata = gold.get("meta") if isinstance(gold.get("meta"), dict) else {}

    return cap_value(
        {
            "qid": gold["id"],
            "question": gold.get("question") or "",
            "gold": gold.get("answer") or [],
            "metadata": {
                "difficulty": metadata.get("Difficulty"),
                "type": metadata.get("Type"),
                "template": gold.get("template") or "",
                "prolog": gold.get("prolog") or [],
                "prolog_answer": gold.get("prolog_answer") or "",
                "supporting_titles": gold.get("supporting_titles") or [],
            },
            "prediction": prediction,
            "extracted_judge_answer": parsed.get("extracted_final_answer"),
            "answered": answered,
            "correct": correct,
            "score": score,
            "score_kind": "answer_f1" if score is not None else "binary",
            "score_details": (
                {
                    key: judge[key]
                    for key in ("precision", "recall", "f1", "n_pred", "n_gold", "n_hit")
                }
                if score is not None
                else None
            ),
            "judge": {
                "text": (
                    (
                        f"precision={judge['precision']:.4f}, "
                        f"recall={judge['recall']:.4f}, f1={judge['f1']:.4f}"
                    )
                    if score is not None
                    else (judge or {}).get("judge_text")
                ),
                "confidence": parsed.get("confidence"),
                "parse_error": parsed.get("parse_error"),
                "model": (judge or {}).get("judge_model"),
                "finish_reason": (judge or {}).get("finish_reason"),
            },
            "stop_reason": (response or {}).get("stop_reason"),
            "finish_reason": (response or {}).get("finish_reason"),
            "finish_reasons": (response or {}).get("finish_reasons") or [],
            "failure_reason": failure_reason(response, answered),
            "tokens": normalized_tokens(response),
            "turns": (response or {}).get("turns"),
            "tool_call_counts": tool_counts,
            "tool_calls": sum(
                value for value in tool_counts.values() if isinstance(value, int)
            ),
            "model": (response or {}).get("model"),
            "mode": (response or {}).get("mode"),
            "events": compact_events((response or {}).get("events")),
        }
    )


def status_for(record: dict[str, Any]) -> str:
    if not record["answered"]:
        return "missing"
    if isinstance(record.get("score"), (int, float)) and 0 < record["score"] < 1:
        return "partial"
    return "correct" if record["correct"] else "incorrect"


def run_summary(record: dict[str, Any]) -> dict[str, Any]:
    return {
        "qid": record["qid"],
        "question": record["question"],
        "gold": record["gold"],
        "answered": record["answered"],
        "correct": record["correct"],
        "score": record.get("score"),
        "score_kind": record.get("score_kind"),
        "status": status_for(record),
        "failure_reason": record["failure_reason"],
    }


def compare_projection(record: dict[str, Any], slot: str) -> dict[str, Any]:
    return {
        key: record[key]
        for key in (
            "prediction",
            "extracted_judge_answer",
            "answered",
            "correct",
            "score",
            "score_kind",
            "score_details",
            "judge",
            "stop_reason",
            "finish_reason",
            "finish_reasons",
            "failure_reason",
            "tokens",
            "turns",
            "tool_call_counts",
            "tool_calls",
            "model",
            "mode",
        )
    } | {"events_path": f"runs/{slot}/records/{record['qid']}.json"}


def write_json(path: Path, value: Any, *, indent: int | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, indent=indent)


def clear_outputs(output_root: Path, compare_root: Path) -> None:
    if output_root.exists():
        shutil.rmtree(output_root)
    if compare_root.exists():
        shutil.rmtree(compare_root)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--eval-file", type=Path, default=DEFAULT_EVAL)
    parser.add_argument("--output-root", type=Path, default=ROOT / "runs")
    parser.add_argument("--compare-root", type=Path, default=ROOT / "compare")
    for slot, source in DEFAULT_SOURCES.items():
        parser.add_argument(
            f"--{slot}-response", type=Path, default=source["response"]
        )
        if source.get("recovery_response"):
            parser.add_argument(
                f"--{slot}-recovery-response",
                type=Path,
                default=source["recovery_response"],
            )
        parser.add_argument(f"--{slot}-judge", type=Path, default=source["judge"])
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    gold_rows = json.loads(args.eval_file.read_text(encoding="utf-8"))
    if len(gold_rows) != 500:
        raise ValueError(f"{args.eval_file}: expected 500 eval rows, got {len(gold_rows)}")
    qids = [str(row["id"]) for row in gold_rows]
    if len(qids) != len(set(qids)):
        raise ValueError(f"{args.eval_file}: duplicate eval qids")
    qid_set = set(qids)

    clear_outputs(args.output_root, args.compare_root)
    records_by_slot: dict[str, dict[str, dict[str, Any]]] = {}
    manifest_slots = []

    for slot, source in DEFAULT_SOURCES.items():
        response_path = getattr(args, f"{slot}_response")
        recovery_response_path = getattr(
            args, f"{slot}_recovery_response", None
        )
        judge_path = getattr(args, f"{slot}_judge")
        responses = load_dataset_rows(response_path)
        if recovery_response_path:
            responses.update(load_dataset_rows(recovery_response_path))
        judges = (
            load_f1_rows(judge_path)
            if source.get("judge_format") == "f1_csv"
            else load_dataset_rows(judge_path)
        )
        extra_response_qids = set(responses) - qid_set
        if extra_response_qids:
            raise ValueError(
                f"{slot}: response file contains "
                f"{len(extra_response_qids)} qids outside eval"
            )
        if not set(judges).issubset(qid_set):
            raise ValueError(f"{slot}: judge file contains qids outside eval")

        slot_records: dict[str, dict[str, Any]] = {}
        summaries = []
        records_dir = args.output_root / slot / "records"
        for gold in gold_rows:
            qid = str(gold["id"])
            record = normalize_record(gold, responses.get(qid), judges.get(qid))
            slot_records[qid] = record
            summaries.append(run_summary(record))
            write_json(records_dir / f"{qid}.json", record)

        answered = sum(record["answered"] for record in slot_records.values())
        correct = sum(record["correct"] for record in slot_records.values())
        canonical_score = sum(
            record["score"]
            if isinstance(record.get("score"), (int, float))
            else float(record["correct"])
            for record in slot_records.values()
        ) / len(qids)
        index = {
            "slot": slot,
            "label": source["label"],
            "scope": SCOPE,
            "total": len(qids),
            "answered": answered,
            "correct": correct,
            "coverage": answered / len(qids),
            "score": canonical_score,
            "score_label": source.get("score_label", "Canonical score"),
            "score_detail": (
                f"{correct} perfect · "
                f"{sum((record.get('score') or 0) > 0 for record in slot_records.values())} "
                "with answer overlap"
                if source.get("judge_format") == "f1_csv"
                else f"{correct} correct"
            ),
            "records": summaries,
        }
        write_json(args.output_root / slot / "index.json", index)
        records_by_slot[slot] = slot_records
        manifest_slots.append(
            {
                key: index[key]
                for key in (
                    "slot",
                    "label",
                    "scope",
                    "total",
                    "answered",
                    "correct",
                    "coverage",
                    "score",
                    "score_label",
                    "score_detail",
                )
            }
            | {
                "index": f"runs/{slot}/index.json",
                "records": f"runs/{slot}/records",
                "response_source": str(response_path),
                "recovery_response_source": (
                    str(recovery_response_path) if recovery_response_path else None
                ),
                "judge_source": str(judge_path),
            }
        )
        print(
            f"{slot}: score={canonical_score:.2%}, perfect/correct={correct}/{len(qids)}, "
            f"answered={answered}/{len(qids)} ({answered / len(qids):.2%})"
        )

    compare_summaries = []
    for gold in gold_rows:
        qid = str(gold["id"])
        run_records = {
            slot: records_by_slot[slot][qid] for slot in DEFAULT_SOURCES
        }
        correct_flags = [record["correct"] for record in run_records.values()]
        answered_flags = [record["answered"] for record in run_records.values()]
        summary = {
            "qid": qid,
            "question": gold.get("question") or "",
            "correct": {
                slot: record["correct"] for slot, record in run_records.items()
            },
            "answered": {
                slot: record["answered"] for slot, record in run_records.items()
            },
            "correct_count": sum(correct_flags),
            "answered_count": sum(answered_flags),
            "disagreement": len(set(correct_flags)) > 1,
            "any_missing": not all(answered_flags),
            "only_e2e_correct": (
                run_records["e2e"]["correct"]
                and all(
                    not record["correct"]
                    for slot, record in run_records.items()
                    if slot != "e2e"
                )
            ),
        }
        compare_summaries.append(summary)
        compare_record = {
            "qid": qid,
            "question": gold.get("question") or "",
            "gold": gold.get("answer") or [],
            "metadata": run_records["c1"]["metadata"],
            "runs": {
                slot: compare_projection(record, slot)
                for slot, record in run_records.items()
            },
        }
        write_json(args.compare_root / "records" / f"{qid}.json", compare_record)

    write_json(
        args.compare_root / "index.json",
        {
            "scope": SCOPE,
            "total": len(qids),
            "slots": list(DEFAULT_SOURCES),
            "records": compare_summaries,
        },
    )
    write_json(
        args.output_root / "manifest.json",
        {
            "dataset": DATASET,
            "scope": SCOPE,
            "scope_label": "depth_20_size_5000_seed_1",
            "eval_source": str(args.eval_file),
            "total": len(qids),
            "slots": manifest_slots,
            "compare_index": "compare/index.json",
            "scoring": (
                "Run-specific canonical score over all 500 eval qids; "
                "missing or unanswered is zero"
            ),
        },
        indent=2,
    )
    print(f"compare: {len(compare_summaries)} records")


if __name__ == "__main__":
    main()