#!/usr/bin/env python3
"""Validate Qwen235 serving collectives and materialize Frontier's CC CSV."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from pathlib import Path
from typing import Any


VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
TOKEN_POINTS = (
    1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 63, 64, 65, 72, 80, 88, 96,
    104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200,
    208, 216, 224, 232, 240, 248, 255, 256, 257, 512, 1024, 2048, 4096,
    8192,
)
TP_SIZES = (4, 8)
HIDDEN_DIM = 4096
ALLOWED_BACKENDS = {
    "flashinfer_trtllm_fused_projection",
    "pynccl_symmetric_with_copy",
    "torch_symmetric_memory",
    "pynccl",
}
FIELDS = (
    "time_stats.all_reduce.min",
    "time_stats.all_reduce.max",
    "time_stats.all_reduce.mean",
    "time_stats.all_reduce.median",
    "time_stats.all_reduce.std",
    "rank",
    "num_workers",
    "size",
    "collective",
    "devices_per_node",
    "max_devices_per_node",
)


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


def materialize(
    input_paths: list[Path], output_path: Path, manifest_path: Path
) -> dict[str, Any]:
    if len(input_paths) != 2:
        raise ValueError("exactly two raw inputs are required for TP4 and TP8")
    rows = []
    seen = set()
    backends_by_tp: dict[int, set[str]] = {}
    environments = []
    for path in input_paths:
        payload = json.loads(path.read_text())
        if payload.get("schema_version") != "vllm020_allreduce_raw.v2":
            raise ValueError(f"unexpected raw collective schema in {path}")
        environment = payload["environment"]
        if environment.get("vllm_version") != "0.20.0":
            raise ValueError(f"vLLM version mismatch in {path}")
        if environment.get("vllm_source_commit") != VLLM_COMMIT:
            raise ValueError(f"vLLM commit mismatch in {path}")
        if environment.get("gpu") != "NVIDIA H20":
            raise ValueError(f"GPU mismatch in {path}: {environment.get('gpu')!r}")
        if Path(environment.get("model", "")).name != "Qwen3-235B-A22B-FP8":
            raise ValueError(f"model mismatch in {path}")
        if environment.get("collective_contract") != "qwen235-serving-projected":
            raise ValueError(f"collective contract mismatch in {path}")
        if environment.get("disable_custom_all_reduce") is not True:
            raise ValueError(f"custom all-reduce was not disabled in {path}")
        backend_env = environment.get("backend_env", {})
        if backend_env.get("VLLM_ALLREDUCE_USE_FLASHINFER") != "1":
            raise ValueError(f"FlashInfer projection was not enabled in {path}")
        if backend_env.get("VLLM_ALLREDUCE_USE_SYMM_MEM") != "1":
            raise ValueError(f"symmetric-memory contract mismatch in {path}")
        environments.append(environment)

        for source in payload["rows"]:
            tp = int(source["tensor_parallel_size"])
            tokens = int(source["num_tokens"])
            key = (tp, tokens)
            if key in seen:
                raise ValueError(f"duplicate collective row: {key}")
            seen.add(key)
            if tp not in TP_SIZES:
                raise ValueError(f"unexpected TP: {tp}")
            if int(source["hidden_dim"]) != HIDDEN_DIM:
                raise ValueError(f"hidden-dim mismatch for {key}")
            expected_bytes = tokens * HIDDEN_DIM * 2
            if int(source["payload_bytes"]) != expected_bytes:
                raise ValueError(f"payload mismatch for {key}")
            backend = source["selected_backend"]
            if backend not in ALLOWED_BACKENDS:
                raise ValueError(f"unexpected backend for {key}: {backend!r}")
            fusion_limit = {4: 2 * 1024 * 1024, 8: 512 * 1024}[tp]
            if int(source["real_fusion_limit_bytes"]) != fusion_limit:
                raise ValueError(f"fusion limit mismatch for {key}")
            expected_backend = (
                "flashinfer_trtllm_fused_projection"
                if expected_bytes <= fusion_limit
                else None
            )
            if expected_backend is not None and backend != expected_backend:
                raise ValueError(f"missing fused projection for {key}")
            if expected_backend is None and backend == "flashinfer_trtllm_fused_projection":
                raise ValueError(f"fused projection exceeds runtime limit for {key}")
            backends_by_tp.setdefault(tp, set()).add(backend)
            median = float(source["critical_path_median_ms"])
            if not math.isfinite(median) or median <= 0:
                raise ValueError(f"invalid critical-path median for {key}")
            if int(source.get("trials", 0)) != 3 or int(
                source.get("repeats_per_trial", 0)
            ) != 20:
                raise ValueError(f"repeat contract mismatch for {key}")
            samples = source.get("per_trial_rank_samples_ms", [])
            if (
                len(samples) != 3
                or any(len(trial) != tp for trial in samples)
                or any(
                    len(rank_samples) != 20
                    or any(
                        not math.isfinite(float(value)) or float(value) <= 0
                        for value in rank_samples
                    )
                    for trial in samples
                    for rank_samples in trial
                )
            ):
                raise ValueError(f"raw sample coverage mismatch for {key}")
            critical_samples = [
                max(float(trial[rank][repeat]) for rank in range(tp))
                for trial in samples
                for repeat in range(20)
            ]
            ordered = sorted(critical_samples)
            measured_median = (ordered[29] + ordered[30]) / 2
            if not math.isclose(median, measured_median, rel_tol=1e-6, abs_tol=1e-9):
                raise ValueError(f"critical-path median/sample mismatch for {key}")
            rows.append(
                {
                    "time_stats.all_reduce.min": median,
                    "time_stats.all_reduce.max": median,
                    "time_stats.all_reduce.mean": median,
                    "time_stats.all_reduce.median": median,
                    "time_stats.all_reduce.std": 0.0,
                    "rank": 0,
                    "num_workers": tp,
                    "size": expected_bytes,
                    "collective": "all_reduce",
                    "devices_per_node": tp,
                    "max_devices_per_node": 8,
                }
            )

    expected = {(tp, tokens) for tp in TP_SIZES for tokens in TOKEN_POINTS}
    if seen != expected:
        raise ValueError(
            f"collective coverage mismatch: missing={expected - seen}, extra={seen - expected}"
        )
    if output_path.exists() or manifest_path.exists():
        raise FileExistsError("refusing to overwrite immutable collective profile")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
        writer.writeheader()
        writer.writerows(sorted(rows, key=lambda row: (row["num_workers"], row["size"])))

    manifest = {
        "schema": "qwen235-v020-serving-allreduce-profile-v1",
        "inputs": {str(path.resolve()): sha256(path) for path in input_paths},
        "output": str(output_path.resolve()),
        "output_sha256": sha256(output_path),
        "vllm_source_commit": VLLM_COMMIT,
        "hardware": "NVIDIA H20",
        "model": "Qwen3-235B-A22B-FP8",
        "dtype": "bfloat16",
        "hidden_dim": HIDDEN_DIM,
        "collective_contract": "qwen235-serving-projected",
        "tp_coverage": list(TP_SIZES),
        "token_points": list(TOKEN_POINTS),
        "rows": len(rows),
        "observed_backends_by_tp": {
            str(tp): sorted(backends) for tp, backends in sorted(backends_by_tp.items())
        },
        "frontier_target": "time_stats.all_reduce.median",
        "unused_stat_policy": "repeat critical_path_median; std=0",
        "raw_environments": environments,
    }
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, action="append", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    args = parser.parse_args()
    manifest = materialize(args.input, args.output, args.manifest)
    print(json.dumps(manifest, sort_keys=True))


if __name__ == "__main__":
    main()
