#!/usr/bin/env python3
"""Project base workload cases to a constant per-GPU request-rate contract."""

from __future__ import annotations

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-manifest", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--tp", type=int, action="append", default=None)
    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 rho_label(rho: float) -> str:
    return format(rho, ".12g").replace(".", "p")


def project_case(
    case: dict[str, Any], output_root: Path, tp: int
) -> dict[str, Any]:
    source_public = Path(case["public_csv"])
    source_private = Path(case["private_jsonl"])
    if sha256(source_public) != case["public_csv_sha256"]:
        raise ValueError(f"source public digest mismatch: {source_public}")
    if sha256(source_private) != case["private_jsonl_sha256"]:
        raise ValueError(f"source private digest mismatch: {source_private}")

    case_root = output_root / f"tp{tp}" / case["family"] / f"rho{rho_label(case['rho'])}"
    public_root = case_root / "public"
    private_root = case_root / "private"
    public_root.mkdir(parents=True, exist_ok=True)
    private_root.mkdir(parents=True, exist_ok=True)
    public_path = public_root / "frontier.csv"
    private_path = private_root / "real_requests.jsonl"

    with source_public.open(newline="") as source:
        reader = csv.DictReader(source)
        rows = list(reader)
        fieldnames = reader.fieldnames
    if not rows or fieldnames is None:
        raise ValueError(f"empty public source: {source_public}")
    with public_path.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=fieldnames, lineterminator="\n")
        writer.writeheader()
        for row in rows:
            row["arrived_at"] = f"{float(row['arrived_at']) / tp:.12f}"
            writer.writerow(row)

    private_rows = [
        json.loads(line) for line in source_private.open() if line.strip()
    ]
    if len(private_rows) != len(rows):
        raise ValueError(f"public/private row mismatch: {source_public}")
    with private_path.open("w") as output:
        for row in private_rows:
            row["arrived_at"] = float(row["arrived_at"]) / tp
            output.write(json.dumps(row, separators=(",", ":")) + "\n")

    projected = dict(case)
    projected.update(
        {
            "load_contract": "constant_per_gpu_request_rate",
            "tp": tp,
            "rho_per_gpu": case["rho"],
            "per_gpu_offered_request_rate": case["global_offered_request_rate"],
            "global_offered_request_rate": case["global_offered_request_rate"] * tp,
            "empirical_interarrival_rate": case["empirical_interarrival_rate"] * tp,
            "decode_offered_tokens_per_second_per_gpu": case[
                "decode_offered_tokens_per_second"
            ],
            "decode_offered_tokens_per_second": case[
                "decode_offered_tokens_per_second"
            ]
            * tp,
            "last_arrival_s": case["last_arrival_s"] / tp,
            "source_case_public_csv": str(source_public),
            "source_case_public_csv_sha256": case["public_csv_sha256"],
            "source_case_private_jsonl": str(source_private),
            "source_case_private_jsonl_sha256": case["private_jsonl_sha256"],
            "public_csv": str(public_path.resolve()),
            "public_csv_sha256": sha256(public_path),
            "private_jsonl": str(private_path.resolve()),
            "private_jsonl_sha256": sha256(private_path),
        }
    )
    (public_root / "manifest.json").write_text(
        json.dumps(projected, indent=2, sort_keys=True) + "\n"
    )
    return projected


def main() -> None:
    args = parse_args()
    tps = args.tp or [1, 2, 4]
    if any(tp not in (1, 2, 4) for tp in tps):
        raise ValueError("TP must be 1, 2, or 4")
    suite = json.loads(args.source_manifest.read_text())
    projected = [
        project_case(case, args.output_root, tp)
        for tp in tps
        for case in suite["cases"]
    ]
    output = {
        "schema": "frontier-workload-regime-per-gpu-v1",
        "load_contract": "constant_per_gpu_request_rate",
        "source_manifest": str(args.source_manifest.resolve()),
        "source_manifest_sha256": sha256(args.source_manifest),
        "tensor_parallel_sizes": tps,
        "cases": projected,
    }
    args.output_root.mkdir(parents=True, exist_ok=True)
    path = args.output_root / "manifest.json"
    path.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
    print(path)


if __name__ == "__main__":
    main()