#!/usr/bin/env python3
"""Append measured true-mixed attention rows to an immutable Qwen235 profile."""

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()
    parser.add_argument("--base-frozen", type=Path, required=True)
    parser.add_argument("--cuda-tp4", type=Path, required=True)
    parser.add_argument("--cuda-tp8", type=Path, required=True)
    parser.add_argument("--kernel-tp4", type=Path, required=True)
    parser.add_argument("--kernel-tp8", 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 read_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]:
    with path.open(newline="") as source:
        reader = csv.DictReader(source)
        return list(reader.fieldnames or []), list(reader)


def merge_csv(inputs: list[Path], output: Path, measurement_type: str) -> int:
    fields: list[str] = []
    rows: list[dict[str, str]] = []
    true_mixed_count = 0
    for index, path in enumerate(inputs):
        input_fields, input_rows = read_rows(path)
        for field in input_fields:
            if field not in fields:
                fields.append(field)
        if index:
            for row in input_rows:
                if row.get("is_true_mixed_batch", "").lower() != "true":
                    raise ValueError(f"non-true-mixed row in {path}")
                if row.get("measurement_type") != measurement_type:
                    raise ValueError(f"measurement type mismatch in {path}")
            true_mixed_count += len(input_rows)
        rows.extend(input_rows)
    if not fields or not rows or not true_mixed_count:
        raise ValueError("cannot assemble an empty true-mixed augmentation")
    with output.open("w", newline="") as target:
        writer = csv.DictWriter(target, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)
    return true_mixed_count


def main() -> None:
    args = parse_args()
    output = args.output_root.resolve()
    if output.exists():
        raise FileExistsError(output)
    output.mkdir(parents=True)
    for name in ("linear_op.csv", "linear_op_kernel_only.csv", "moe.csv", "moe_kernel_only.csv"):
        shutil.copy2(args.base_frozen / name, output / name)
    cuda_count = merge_csv(
        [args.base_frozen / "attention.csv", args.cuda_tp4, args.cuda_tp8],
        output / "attention.csv",
        "CUDA_EVENT",
    )
    kernel_count = merge_csv(
        [args.base_frozen / "attention_kernel_only.csv", args.kernel_tp4, args.kernel_tp8],
        output / "attention_kernel_only.csv",
        "KERNEL_ONLY",
    )
    base_manifest = json.loads((args.base_frozen / "manifest.json").read_text())
    manifest = {
        **base_manifest,
        "schema": "qwen235-v020-frontier-profile-v2-true-mixed",
        "true_mixed_inputs": {
            name: str(getattr(args, name).resolve())
            for name in ("cuda_tp4", "cuda_tp8", "kernel_tp4", "kernel_tp8")
        },
        "true_mixed_rows": {"CUDA_EVENT": cuda_count, "KERNEL_ONLY": kernel_count},
        "outputs": {path.name: digest(path) for path in sorted(output.glob("*.csv"))},
    }
    (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()