#!/usr/bin/env python3
"""Derive a per-GPU-normalized real/Frontier Qwen30 trace pair.

Arrival times change as ``t' = t / TP``.  The private JSONL preserves the
complete prompt identity vector for real replay; the Frontier CSV exports only
the completed 16-token prompt blocks that vLLM can legally reuse from prefix
cache. Prompt bodies remain exclusively in the private output.
"""

from __future__ import annotations

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


CSV_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 update_digest(digest: Any, values: list[Any]) -> None:
    digest.update(json.dumps(values, separators=(",", ":")).encode())
    digest.update(b"\n")


def full_prompt_block_ids(
    runtime_block_ids: list[int], input_tokens: int, block_size: int
) -> list[int]:
    """Project exact prompt identities to vLLM-cacheable full blocks only."""
    if input_tokens <= 0 or block_size <= 0:
        raise ValueError("input_tokens and block_size must be positive")
    expected_runtime_blocks = math.ceil(input_tokens / block_size)
    if len(runtime_block_ids) != expected_runtime_blocks:
        raise ValueError(
            "runtime block identity count does not match the exact prompt: "
            f"got {len(runtime_block_ids)}, expected {expected_runtime_blocks}"
        )
    return runtime_block_ids[: input_tokens // block_size]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-public-csv", type=Path, required=True)
    parser.add_argument("--base-private-jsonl", type=Path, required=True)
    parser.add_argument("--tp", type=int, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--base-window-seconds", type=float, default=600.0)
    parser.add_argument("--runtime-block-size", type=int, default=16)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if (
        args.tp <= 0
        or args.runtime_block_size <= 0
        or not math.isfinite(args.base_window_seconds)
        or args.base_window_seconds <= 0
    ):
        raise ValueError("TP, runtime block size, and base window must be positive")

    public_rows = list(csv.DictReader(args.base_public_csv.open(newline="")))
    private_rows = [
        json.loads(line) for line in args.base_private_jsonl.open() if line.strip()
    ]
    if not public_rows or len(public_rows) != len(private_rows):
        raise ValueError("public/private request counts must match and be non-zero")
    if set(public_rows[0]) != set(CSV_FIELDS):
        raise ValueError("unexpected public trace schema")

    normalized_digest = hashlib.sha256()
    semantic_digest = hashlib.sha256()
    normalized_private: list[dict[str, Any]] = []
    normalized_public: list[dict[str, str]] = []
    original_arrivals: list[float] = []
    frontier_prefix_digest = hashlib.sha256()
    partial_prompt_count = 0
    raw_runtime_block_count = 0
    frontier_full_block_count = 0

    for index, (public, private) in enumerate(
        zip(public_rows, private_rows, strict=True)
    ):
        original_arrival = float(public["arrived_at"])
        if not math.isfinite(original_arrival) or original_arrival < 0:
            raise ValueError(f"invalid arrival at row {index}")
        if not math.isclose(original_arrival, float(private["arrived_at"]), abs_tol=1e-9):
            raise ValueError(f"public/private arrival mismatch at row {index}")
        runtime_block_ids = [int(value) for value in private["runtime_block_ids"]]
        expected = (
            int(public["num_prefill_tokens"]),
            int(public["num_decode_tokens"]),
            int(public["session_id"]),
            public["block_hash_ids"],
        )
        actual = (
            int(private["input_length"]),
            int(private["output_length"]),
            int(private["session_id"]),
            "|".join(str(value) for value in runtime_block_ids),
        )
        if expected != actual:
            raise ValueError(f"public/private payload mismatch at row {index}")

        input_tokens = int(private["input_length"])
        frontier_block_ids = full_prompt_block_ids(
            runtime_block_ids, input_tokens, args.runtime_block_size
        )
        if not frontier_block_ids:
            raise ValueError(
                "Frontier prefix-cache replay requires at least one complete "
                f"runtime block; request {index} has ISL={input_tokens}"
            )
        partial_prompt_count += int(input_tokens % args.runtime_block_size != 0)
        raw_runtime_block_count += len(runtime_block_ids)
        frontier_full_block_count += len(frontier_block_ids)

        normalized_arrival = original_arrival / args.tp
        updated = dict(private)
        updated["arrived_at"] = normalized_arrival
        normalized_private.append(updated)
        normalized_public.append(
            {
                "arrived_at": f"{normalized_arrival:.12f}",
                "num_prefill_tokens": public["num_prefill_tokens"],
                "num_decode_tokens": public["num_decode_tokens"],
                "session_id": public["session_id"],
                "block_hash_ids": "|".join(str(value) for value in frontier_block_ids),
            }
        )
        update_digest(
            semantic_digest,
            [
                int(private["source_index"]),
                int(private["input_length"]),
                int(private["output_length"]),
                int(private["session_id"]),
                runtime_block_ids,
            ],
        )
        update_digest(
            normalized_digest,
            [
                int(private["source_index"]),
                normalized_arrival,
                int(private["input_length"]),
                int(private["output_length"]),
                int(private["session_id"]),
                runtime_block_ids,
            ],
        )
        update_digest(
            frontier_prefix_digest,
            [
                int(private["source_index"]),
                normalized_arrival,
                input_tokens,
                frontier_block_ids,
            ],
        )
        original_arrivals.append(original_arrival)

    if any(
        right < left for left, right in zip(original_arrivals, original_arrivals[1:])
    ):
        raise ValueError("base arrival order is not monotonic")

    public_dir = args.output_root / "public"
    private_dir = args.output_root / "private"
    public_dir.mkdir(parents=True, exist_ok=True)
    private_dir.mkdir(parents=True, exist_ok=True)
    public_path = public_dir / "frontier.csv"
    private_path = private_dir / "real_requests.jsonl"
    with public_path.open("w", newline="") as output:
        writer = csv.DictWriter(output, fieldnames=CSV_FIELDS, lineterminator="\n")
        writer.writeheader()
        writer.writerows(normalized_public)
    with private_path.open("w") as output:
        for row in normalized_private:
            output.write(json.dumps(row, separators=(",", ":")) + "\n")

    manifest = {
        "schema": "qwen30-tp-normalized-trace-v2",
        "privacy": "prompt text exists only under private/ and must not be harvested",
        "transform": "arrival_seconds_prime = arrival_seconds / tensor_parallel_size",
        "tensor_parallel_size": args.tp,
        "requests": len(normalized_private),
        "base_window_seconds": args.base_window_seconds,
        "global_offered_request_rate": len(normalized_private)
        / args.base_window_seconds
        * args.tp,
        "per_gpu_offered_request_rate": len(normalized_private)
        / args.base_window_seconds,
        "frontier_prefix_block_projection": {
            "runtime_block_size": args.runtime_block_size,
            "rule": "full_prompt_blocks_only=floor(input_tokens/runtime_block_size)",
            "partial_prompt_count": partial_prompt_count,
            "raw_runtime_block_count": raw_runtime_block_count,
            "frontier_full_block_count": frontier_full_block_count,
        },
        "original_first_arrival_s": original_arrivals[0],
        "original_last_arrival_s": original_arrivals[-1],
        "normalized_first_arrival_s": original_arrivals[0] / args.tp,
        "normalized_last_arrival_s": original_arrivals[-1] / args.tp,
        "base_public_csv": str(args.base_public_csv.resolve()),
        "base_public_csv_sha256": sha256(args.base_public_csv),
        "base_private_jsonl": str(args.base_private_jsonl.resolve()),
        "base_private_jsonl_sha256": sha256(args.base_private_jsonl),
        "public_csv": str(public_path.resolve()),
        "public_csv_sha256": sha256(public_path),
        "private_jsonl": str(private_path.resolve()),
        "private_jsonl_sha256": sha256(private_path),
        "semantic_vector_sha256": semantic_digest.hexdigest(),
        "normalized_row_vector_sha256": normalized_digest.hexdigest(),
        "frontier_prefix_vector_sha256": frontier_prefix_digest.hexdigest(),
    }
    manifest_path = public_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest, sort_keys=True))


if __name__ == "__main__":
    main()
