#!/usr/bin/env python3
"""Materialize deterministic, prefix-disjoint fixed-shape Frontier traces."""

from __future__ import annotations

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


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--input-tokens", type=int, required=True)
    parser.add_argument("--output-tokens", type=int, required=True)
    parser.add_argument("--requests", type=int, default=64)
    parser.add_argument("--duration-seconds", type=float)
    parser.add_argument("--rate", type=float, action="append", required=True)
    parser.add_argument("--block-size", type=int, default=16)
    return parser.parse_args()


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def rate_key(rate: float) -> str:
    return f"r{rate:g}".replace(".", "p")


def materialize(
    root: Path,
    *,
    input_tokens: int,
    output_tokens: int,
    requests: int,
    rates: list[float],
    block_size: int,
    duration_seconds: float | None = None,
) -> dict[str, object]:
    values = [input_tokens, output_tokens, requests, block_size]
    if any(value <= 0 for value in values):
        raise ValueError("token counts, requests, and block size must be positive")
    if input_tokens + output_tokens > 40960:
        raise ValueError("shape exceeds the frozen max model length")
    if not rates or any(rate <= 0 or not math.isfinite(rate) for rate in rates):
        raise ValueError("rates must be positive and finite")
    if len(set(rates)) != len(rates):
        raise ValueError("rates must be unique")
    if duration_seconds is not None and duration_seconds <= 0:
        raise ValueError("duration must be positive")
    root.mkdir(parents=True, exist_ok=True)
    anchors = []
    blocks = math.ceil(input_tokens / block_size)
    for rate in rates:
        request_count = max(
            requests,
            math.ceil(rate * duration_seconds) if duration_seconds is not None else 0,
        )
        path = root / f"{rate_key(rate)}.csv"
        with path.open("w", newline="") as output:
            writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
            writer.writeheader()
            for request_id in range(request_count):
                first_block = request_id * blocks + 1
                writer.writerow(
                    {
                        "arrived_at": f"{request_id / rate:.12f}",
                        "num_prefill_tokens": input_tokens,
                        "num_decode_tokens": output_tokens,
                        "session_id": request_id,
                        "block_hash_ids": "|".join(
                            str(first_block + offset) for offset in range(blocks)
                        ),
                    }
                )
        anchors.append(
            {
                "label": rate_key(rate),
                "rate": rate,
                "requests": request_count,
                "path": path.name,
                "sha256": sha256(path),
            }
        )
    manifest = {
        "schema": "qwen30-fixed-frontier-traces-v1",
        "contract": {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "minimum_requests_per_anchor": requests,
            "minimum_steady_duration_seconds": duration_seconds,
            "block_size": block_size,
            "arrival": "open_loop_uniform",
            "prefix_caching": False,
            "prefix_relation": "all request block identities are disjoint",
        },
        "anchors": anchors,
    }
    (root / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n"
    )
    return manifest


def main() -> None:
    args = parse_args()
    manifest = materialize(
        args.output_root.resolve(),
        input_tokens=args.input_tokens,
        output_tokens=args.output_tokens,
        requests=args.requests,
        rates=args.rate,
        block_size=args.block_size,
        duration_seconds=args.duration_seconds,
    )
    print(json.dumps(manifest, sort_keys=True))


if __name__ == "__main__":
    main()