#!/usr/bin/env python3
"""Freeze the EXP-TP2-PREFILL-SERVING entry audit from simulator ledgers."""

from __future__ import annotations

import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
REPLAY = REPO / "runs/frontier-attn-structured-v0/replay"
REAL_CHUNK_MS = {2: 410.0, 4: 231.0}


def find_one(root: Path, name: str) -> Path:
    matches = list(root.rglob(name))
    if len(matches) != 1:
        raise ValueError(f"expected one {name} below {root}: {matches}")
    return matches[0]


def chunk_zero(root: Path) -> dict[str, float | int]:
    ledger = find_one(root, "frontier_stage_batch_ledger.jsonl")
    chunks = []
    processed: dict[str, int] = {}
    with ledger.open() as stream:
        for line in stream:
            row = json.loads(line)
            request_ids = row.get("request_ids") or []
            request_tokens = row.get("request_num_tokens") or []
            if len(request_ids) != 1 or request_tokens != [8192]:
                continue
            request_id = request_ids[0]
            before = processed.get(request_id, 0)
            processed[request_id] = before + 8192
            if before == 0:
                components = row["execution_time"]["component_ledger_ms"]
                chunks.append(
                    {
                        "total_ms": row["execution_time"]["total_time_ms"],
                        "attention_prefill_ms": components[
                            "attention_prefill_execution_time"
                        ],
                        "moe_grouped_gemm_ms": components[
                            "moe_grouped_gemm_time"
                        ],
                    }
                )
            if len(chunks) == 9:
                break
    if len(chunks) != 9:
        raise ValueError(f"expected 9 initial q8k chunks, got {len(chunks)}")
    return {
        name: sum(float(row[name]) for row in chunks) / len(chunks)
        for name in chunks[0]
    } | {"samples": len(chunks)}


def main() -> None:
    cells = {}
    for tp, label in ((2, "tp2_rho0p0025"), (4, "tp4_rho0p0025")):
        measured = chunk_zero(REPLAY / label)
        real = REAL_CHUNK_MS[tp]
        measured["real_total_ms"] = real
        measured["total_bias"] = (measured["total_ms"] - real) / real
        cells[f"tp{tp}"] = measured
    tp2 = cells["tp2"]
    required_moe = (
        tp2["moe_grouped_gemm_ms"]
        + tp2["real_total_ms"]
        - tp2["total_ms"]
    )
    payload = {
        "schema": "frontier-tp2-prefill-serving-entry-v1",
        "cells": cells,
        "tp2_required_moe_if_residual_is_all_moe_ms": required_moe,
        "tp2_required_moe_shift": (
            required_moe / tp2["moe_grouped_gemm_ms"] - 1
        ),
        "entry_gate": abs(tp2["total_bias"]) >= 0.10,
    }
    results = ROOT / "results"
    results.mkdir(exist_ok=True)
    (results / "entry-audit.json").write_text(json.dumps(payload, indent=2))
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()