#!/usr/bin/env python3
"""Pair P1 engine intervals with detailed Frontier state summaries."""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Any


HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

from common_state import load_jsonl, numeric, residual, summarize_engine  # noqa: E402


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    temporary.replace(path)


def adjudicated_labels(path: Path) -> dict[tuple[str, str], bool]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {
        (row["cell"], row["level"]): bool(row["adjudicated_feasible"])
        for row in payload["examples"]
    }


def state_runs(root: Path) -> list[Path]:
    candidates = list(root.glob("*/result.json"))
    if (root / "result.json").is_file():
        candidates.append(root / "result.json")
    return sorted(
        path
        for path in candidates
        if json.loads(path.read_text(encoding="utf-8")).get("schema")
        == "telemetry-residual-frontier-state-result-v1"
    )


def execute(args: argparse.Namespace) -> dict[str, Any]:
    labels = adjudicated_labels(args.pilot_metrics)
    examples = []
    red_flags = []
    for state_result_path in state_runs(args.sim_state_root):
        run_root = state_result_path.parent
        manifest = json.loads((run_root / "run_manifest.json").read_text(encoding="utf-8"))
        cell = manifest["entry"]["cell"]
        role = manifest["entry"]["role"]
        level = "low" if role.startswith("low") else "high"
        real_run = args.real_root / "cells" / cell / f"{level}-rep1"
        real_result = json.loads((real_run / "result.json").read_text(encoding="utf-8"))
        stream_paths = sorted((args.real_root / "cells" / cell / "opprof").glob("*.jsonl"))
        if len(stream_paths) != 1:
            raise ValueError(f"expected one engine stream for {cell}, found {len(stream_paths)}")
        engine = summarize_engine(
            load_jsonl(stream_paths[0]),
            start_ns=int(real_result["interval"]["start_mono_ns"]),
            end_ns=int(real_result["interval"]["end_mono_ns"]),
            request_count=int(real_result["selection"]["count"]),
        )
        simulator = json.loads((run_root / "common-state.json").read_text(encoding="utf-8"))
        scorer = json.loads((run_root / "scorer_output.json").read_text(encoding="utf-8"))
        if engine["interval"]["request_count"] != simulator["interval"]["request_count"]:
            red_flags.append(f"request_count_mismatch:{cell}:{role}")
        difference = residual(engine, simulator)
        if any(not math.isfinite(float(value)) for value in difference["values"].values()):
            red_flags.append(f"nonfinite_residual:{cell}:{role}")
        real_feasible = labels[(cell, level)]
        sim_feasible = bool(scorer["slo"]["feasible"])
        examples.append(
            {
                "cell": cell,
                "role": role,
                "level": level,
                "tp": int(cell.split("_")[0][2:]),
                "mns": int(cell.split("_")[1][3:]),
                "request_count": engine["interval"]["request_count"],
                "real_feasible": real_feasible,
                "sim_feasible": sim_feasible,
                "simulator_error": sim_feasible != real_feasible,
                "simulator_false_feasible": sim_feasible and not real_feasible,
                "real_pass_rate_rep1": float(real_result["pass_rate"]),
                "offered_req_s": float(real_result["selection"]["offered_req_s"]),
                "offered_req_s_per_gpu": float(
                    real_result["selection"]["offered_req_s_per_gpu"]
                ),
                "sim_pass_rate": float(scorer["slo"]["pass_rate"]),
                "pass_rate_residual": float(real_result["pass_rate"])
                - float(scorer["slo"]["pass_rate"]),
                "engine": engine,
                "simulator": simulator,
                "state_residual": difference,
                "paths": {
                    "real_result": str((real_run / "result.json").resolve()),
                    "engine_stream": str(stream_paths[0].resolve()),
                    "simulator_result": str(state_result_path.resolve()),
                },
            }
        )
    if not examples:
        red_flags.append("no_state_examples")
    request_counts = [int(row["request_count"]) for row in examples]
    pass_rate_residuals = [float(row["pass_rate_residual"]) for row in examples]
    result = {
        "schema": "telemetry-residual-p1-state-pairs-v1",
        "status": "PASS" if not red_flags else "STOP",
        "scope": "P1 development plumbing; not held-out contribution evidence",
        "examples": examples,
        "red_flags": red_flags,
        "sanity": {
            "n": len(examples),
            "request_count": numeric(request_counts) if request_counts else None,
            "pass_rate_residual": numeric(pass_rate_residuals)
            if pass_rate_residuals
            else None,
            "simulator_errors": sum(bool(row["simulator_error"]) for row in examples),
            "invariants": {
                "request_counts_match": not any(
                    flag.startswith("request_count_mismatch") for flag in red_flags
                ),
                "finite_residuals": not any(
                    flag.startswith("nonfinite_residual") for flag in red_flags
                ),
                "ratios_bounded": all(
                    0.0 <= row["real_pass_rate_rep1"] <= 1.0
                    and 0.0 <= row["sim_pass_rate"] <= 1.0
                    for row in examples
                ),
                "per_config_not_identical": len(set(pass_rate_residuals)) > 1
                if len(pass_rate_residuals) > 1
                else None,
            },
        },
    }
    atomic_json(args.output, result)
    if result["status"] != "PASS":
        raise RuntimeError(red_flags)
    return result


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--real-root", type=Path, required=True)
    result.add_argument("--sim-state-root", type=Path, required=True)
    result.add_argument("--pilot-metrics", type=Path, required=True)
    result.add_argument("--output", type=Path, required=True)
    return result


def main() -> None:
    result = execute(parser().parse_args())
    print(
        json.dumps(
            {
                "status": result["status"],
                "examples": len(result["examples"]),
                "simulator_errors": result["sanity"]["simulator_errors"],
                "sanity": result["sanity"],
                "red_flags": result["red_flags"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()