#!/usr/bin/env python3
"""Audit one completed Qwen30 latency-selection real surface.

This reader intentionally treats a prefill-only case as a four-objective
surface: TTFT/E2E mean and p90.  TPOT is omitted rather than coerced to zero.
"""

from __future__ import annotations

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


CONFIGS = tuple(f"tp{tp}_mns{mns}" for tp in (1, 2, 4) for mns in (8, 16, 32, 64))
TRIALS = (1, 2, 3)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--case-root", type=Path, required=True)
    parser.add_argument(
        "--traces-root",
        type=Path,
        help="Frozen TP-normalized traces; defaults to <case-root>/traces.",
    )
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    return parser.parse_args()


def nearest_rank(values: list[float], fraction: float) -> float:
    if not values:
        raise ValueError("cannot calculate a percentile of an empty metric")
    return sorted(values)[math.ceil(len(values) * fraction) - 1]


def finite(value: Any, field: str) -> float:
    if not isinstance(value, (int, float)) or isinstance(value, bool):
        raise ValueError(f"{field} is not numeric: {value!r}")
    value = float(value)
    if not math.isfinite(value) or value < 0:
        raise ValueError(f"{field} is invalid: {value!r}")
    return value


def read_manifest(traces_root: Path, tp: int) -> dict[str, Any]:
    path = traces_root / f"tp{tp}" / "public" / "manifest.json"
    manifest = json.loads(path.read_text())
    if manifest.get("schema") != "qwen30-latency-case-v1":
        raise ValueError(f"unexpected trace schema in {path}")
    if manifest.get("tensor_parallel_size") != tp or int(manifest.get("requests", 0)) <= 0:
        raise ValueError(f"invalid trace contract in {path}")
    outputs = manifest.get("output_tokens")
    if not isinstance(outputs, list) or len(outputs) != 1 or int(outputs[0]) <= 0:
        raise ValueError(f"non-uniform output contract in {path}")
    return manifest


def metric_stats(values: list[float]) -> dict[str, float]:
    return {"samples": len(values), "mean_ms": statistics.fmean(values), "p90_ms": nearest_rank(values, 0.90)}


def validate_trial(path: Path, manifest: dict[str, Any]) -> tuple[dict[str, list[float]], dict[str, Any]]:
    payload = json.loads(path.read_text())
    if payload.get("schema") != "qwen30-exact-trace-anchor-v1":
        raise ValueError(f"unexpected real result schema in {path}")
    contract = payload.get("contract")
    summary = payload.get("summary")
    records = payload.get("requests")
    if not isinstance(contract, dict) or not isinstance(summary, dict) or not isinstance(records, list):
        raise ValueError(f"malformed result payload: {path}")
    expected = int(manifest["requests"])
    expected_contract = {
        "requests": expected,
        "requests_file_sha256": manifest["private_jsonl_sha256"],
        "row_vector_sha256": manifest["row_vector_sha256"],
        "first_arrival_s": manifest["first_arrival_s"],
        "last_arrival_s": manifest["last_arrival_s"],
    }
    for name, wanted in expected_contract.items():
        actual = contract.get(name)
        if isinstance(wanted, float):
            if not isinstance(actual, (int, float)) or not math.isclose(float(actual), wanted, abs_tol=1e-9):
                raise ValueError(f"{path}: contract drift for {name}")
        elif actual != wanted:
            raise ValueError(f"{path}: contract drift for {name}")
    if len(records) != expected or summary.get("completed") != expected or summary.get("failed") != 0:
        raise ValueError(f"{path}: incomplete real replay")

    metrics: dict[str, list[float]] = {"ttft_ms": [], "tpot_ms": [], "e2e_ms": []}
    indices: set[int] = set()
    expected_output = int(manifest["output_tokens"][0])
    for record in records:
        if record.get("success") is not True:
            raise ValueError(f"{path}: failed request record")
        index = record.get("source_index")
        if not isinstance(index, int) or index in indices:
            raise ValueError(f"{path}: duplicate/non-integer source index")
        indices.add(index)
        if record.get("actual_input_tokens") != record.get("input_tokens"):
            raise ValueError(f"{path}: input usage drift")
        if record.get("requested_output_tokens") != expected_output or record.get("actual_output_tokens") != expected_output:
            raise ValueError(f"{path}: output usage drift")
        for metric in ("ttft_ms", "e2e_ms"):
            metrics[metric].append(finite(record.get(metric), metric))
        tpot = record.get("tpot_ms")
        if expected_output == 1:
            if tpot is not None:
                raise ValueError(f"{path}: OSL=1 must report TPOT=null")
        else:
            metrics["tpot_ms"].append(finite(tpot, "tpot_ms"))
    if len(indices) != expected:
        raise ValueError(f"{path}: missing source rows")
    return metrics, {"result_path": str(path), "requests": expected}


