#!/usr/bin/env python3
"""Materialize Qwen30 Fixed/Trace latency cases without changing their contract.

The private request file deliberately contains either an exact source prompt
(trace cases) or deterministic token IDs (fixed cases).  The public Frontier
fixture contains only arrivals, shapes, sessions, and legal complete prefix
block identities.
"""

from __future__ import annotations

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


FIELDS = (
    "arrived_at",
    "num_prefill_tokens",
    "num_decode_tokens",
    "session_id",
    "block_hash_ids",
)


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 row_digest(rows: list[dict[str, Any]]) -> str:
    digest = hashlib.sha256()
    for row in rows:
        digest.update(
            json.dumps(
                [
                    row["source_index"],
                    row["arrived_at"],
                    row["input_length"],
                    row["output_length"],
                    row["session_id"],
                    row["runtime_block_ids"],
                ],
                separators=(",", ":"),
            ).encode()
        )
        digest.update(b"\n")
    return digest.hexdigest()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="kind", required=True)

    trace = sub.add_parser("trace")
    trace.add_argument("--base-public", type=Path, required=True)
    trace.add_argument("--base-private", type=Path, required=True)
    trace.add_argument("--output-tokens", type=int, required=True)
    trace.add_argument("--tp", type=int, required=True)
    trace.add_argument("--output-root", type=Path, required=True)

    fixed = sub.add_parser("fixed")
    fixed.add_argument("--model", type=Path, required=True)
    fixed.add_argument("--input-tokens", type=int, required=True)
    fixed.add_argument("--output-tokens", type=int, required=True)
    fixed.add_argument("--requests", type=int, required=True)
    fixed.add_argument("--per-gpu-rate", type=float, required=True)
    fixed.add_argument("--tp", type=int, required=True)
    fixed.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def materialize_trace(args: argparse.Namespace) -> tuple[list[dict[str, Any]], bool, str, float]:
    with args.base_public.open(newline="") as source:
        public = list(csv.DictReader(source))
    private = [json.loads(line) for line in args.base_private.open() if line.strip()]
    if not public or len(public) != len(private):
        raise ValueError("trace public/private request count mismatch")
    rows: list[dict[str, Any]] = []
    for public_row, private_row in zip(public, private, strict=True):
        if int(public_row["num_prefill_tokens"]) != int(private_row["input_length"]):
            raise ValueError("trace input-length drift")
        if float(public_row["arrived_at"]) != float(private_row["arrived_at"]):
            raise ValueError("trace arrival drift")
        # The real vLLM request artifact retains its final partial prompt block
        # for provenance, while the Frontier CSV deliberately projects only
        # legal complete cache blocks.  Use the latter as the shared
        # simulator-facing identity vector; vLLM itself derives cache hashes
        # from the exact private prompt text.
        runtime_ids = [int(value) for value in public_row["block_hash_ids"].split("|") if value]
        if len(runtime_ids) != int(private_row["input_length"]) // 16:
            raise ValueError("public trace exposes an incomplete runtime prefix block")
        if int(public_row["session_id"]) != int(private_row["session_id"]):
            raise ValueError("trace session-id drift")
        body = dict(private_row["body"])
        body.update(
            {
                "min_tokens": args.output_tokens,
                "max_tokens": args.output_tokens,
                "ignore_eos": True,
            }
        )
        rows.append(
            {
                "source_index": int(private_row["source_index"]),
                "arrived_at": float(private_row["arrived_at"]),
                "input_length": int(private_row["input_length"]),
                "output_length": args.output_tokens,
                "session_id": int(private_row["session_id"]),
                "runtime_block_ids": runtime_ids,
                "body": body,
            }
        )
    return (
        rows,
        True,
        "trace-derived: exact input/arrival/session/prefix; output override only",
        len(rows) / 600.0 * args.tp,
    )


def materialize_fixed(args: argparse.Namespace) -> tuple[list[dict[str, Any]], bool, str, float]:
    if min(args.input_tokens, args.output_tokens, args.requests, args.tp) <= 0:
        raise ValueError("fixed dimensions must be positive")
    if args.input_tokens + args.output_tokens > 40960:
        raise ValueError("fixed shape exceeds max model length")
    if not math.isfinite(args.per_gpu_rate) or args.per_gpu_rate <= 0:
        raise ValueError("per-GPU rate must be positive")
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    candidates = [
        token for token in range(tokenizer.vocab_size) if token not in set(tokenizer.all_special_ids)
    ]
    if len(candidates) < args.requests + 1:
        raise ValueError("tokenizer does not provide enough non-special tokens")
    rate = args.per_gpu_rate * args.tp
    base = candidates[0]
    rows = []
    for index in range(args.requests):
        rows.append(
            {
                "source_index": index,
                "arrived_at": index / rate,
                "input_length": args.input_tokens,
                "output_length": args.output_tokens,
                "session_id": index,
                "runtime_block_ids": [],
                "body": {
                    "prompt": [candidates[index + 1], *([base] * (args.input_tokens - 1))],
                    "min_tokens": args.output_tokens,
                    "max_tokens": args.output_tokens,
                    "ignore_eos": True,
                },
            }
        )
    return (
        rows,
        False,
        "fixed-shape: deterministic token IDs, uniform TP-normalized QPS, no prefix reuse",
        rate,
    )


def write_case(
    rows: list[dict[str, Any]],
    *,
    prefix_caching: bool,
    description: str,
    global_rate: float,
    tp: int,
    root: Path,
) -> None:
    if not rows:
        raise ValueError("empty case")
    root = root.resolve()
    public_root = root / "public"
    private_root = 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 public_path.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
        writer.writeheader()
        for row in rows:
            writer.writerow(
                {
                    "arrived_at": f"{row['arrived_at']:.12f}",
                    "num_prefill_tokens": row["input_length"],
                    "num_decode_tokens": row["output_length"],
                    "session_id": row["session_id"],
                    "block_hash_ids": "|".join(str(value) for value in row["runtime_block_ids"]),
                }
            )
    with private_path.open("w") as output:
        for row in rows:
            output.write(json.dumps(row, separators=(",", ":")) + "\n")
    arrivals = [float(row["arrived_at"]) for row in rows]
    manifest = {
        "schema": "qwen30-latency-case-v1",
        "description": description,
        "tensor_parallel_size": tp,
        "requests": len(rows),
        "output_tokens": sorted({int(row["output_length"]) for row in rows}),
        "prefix_caching": prefix_caching,
        "public_csv": str(public_path),
        "public_csv_sha256": sha256(public_path),
        "private_jsonl": str(private_path),
        "private_jsonl_sha256": sha256(private_path),
        "row_vector_sha256": row_digest(rows),
        "first_arrival_s": arrivals[0],
        "last_arrival_s": arrivals[-1],
        "global_offered_request_rate": global_rate,
        "per_gpu_offered_request_rate": global_rate / tp,
    }
    (public_root / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest, sort_keys=True))


def main() -> None:
    args = parse_args()
    if args.kind == "trace":
        rows, prefix, description, global_rate = materialize_trace(args)
    else:
        rows, prefix, description, global_rate = materialize_fixed(args)
    write_case(
        rows,
        prefix_caching=prefix,
        description=description,
        global_rate=global_rate,
        tp=args.tp,
        root=args.output_root,
    )


if __name__ == "__main__":
    main()
