#!/usr/bin/env python3
"""Assemble immutable flat profiles for the Qwen235 Frontier runner."""

from __future__ import annotations

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


MODEL = "Qwen3-235B-A22B"


def parse_args():
    parser = argparse.ArgumentParser()
    for name in ("cuda_common", "cuda_moe_tp4", "cuda_moe_ep8", "kernel_common", "kernel_moe_tp4", "kernel_moe_ep8"):
        parser.add_argument(f"--{name.replace('_', '-')}", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def digest(path: Path) -> str:
    value = hashlib.sha256()
    value.update(path.read_bytes())
    return value.hexdigest()


def model_file(root: Path, name: str) -> Path:
    path = root / "compute" / "h20" / MODEL / name
    if not path.is_file():
        raise FileNotFoundError(path)
    return path


def merge_csv(inputs: list[Path], output: Path) -> None:
    fields = None
    rows = []
    for path in inputs:
        with path.open(newline="") as source:
            reader = csv.DictReader(source)
            if fields is None:
                fields = reader.fieldnames
            elif reader.fieldnames != fields:
                raise ValueError(f"CSV schema mismatch: {path}")
            rows.extend(reader)
    if not fields or not rows:
        raise ValueError("cannot merge empty profile CSV")
    with output.open("w", newline="") as target:
        writer = csv.DictWriter(target, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    args = parse_args()
    output = args.output_root.resolve()
    if output.exists():
        raise FileExistsError(output)
    output.mkdir(parents=True)
    sources = {
        "linear_op.csv": model_file(args.cuda_common, "linear_op.csv"),
        "attention.csv": model_file(args.cuda_common, "attention.csv"),
        "linear_op_kernel_only.csv": model_file(args.kernel_common, "linear_op_kernel_only.csv"),
        "attention_kernel_only.csv": model_file(args.kernel_common, "attention_kernel_only.csv"),
    }
    for name, source in sources.items():
        shutil.copy2(source, output / name)
    merge_csv(
        [model_file(args.cuda_moe_tp4, "moe.csv"), model_file(args.cuda_moe_ep8, "moe.csv")],
        output / "moe.csv",
    )
    merge_csv(
        [model_file(args.kernel_moe_tp4, "moe_kernel_only.csv"), model_file(args.kernel_moe_ep8, "moe_kernel_only.csv")],
        output / "moe_kernel_only.csv",
    )
    outputs = {path.name: digest(path) for path in sorted(output.glob("*.csv"))}
    manifest = {
        "schema": "qwen235-v020-frontier-profile-v1",
        "model": MODEL,
        "measurement_families": ["CUDA_EVENT", "KERNEL_ONLY"],
        "moe_runtime_paths": {"tp4_ep1": "TRITON", "tp8_ep8": "FLASHINFER_CUTLASS"},
        "inputs": {name: str(getattr(args, name).resolve()) for name in ("cuda_common", "cuda_moe_tp4", "cuda_moe_ep8", "kernel_common", "kernel_moe_tp4", "kernel_moe_ep8")},
        "outputs": outputs,
    }
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest, sort_keys=True))


if __name__ == "__main__":
    main()