#!/usr/bin/env python3
"""Diagnose Qwen30 Fixed-PD TTFT ranking through admission-state accounting."""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
from pathlib import Path
from statistics import fmean
from typing import Any


CONFIGS = {
    "tp2_mns64": {"tp": 2, "mns": 64},
    "tp4_mns16": {"tp": 4, "mns": 16},
    "tp4_mns32": {"tp": 4, "mns": 32},
    "tp4_mns64": {"tp": 4, "mns": 64},
}
REAL_LOG_PATTERN = re.compile(
    r"Avg prompt throughput: (?P<prompt>[0-9.]+) tokens/s, "
    r"Avg generation throughput: (?P<generation>[0-9.]+) tokens/s, "
    r"Running: (?P<running>[0-9]+) reqs, Waiting: (?P<waiting>[0-9]+) reqs"
)


def atomic_write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(text)
    os.replace(temporary, path)


def load_request_metrics(sim_root: Path, config: str, tp: int) -> list[dict[str, str]]:
    matches = sorted(
        (sim_root / "runs" / config / f"tp{tp}" / "metrics").glob(
            "**/request_metrics.csv"
        )
    )
    if len(matches) != 1:
        raise ValueError(f"expected one request_metrics.csv for {config}: {matches}")
    with matches[0].open(newline="") as source:
        rows = list(csv.DictReader(source))
    if len(rows) != 257:
        raise ValueError(f"expected 257 simulator requests for {config}, got {len(rows)}")
    return rows


def mean_column(rows: list[dict[str, str]], name: str) -> float:
    values = [float(row[name]) for row in rows]
    if len(values) != len(rows):
        raise ValueError(f"missing values in {name}")
    return fmean(values)


def real_queue_summary(real_root: Path, config: str) -> dict[str, float | int]:
    rows: list[tuple[int, int]] = []
    for trial in ("trial1", "trial2", "trial3"):
        path = real_root / config / trial / "logs/server.log"
        trial_rows = []
        for line in path.read_text(errors="replace").splitlines():
            match = REAL_LOG_PATTERN.search(line)
            if match is not None and float(match.group("generation")) > 0:
                trial_rows.append(
                    (int(match.group("running")), int(match.group("waiting")))
                )
        if len(trial_rows) < 2:
            raise ValueError(f"insufficient periodic queue samples: {path}")
        rows.extend(trial_rows[1:])
    return {
        "samples": len(rows),
        "running_mean": fmean(row[0] for row in rows),
        "running_max": max(row[0] for row in rows),
        "waiting_mean": fmean(row[1] for row in rows),
        "waiting_max": max(row[1] for row in rows),
    }


