#!/usr/bin/env python3
"""Convert frozen vLLM collective measurements to Frontier Vidur CC CSV."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
from pathlib import Path


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 convert(input_path: Path, output_path: Path) -> dict[str, object]:
    payload = json.loads(input_path.read_text())
    if payload.get("schema_version") != "qwen30_vllm020_allreduce_frozen.v1":
        raise ValueError(f"unexpected input schema: {payload.get('schema_version')!r}")

    rows = []
    seen = set()
    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 (2, 4):
            raise ValueError(f"unsupported TP: {tp}")
        expected_bytes = tokens * int(source["hidden_dim"]) * 2
        if int(source["payload_bytes"]) != expected_bytes:
            raise ValueError(f"payload mismatch for {key}")

        # Frontier Vidur consumes only the median target. The raw profiler kept
        # per-rank distributions but not aligned per-repeat critical-path
        # samples, so do not invent critical-path min/mean/max/std. Repeating
        # the observed critical-path median in the unused fields keeps the CSV
        # schema explicit without changing the trained target.
        median = float(source["critical_path_median_ms"])
        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 (2, 4) for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)}
    if seen != expected:
        raise ValueError(f"collective coverage mismatch: missing={expected - seen}, extra={seen - expected}")

    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"])))

    return {
        "schema": "frontier-vidur-allreduce-materialization-v1",
        "source": str(input_path.resolve()),
        "source_sha256": sha256(input_path),
        "output": str(output_path.resolve()),
        "output_sha256": sha256(output_path),
        "rows": len(rows),
        "tp_coverage": [2, 4],
        "target": "time_stats.all_reduce.median",
        "unused_stat_policy": "repeat critical_path_median; std=0",
        "payload_contract": "size=num_tokens*hidden_dim*2_bytes",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    args = parser.parse_args()
    manifest = convert(args.input, args.output)
    args.manifest.parent.mkdir(parents=True, exist_ok=True)
    args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest, sort_keys=True))


if __name__ == "__main__":
    main()