#!/usr/bin/env python3
"""Compact code-trace Frontier rho calibration into a reviewable artifact."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from pathlib import Path


CELL_RE = re.compile(r"r(?P<rho>[0-9p]+)-tp(?P<tp>[24])-v3")
ATTENTION_FLAG = "--random_forrest_execution_time_predictor_config_atten_input_file"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sim-root", type=Path, required=True)
    parser.add_argument("--profile-manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    args = parse_args()
    profile = json.loads(args.profile_manifest.read_text())
    expected_profile_sha = profile["output_sha256"]
    cells = []
    for root in sorted(args.sim_root.iterdir()):
        match = CELL_RE.fullmatch(root.name)
        if match is None or not (root / "summary.json").is_file():
            continue
        manifest = json.loads((root / "manifest.json").read_text())
        observed_profile_sha = manifest.get("attention_profile_sha256")
        if observed_profile_sha is None:
            argv = manifest["argv"]
            profile_path = Path(argv[argv.index(ATTENTION_FLAG) + 1])
            observed_profile_sha = sha256(profile_path)
        if observed_profile_sha != expected_profile_sha:
            raise ValueError(
                f"{root}: attention profile {observed_profile_sha} "
                f"!= profile-v6 {expected_profile_sha}"
            )
        summary = json.loads((root / "summary.json").read_text())
        drain_fraction = (
            summary["drain"]["tail_after_last_arrival_s"] / summary["duration_s"]
        )
        cells.append(
            {
                "topology": f"tp{match.group('tp')}_mns16",
                "rho": float(match.group("rho").replace("p", ".")),
                "requests": summary["requests"],
                "offered_load": summary["offered_load"],
                "waiting_ms": summary["latency_ms"]["waiting"],
                "ttft_ms": summary["latency_ms"]["ttft"],
                "tpot_ms": summary["latency_ms"]["tpot"],
                "decode_batch": summary["decode_batch"],
                "drain_tail_s": summary["drain"]["tail_after_last_arrival_s"],
                "drain_fraction": drain_fraction,
                "subcritical_gate": drain_fraction <= 0.1,
                "prefix_cache_hit_ratio": summary["prefix_cache"]["hit_ratio"],
                "trace_sha256": manifest["trace_sha256"],
            }
        )
    if not cells:
        raise ValueError(f"no completed v3 calibration cells below {args.sim_root}")
    payload = {
        "schema": "frontier-code-trace-calibration-summary-v1",
        "attention_profile_sha256": expected_profile_sha,
        "subcritical_rule": "drain tail <= 10% of the 3660s arrival window",
        "cells": cells,
    }
    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({"cells": len(cells), "output": str(args.output)}, sort_keys=True))


if __name__ == "__main__":
    main()