def analyze(campaign_root: Path) -> dict[str, Any]:
    sim_root = campaign_root / "sim/fixed-pd"
    real_root = campaign_root / "real/fixed-pd/real"
    comparison = json.loads(
        (campaign_root / "analysis/fixed-pd/comparison.json").read_text()
    )
    real_audit = json.loads(
        (campaign_root / "analysis/fixed-pd/real-audit.json").read_text()
    )
    cells = {}
    for config, knobs in CONFIGS.items():
        tp = knobs["tp"]
        mns = knobs["mns"]
        rate = 1.125 * tp
        rows = load_request_metrics(sim_root, config, tp)
        sim = {
            "ttft_ms": mean_column(rows, "ttft"),
            "first_scheduling_delay_ms": mean_column(
                rows, "request_first_scheduling_delay"
            ),
            "queue_free_ttft_ms": mean_column(
                rows, "prefill_execution_plus_preemption"
            ),
            "request_execution_ms": mean_column(rows, "request_execution_time"),
            "tpot_ms": mean_column(rows, "tpot"),
        }
        sim["required_slots"] = rate * sim["request_execution_ms"] / 1000.0
        sim["slot_margin"] = mns - sim["required_slots"]
        sim["ttft_queue_fraction"] = (
            sim["first_scheduling_delay_ms"] / sim["ttft_ms"]
        )
        real_metrics = real_audit["configs"][config]["metrics"]
        real = {
            "ttft_ms": real_metrics["ttft_ms"]["pooled_mean_ms"],
            "e2e_ms": real_metrics["e2e_ms"]["pooled_mean_ms"],
            "tpot_ms": real_metrics["tpot_ms"]["pooled_mean_ms"],
            "periodic_queue": real_queue_summary(real_root, config),
        }
        real["required_slots_upper_bound"] = rate * real["e2e_ms"] / 1000.0
        real["slot_margin_lower_bound"] = mns - real["required_slots_upper_bound"]
        cells[config] = {
            "tp": tp,
            "mns": mns,
            "arrival_rate_requests_per_s": rate,
            "simulator": sim,
            "real": real,
            "tpot_overprediction_ratio": sim["tpot_ms"] / real["tpot_ms"],
        }

    tp2 = cells["tp2_mns64"]
    tp4 = cells["tp4_mns32"]
    contrasts = {
        "observed_sim_ttft_tp4_minus_tp2_ms": (
            tp4["simulator"]["ttft_ms"] - tp2["simulator"]["ttft_ms"]
        ),
        "queue_free_sim_ttft_tp4_minus_tp2_ms": (
            tp4["simulator"]["queue_free_ttft_ms"]
            - tp2["simulator"]["queue_free_ttft_ms"]
        ),
        "real_ttft_tp4_minus_tp2_ms": (
            tp4["real"]["ttft_ms"] - tp2["real"]["ttft_ms"]
        ),
    }
    if not (
        contrasts["observed_sim_ttft_tp4_minus_tp2_ms"] > 0
        and contrasts["queue_free_sim_ttft_tp4_minus_tp2_ms"] < 0
        and contrasts["real_ttft_tp4_minus_tp2_ms"] < 0
        and cells["tp2_mns64"]["simulator"]["slot_margin"] > 0
        and cells["tp4_mns64"]["simulator"]["slot_margin"] < 0
        and cells["tp4_mns16"]["real"]["periodic_queue"]["waiting_max"] == 0
    ):
        raise ValueError("admission-amplification invariants did not hold")
    return {
        "schema": "qwen30-fixed-pd-ttft-admission-diagnosis-v1",
        "status": "PASS",
        "workload": "Qwen3-30B Fixed-PD 4096->256 at 1.125 req/s/GPU",
        "selection": comparison["selection"]["ttft_ms:mean"],
        "cells": cells,
        "contrasts": contrasts,
        "verdict": (
            "Frontier overpredicts decode service time, so TP4's higher global arrival "
            "rate crosses the MNS admission cap only in simulation. The resulting first-"
            "scheduling queue reverses TTFT ranking; without that queue Frontier itself "
            "predicts the correct TP4 winner."
        ),
        "boundary": (
            "This identifies where the selection error is created, but the frozen real "
            "logs do not attribute service-time overprediction to one operator."
        ),
    }


def markdown(result: dict[str, Any]) -> str:
    lines = [
        "# Qwen30 Fixed-PD TTFT admission diagnosis",
        "",
        f"**Verdict:** {result['verdict']}",
        "",
        "| Config | Sim TPOT / real | Sim required slots / MNS | Real slots upper bound / MNS | Sim first-schedule wait | Real waiting max |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for config, cell in result["cells"].items():
        lines.append(
            f"| {config} | {cell['tpot_overprediction_ratio']:.2f}x "
            f"| {cell['simulator']['required_slots']:.1f}/{cell['mns']} "
            f"| {cell['real']['required_slots_upper_bound']:.1f}/{cell['mns']} "
            f"| {cell['simulator']['first_scheduling_delay_ms']:.1f} ms "
            f"| {cell['real']['periodic_queue']['waiting_max']} |"
        )
    contrast = result["contrasts"]
    lines.extend(
        [
            "",
            "## TP4/MNS32 minus TP2/MNS64 TTFT",
            "",
            f"- Observed Frontier: {contrast['observed_sim_ttft_tp4_minus_tp2_ms']:+.1f} ms.",
            f"- Frontier without first-scheduling queue: {contrast['queue_free_sim_ttft_tp4_minus_tp2_ms']:+.1f} ms.",
            f"- Real hardware: {contrast['real_ttft_tp4_minus_tp2_ms']:+.1f} ms.",
            "",
            f"**Boundary:** {result['boundary']}",
            "",
        ]
    )
    return "\n".join(lines)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--campaign-root", type=Path, required=True)
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    result = analyze(args.campaign_root.resolve())
    atomic_write(args.json_output, json.dumps(result, indent=2, sort_keys=True) + "\n")
    atomic_write(args.markdown_output, markdown(result))
    print(json.dumps({"status": result["status"], "verdict": result["verdict"]}))


if __name__ == "__main__":
    main()