#!/usr/bin/env python3
"""Create an immutable Frontier profile root with measured decode rows."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import shutil
import tempfile
from pathlib import Path
from typing import Any


MODEL = "Qwen3-235B-A22B-FP8"
ATTENTION_RELATIVE_PATH = Path("compute/h20") / MODEL / "attention.csv"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-profile-root", type=Path, required=True)
    parser.add_argument("--decode-attention-csv", type=Path, required=True)
    parser.add_argument("--true-mixed-attention-csv", type=Path)
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


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 read_csv(path: Path) -> tuple[list[str], list[dict[str, str]]]:
    with path.open(newline="") as source:
        reader = csv.DictReader(source)
        if reader.fieldnames is None:
            raise ValueError(f"missing CSV header: {path}")
        return reader.fieldnames, list(reader)


def is_true(value: str) -> bool:
    return value.strip().lower() == "true"


def profile_hashes(root: Path) -> dict[str, str]:
    return {
        str(path.relative_to(root)): sha256(path)
        for path in sorted(root.rglob("*"))
        if path.is_file() and path.name != "profile_closure_manifest.json"
    }


def write_json(path: Path, payload: Any) -> None:
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def main() -> None:
    args = parse_args()
    base_root = args.base_profile_root.resolve()
    decode_csv = args.decode_attention_csv.resolve()
    output_root = args.output_root.resolve()
    base_attention = base_root / ATTENTION_RELATIVE_PATH

    if output_root.exists():
        raise FileExistsError(f"refusing to overwrite profile root: {output_root}")
    for path in (base_attention, decode_csv):
        if not path.is_file():
            raise FileNotFoundError(path)

    base_fields, base_rows = read_csv(base_attention)
    decode_fields, decode_source_rows = read_csv(decode_csv)
    if base_fields != decode_fields:
        raise ValueError("base and decode attention CSV schemas differ")
    if not base_rows or any(not is_true(row["is_prefill"]) for row in base_rows):
        raise ValueError("base attention profile must contain only prefill rows")
    decode_rows = [row for row in decode_source_rows if not is_true(row["is_prefill"])]
    if not decode_rows:
        raise ValueError("decode attention profile contains no decode rows")
    if any(not row["time_stats.attn_decode.median"] for row in decode_rows):
        raise ValueError("decode attention profile has an empty median")

    dimensions = {
        (
            int(row["num_tensor_parallel_workers"]),
            int(row["batch_size"]),
            int(row["kv_cache_size"]),
            row["attention_backend"],
        )
        for row in decode_rows
    }
    if len(dimensions) != len(decode_rows):
        raise ValueError("decode attention profile has duplicate coverage coordinates")
    if any(dimension[-1] != "FLASHINFER" for dimension in dimensions):
        raise ValueError("decode attention profile is not entirely FlashInfer")

    true_mixed_csv = (
        args.true_mixed_attention_csv.resolve()
        if args.true_mixed_attention_csv is not None
        else None
    )
    true_mixed_rows: list[dict[str, str]] = []
    output_fields = list(base_fields)
    if true_mixed_csv is not None:
        if not true_mixed_csv.is_file():
            raise FileNotFoundError(true_mixed_csv)
        true_mixed_fields, true_mixed_rows = read_csv(true_mixed_csv)
        required = {
            "is_true_mixed_batch",
            "decode_batch_size",
            "decode_avg_kv_cache_size",
            "num_prefill_seqs",
            "time_stats.attn_decode.median",
            "time_stats.attn_prefill.median",
        }
        missing = required - set(true_mixed_fields)
        if missing:
            raise ValueError(f"true-mixed attention CSV lacks columns: {sorted(missing)}")
        if not true_mixed_rows or any(
            not is_true(row["is_true_mixed_batch"]) for row in true_mixed_rows
        ):
            raise ValueError("true-mixed attention CSV has non-mixed rows")
        if any(
            not row["time_stats.attn_decode.median"]
            or not row["time_stats.attn_prefill.median"]
            for row in true_mixed_rows
        ):
            raise ValueError("true-mixed attention profile has an empty median")
        if {
            int(row["num_tensor_parallel_workers"]) for row in true_mixed_rows
        } != {4, 8}:
            raise ValueError("true-mixed attention profile must cover TP4 and TP8")
        output_fields.extend(
            field for field in true_mixed_fields if field not in output_fields
        )

    for row in [*base_rows, *decode_rows]:
        if "is_true_mixed_batch" in output_fields:
            row["is_true_mixed_batch"] = "False"

    output_root.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(
        prefix=f".{output_root.name}.", dir=output_root.parent
    ) as temporary:
        temporary_root = Path(temporary) / output_root.name
        shutil.copytree(base_root, temporary_root)
        merged_attention = temporary_root / ATTENTION_RELATIVE_PATH
        with merged_attention.open("w", newline="") as output:
            writer = csv.DictWriter(output, fieldnames=output_fields, lineterminator="\n")
            writer.writeheader()
            writer.writerows([*base_rows, *decode_rows, *true_mixed_rows])

        payload = {
            "schema": "frontier-profile-closure-v1",
            "model": MODEL,
            "base_profile_root": str(base_root),
            "base_attention_sha256": sha256(base_attention),
            "decode_attention_csv": str(decode_csv),
            "decode_attention_sha256": sha256(decode_csv),
            "base_prefill_rows": len(base_rows),
            "decode_source_rows": len(decode_source_rows),
            "added_decode_rows": len(decode_rows),
            "true_mixed_attention_csv": (
                str(true_mixed_csv) if true_mixed_csv is not None else None
            ),
            "true_mixed_attention_sha256": (
                sha256(true_mixed_csv) if true_mixed_csv is not None else None
            ),
            "added_true_mixed_rows": len(true_mixed_rows),
            "merged_attention_rows": (
                len(base_rows) + len(decode_rows) + len(true_mixed_rows)
            ),
            "decode_dimensions": [list(values) for values in sorted(dimensions)],
            "output_files_sha256": profile_hashes(temporary_root),
        }
        write_json(temporary_root / "profile_closure_manifest.json", payload)
        temporary_root.rename(output_root)

    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