def main() -> None:
    args = parse_args()
    root = args.case_root.resolve()
    traces_root = (args.traces_root or root / "traces").resolve()
    manifests = {tp: read_manifest(traces_root, tp) for tp in (1, 2, 4)}
    prefill_only = int(manifests[1]["output_tokens"][0]) == 1
    if any((int(manifest["output_tokens"][0]) == 1) != prefill_only for manifest in manifests.values()):
        raise ValueError("TP-specific output contracts differ")
    applicable = ("ttft_ms", "e2e_ms") if prefill_only else ("ttft_ms", "tpot_ms", "e2e_ms")
    configs: dict[str, Any] = {}
    for name in CONFIGS:
        tp = int(name.split("_", 1)[0].removeprefix("tp"))
        trial_rows = []
        pooled = {metric: [] for metric in applicable}
        for trial in TRIALS:
            path = root / "real" / name / f"trial{trial}" / "results" / "result.json"
            values, provenance = validate_trial(path, manifests[tp])
            trial_rows.append({
                **provenance,
                "trial": trial,
                "metrics": {metric: metric_stats(values[metric]) for metric in applicable},
            })
            for metric in applicable:
                pooled[metric].extend(values[metric])
        configs[name] = {
            "trials": trial_rows,
            "metrics": {
                metric: {
                    "pooled_samples": len(pooled[metric]),
                    "pooled_mean_ms": statistics.fmean(pooled[metric]),
                    "pooled_p90_ms": nearest_rank(pooled[metric], 0.90),
                    "trial_mean_of_means_ms": statistics.fmean(
                        row["metrics"][metric]["mean_ms"] for row in trial_rows
                    ),
                    "trial_stddev_of_means_ms": statistics.stdev(
                        row["metrics"][metric]["mean_ms"] for row in trial_rows
                    ),
                }
                for metric in applicable
            },
        }
    winners: dict[str, Any] = {}
    for metric in applicable:
        for statistic in ("pooled_mean_ms", "pooled_p90_ms"):
            ranked = sorted((row["metrics"][metric][statistic], name) for name, row in configs.items())
            winners[f"{metric}:{statistic}"] = {
                "winner": ranked[0][1],
                "winner_value_ms": ranked[0][0],
                "ranking": [name for _, name in ranked],
            }
    payload = {
        "schema": "qwen30-latency-case-real-audit-v1",
        "case_root": str(root),
        "traces_root": str(traces_root),
        "prefill_only": prefill_only,
        "applicable_metrics": list(applicable),
        "trace_manifests": {f"tp{tp}": manifests[tp] for tp in manifests},
        "configs": configs,
        "winners": winners,
    }
    lines = ["# Qwen30 real latency case audit", "", f"Prefill-only: `{prefill_only}`.", ""]
    lines += ["| Objective | Real winner | Value (ms) |", "|---|---|---:|"]
    for objective, winner in winners.items():
        lines.append(f"| {objective} | {winner['winner']} | {winner['winner_value_ms']:.2f} |")
    args.json_output.parent.mkdir(parents=True, exist_ok=True)
    args.markdown_output.parent.mkdir(parents=True, exist_ok=True)
    args.json_output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    args.markdown_output.write_text("\n".join(lines) + "\n")
    print(json.dumps(winners, sort_keys=True))


if __name__ == "__main__":
    main()
