#!/usr/bin/env python3
"""Check fresh-process repeat stability for the code long-context grid."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--first", type=Path, nargs="+", required=True)
    parser.add_argument("--second", type=Path, nargs="+", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max-relative-difference", type=float, default=0.05)
    return parser.parse_args()


def load(paths: list[Path]) -> dict[tuple[int, str], dict]:
    rows: dict[tuple[int, str], dict] = {}
    for path in paths:
        payload = json.loads(path.read_text())
        for row in payload["rows"]:
            if row.get("error"):
                raise ValueError(
                    f"{path}: failed profile row {row['config']['batch_spec']}"
                )
            key = (
                int(row["tensor_parallel_size"]),
                str(row["config"]["batch_spec"]),
            )
            if key in rows:
                raise ValueError(f"duplicate row {key}")
            rows[key] = row
    return rows


def main() -> None:
    args = parse_args()
    first = load(args.first)
    second = load(args.second)
    if first.keys() != second.keys():
        raise ValueError(
            f"repeat key mismatch: first_only={sorted(first.keys()-second.keys())}, "
            f"second_only={sorted(second.keys()-first.keys())}"
        )
    comparisons = []
    for key in sorted(first):
        left = float(first[key]["mean_time"])
        right = float(second[key]["mean_time"])
        relative = abs(left - right) / ((left + right) / 2)
        comparisons.append(
            {
                "tp": key[0],
                "batch_spec": key[1],
                "first_mean_s": left,
                "second_mean_s": right,
                "relative_difference": relative,
                "pass": relative <= args.max_relative_difference,
            }
        )
    maximum = max(item["relative_difference"] for item in comparisons)
    payload = {
        "schema": "frontier-code-longctx-repeat-check-v1",
        "threshold": args.max_relative_difference,
        "maximum_relative_difference": maximum,
        "status": "PASS" if maximum <= args.max_relative_difference else "FAIL",
        "comparisons": comparisons,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"status": payload["status"], "max": maximum}))
    if payload["status"] != "PASS":
        raise SystemExit(1)


if __name__ == "__main__":
    main()