#!/usr/bin/env python3
"""Audit and summarize the TP-normalized Qwen30 real-serving surface.

The script reads only public trace manifests and prompt-free client result
records. It refuses to score an incomplete or contract-drifting trial.
"""

from __future__ import annotations

import argparse
import json
import math
import statistics
from collections.abc import Iterable
from pathlib import Path
from typing import Any


METRICS = ("ttft_ms", "tpot_ms", "e2e_ms")
CONFIGS = tuple((tp, 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("--output-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 nearest_rank(values: Iterable[float], percentile: float) -> float:
    ordered = sorted(values)
    if not ordered:
        raise ValueError("cannot calculate percentile of empty values")
    return ordered[math.ceil(len(ordered) * percentile) - 1]


def number(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 load_manifest(root: Path, tp: int) -> dict[str, Any]:
    path = root / "traces" / f"tp{tp}" / "public" / "manifest.json"
    manifest = json.loads(path.read_text())
    if manifest.get("schema") != "qwen30-tp-normalized-trace-v1":
        raise ValueError(f"unexpected trace schema in {path}")
    if manifest.get("tensor_parallel_size") != tp:
        raise ValueError(f"TP mismatch in {path}")
    if manifest.get("requests") != 129:
        raise ValueError(f"unexpected request count in {path}")
    return manifest


def validate_trial(
    result_path: Path, manifest: dict[str, Any], tp: int, mns: int, trial: int
) -> tuple[dict[str, Any], dict[str, list[float]]]:
    payload = json.loads(result_path.read_text())
    if payload.get("schema") != "qwen30-exact-trace-anchor-v1":
        raise ValueError(f"unexpected result schema: {result_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: {result_path}")
    expected_requests = int(manifest["requests"])
    checks = {
        "requests": (contract.get("requests"), expected_requests),
        "requests_file_sha256": (
            contract.get("requests_file_sha256"),
            manifest["private_jsonl_sha256"],
        ),
        "row_vector_sha256": (
            contract.get("row_vector_sha256"),
            manifest["normalized_row_vector_sha256"],
        ),
        "first_arrival_s": (
            contract.get("first_arrival_s"), manifest["normalized_first_arrival_s"],
        ),
        "last_arrival_s": (
            contract.get("last_arrival_s"), manifest["normalized_last_arrival_s"],
        ),
        "served_model_alias": (
            contract.get("served_model_alias"), "qwen3-30b-exact-trace"
        ),
    }
    for field, (actual, expected) in checks.items():
        if isinstance(expected, float):
            if not isinstance(actual, (int, float)) or not math.isclose(
                float(actual), expected, abs_tol=1e-9
            ):
                raise ValueError(f"{result_path}: contract {field} drift")
        elif actual != expected:
            raise ValueError(f"{result_path}: contract {field} drift")

    if len(records) != expected_requests:
        raise ValueError(f"{result_path}: expected {expected_requests} records")
    if summary.get("completed") != expected_requests or summary.get("failed") != 0:
        raise ValueError(f"{result_path}: incomplete replay summary")

    source_indices: set[int] = set()
    values: dict[str, list[float]] = {metric: [] for metric in METRICS}
    for record in records:
        if record.get("success") is not True:
            raise ValueError(f"{result_path}: failed request record")
        index = record.get("source_index")
        if not isinstance(index, int) or index in source_indices:
            raise ValueError(f"{result_path}: invalid source index")
        source_indices.add(index)
        if record.get("actual_input_tokens") != record.get("input_tokens"):
            raise ValueError(f"{result_path}: input usage mismatch")
        if record.get("actual_output_tokens") != record.get("requested_output_tokens"):
            raise ValueError(f"{result_path}: output usage mismatch")
        for metric in METRICS:
            if metric == "tpot_ms" and record.get(metric) is None:
                # OSL=1 prefill-only traces intentionally have no TPOT samples.
                continue
            values[metric].append(number(record.get(metric), metric))

    if len(source_indices) != expected_requests:
        raise ValueError(f"{result_path}: missing source index")
    for metric, samples in values.items():
        if not samples:
            raise ValueError(f"{result_path}: no {metric} samples")

    stats = {
        "tp": tp,
        "mns": mns,
        "trial": trial,
        "result_path": str(result_path),
        "requests": expected_requests,
        "metrics": {
            metric: {
                "samples": len(samples),
                "mean_ms": statistics.fmean(samples),
                "p90_ms": nearest_rank(samples, 0.90),
            }
            for metric, samples in values.items()
        },
    }
    return stats, values


def aggregate(config_trials: list[dict[str, Any]], pooled: dict[str, list[float]]) -> dict[str, Any]:
    if len(config_trials) != len(TRIALS):
        raise ValueError("aggregate requires exactly three trials")
    metrics: dict[str, Any] = {}
    for metric in METRICS:
        trial_means = [row["metrics"][metric]["mean_ms"] for row in config_trials]
        trial_p90s = [row["metrics"][metric]["p90_ms"] for row in config_trials]
        values = pooled[metric]
        metrics[metric] = {
            "pooled_samples": len(values),
            "pooled_mean_ms": statistics.fmean(values),
            "pooled_p90_ms": nearest_rank(values, 0.90),
            "trial_mean_of_means_ms": statistics.fmean(trial_means),
            "trial_stddev_of_means_ms": statistics.stdev(trial_means),
            "trial_mean_of_p90s_ms": statistics.fmean(trial_p90s),
        }
    return {"trials": config_trials, "metrics": metrics}


def winners(configs: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
    result: dict[str, dict[str, Any]] = {}
    for metric in METRICS:
        for statistic in ("pooled_mean_ms", "pooled_p90_ms"):
            ranked = sorted(
                (
                    (summary["metrics"][metric][statistic], key)
                    for key, summary in configs.items()
                ),
                key=lambda item: (item[0], item[1]),
            )
            result[f"{metric}:{statistic}"] = {
                "winner": ranked[0][1],
                "winner_value_ms": ranked[0][0],
                "ranking": [key for _, key in ranked],
            }
    return result


def format_ms(value: float) -> str:
    return f"{value:.1f}"


def render_markdown(payload: dict[str, Any]) -> str:
    lines = [
        "# Qwen3-30B-A3B TP-normalized Trace-PD: real vLLM audit",
        "",
        "All 36 fresh-server trials passed the trace contract: 129/129 exact-usage requests per trial. "
        "The smoke run is excluded. Values below pool the three trials (387 requests/config); "
        "p90 uses nearest-rank order statistics.",
        "",
        "| Config | TTFT mean/p90 (ms) | TPOT mean/p90 (ms) | E2E mean/p90 (ms) |",
        "|---|---:|---:|---:|",
    ]
    for key, summary in payload["configs"].items():
        metrics = summary["metrics"]
        lines.append(
            "| {key} | {ttft_mean}/{ttft_p90} | {tpot_mean}/{tpot_p90} | {e2e_mean}/{e2e_p90} |".format(
                key=key,
                ttft_mean=format_ms(metrics["ttft_ms"]["pooled_mean_ms"]),
                ttft_p90=format_ms(metrics["ttft_ms"]["pooled_p90_ms"]),
                tpot_mean=format_ms(metrics["tpot_ms"]["pooled_mean_ms"]),
                tpot_p90=format_ms(metrics["tpot_ms"]["pooled_p90_ms"]),
                e2e_mean=format_ms(metrics["e2e_ms"]["pooled_mean_ms"]),
                e2e_p90=format_ms(metrics["e2e_ms"]["pooled_p90_ms"]),
            )
        )
    lines += ["", "## Per-metric winners", ""]
    for target, winner in payload["winners"].items():
        lines.append(
            f"- `{target}`: `{winner['winner']}` ({format_ms(winner['winner_value_ms'])} ms)"
        )
    lines.append("")
    return "\n".join(lines)


def main() -> None:
    args = parse_args()
    manifests = {tp: load_manifest(args.output_root, tp) for tp in (1, 2, 4)}
    configs: dict[str, dict[str, Any]] = {}
    for tp, mns in CONFIGS:
        key = f"tp{tp}_mns{mns}"
        trial_rows: list[dict[str, Any]] = []
        pooled = {metric: [] for metric in METRICS}
        for trial in TRIALS:
            path = args.output_root / "real" / key / f"trial{trial}" / "results" / "result.json"
            trial_stats, values = validate_trial(path, manifests[tp], tp, mns, trial)
            trial_rows.append(trial_stats)
            for metric in METRICS:
                pooled[metric].extend(values[metric])
        configs[key] = aggregate(trial_rows, pooled)
    payload = {
        "schema": "qwen30-tp-normalized-real-surface-audit-v1",
        "trace_manifests": {
            f"tp{tp}": {
                field: manifests[tp][field]
                for field in (
                    "requests",
                    "private_jsonl_sha256",
                    "normalized_row_vector_sha256",
                    "global_offered_request_rate",
                    "per_gpu_offered_request_rate",
                )
            }
            for tp in manifests
        },
        "configs": configs,
        "winners": winners(configs),
    }
    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(render_markdown(payload))
    print(json.dumps(payload["winners"], sort_keys=True))


if __name__ == "__main__":
    main()