#!/usr/bin/env python3
"""Token-exact fixed-duration client for the OpProf Phase-3 protocol."""

from __future__ import annotations

import argparse
import asyncio
import gzip
import hashlib
import json
import math
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import aiohttp

SCHEMA = 1
TOKEN_BASE = 1000
TOKEN_SPAN = 100000


class ManifestExhausted(RuntimeError):
    pass


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def atomic_json(path: Path, value: Any, mode: int = 0o640) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(value, f, sort_keys=True, indent=2)
        f.write("\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def atomic_jsonl(path: Path, rows: list[dict[str, Any]], mode: int = 0o640) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def parse_range(value: str) -> tuple[int, int]:
    lo_text, hi_text = value.split(":", 1)
    lo, hi = int(lo_text), int(hi_text)
    if lo <= 0 or hi < lo:
        raise argparse.ArgumentTypeError(f"invalid positive range: {value}")
    return lo, hi


def _integer_counts(weights: list[float], total: int) -> list[int]:
    raw = [w * total for w in weights]
    counts = [math.floor(x) for x in raw]
    order = sorted(
        range(len(raw)), key=lambda i: raw[i] - counts[i], reverse=True
    )
    for idx in order[: total - sum(counts)]:
        counts[idx] += 1
    return counts


def numeric_sanity(values: list[float | int]) -> dict[str, Any]:
    finite = [float(x) for x in values if math.isfinite(float(x))]
    return {
        "n": len(values),
        "finite_n": len(finite),
        "missing_n": len(values) - len(finite),
        "min": min(finite) if finite else None,
        "max": max(finite) if finite else None,
        "distinct_n": len(set(finite)),
    }


def manifest_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "schema": SCHEMA,
        "rows": len(rows),
        "input_tokens": numeric_sanity([int(r["input_tokens"]) for r in rows]),
        "output_tokens": numeric_sanity([int(r["output_tokens"]) for r in rows]),
        "arrival_values": sorted({str(r["arrival"]) for r in rows}),
        "pattern_values": sorted({str(r["pattern_id"]) for r in rows}),
    }


def materialize(args: argparse.Namespace) -> dict[str, Any]:
    import numpy as np

    rng = np.random.default_rng(args.workload_seed)
    n = args.num_requests
    if args.kind == "prefix-pool":
        if args.num_prefixes <= 0 or args.prefix_len <= 0 or args.suffix_fixed <= 0:
            raise ValueError("prefix-pool requires positive pool/prefix/suffix")
        lengths = np.full(n, args.prefix_len + args.suffix_fixed, dtype=np.int64)
        prefix_ids = np.arange(n, dtype=np.int64) % args.num_prefixes
        rng.shuffle(prefix_ids)
    else:
        prefix_ids = np.full(n, -1, dtype=np.int64)
        if args.input_uniform:
            lo, hi = parse_range(args.input_uniform)
            lengths = rng.integers(lo, hi + 1, n, dtype=np.int64)
        elif args.input_fixed:
            lengths = np.full(n, args.input_fixed, dtype=np.int64)
        elif args.input_mixture:
            spec = json.loads(args.input_mixture)
            if not isinstance(spec, dict) or not spec:
                raise ValueError("input mixture must be a non-empty JSON object")
            keys = list(spec)
            weights = [float(spec[key]) for key in keys]
            if any(w < 0 for w in weights) or not math.isclose(sum(weights), 1.0):
                raise ValueError("mixture weights must be non-negative and sum to 1")
            pieces = []
            for key, count in zip(
                keys, _integer_counts(weights, n), strict=True
            ):
                kind, lo_text, hi_text = key.split(":")
                if kind != "uniform":
                    raise ValueError(f"unsupported mixture component: {key}")
                pieces.append(
                    rng.integers(
                        int(lo_text), int(hi_text) + 1, count, dtype=np.int64
                    )
                )
            lengths = np.concatenate(pieces)
            rng.shuffle(lengths)
        else:
            raise ValueError("exactly one input distribution is required")
    if args.output_fixed <= 0 or args.arrival not in {"steady", "burst:8"}:
        raise ValueError("invalid output length or arrival class")

    rows = []
    for i in range(n):
        row = {
            "schema": SCHEMA,
            "request_id": f"{args.id}-{i:05d}",
            "pattern_id": args.id,
            "kind": args.kind,
            "input_tokens": int(lengths[i]),
            "output_tokens": args.output_fixed,
            "arrival": args.arrival,
            "token_seed": int(args.workload_seed * 1000003 + i),
        }
        if args.kind == "prefix-pool":
            row.update(
                {
                    "prefix_id": int(prefix_ids[i]),
                    "num_prefixes": args.num_prefixes,
                    "prefix_tokens": args.prefix_len,
                }
            )
        rows.append(row)

    out = Path(args.out)
    atomic_jsonl(out, rows, mode=0o600)
    summary = manifest_summary(rows)
    summary.update({"sha256": sha256_file(out), "path": str(out)})
    atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600)
    print(json.dumps(summary, sort_keys=True))
    return summary


def materialize_private(args: argparse.Namespace) -> dict[str, Any]:
    from transformers import AutoTokenizer

    source = Path(args.source)
    selected: list[dict[str, Any]] = []
    with source.open(encoding="utf-8") as f:
        for source_index, line in enumerate(f):
            row = json.loads(line)
            if (
                float(row["sampling_u"]) <= args.sampling_u_max
                and int(row["input_length"]) <= args.max_input_tokens
            ):
                selected.append(
                    {
                        "schema": SCHEMA,
                        "request_id": f"{args.id}-{len(selected):05d}",
                        "pattern_id": args.id,
                        "kind": "private-trace",
                        "input_tokens": int(row["input_length"]),
                        "output_tokens": min(
                            int(row["output_length"]), args.output_cap
                        ),
                        "arrival": args.arrival,
                        "source_index": source_index,
                        "prompt": row["prompt"],
                    }
                )

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    diffs = [
        len(tokenizer.encode(row["prompt"], add_special_tokens=False))
        - row["input_tokens"]
        for row in selected
    ]
    exact = sum(diff == 0 for diff in diffs)
    exact_fraction = exact / len(diffs) if diffs else 0.0
    max_abs = max((abs(diff) for diff in diffs), default=-1)
    if exact_fraction < 0.99 or max_abs > 1:
        raise RuntimeError(
            "tokenizer parity gate failed: "
            f"exact_fraction={exact_fraction:.6f} max_abs_error={max_abs}"
        )

    out = Path(args.out)
    atomic_jsonl(out, selected, mode=0o600)
    summary = manifest_summary(selected)
    summary.update(
        {
            "sha256": sha256_file(out),
            "source_sha256": sha256_file(source),
            "tokenizer_exact_n": exact,
            "tokenizer_exact_fraction": exact_fraction,
            "tokenizer_max_abs_error": max_abs,
            "path": str(out),
        }
    )
    atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600)
    print(json.dumps(summary, sort_keys=True))
    return summary


def load_manifest(path: Path) -> list[dict[str, Any]]:
    rows = [json.loads(line) for line in path.read_text().splitlines() if line]
    required = {
        "request_id",
        "pattern_id",
        "input_tokens",
        "output_tokens",
        "arrival",
    }
    if not rows:
        raise ValueError("empty manifest")
    for row in rows:
        if not required.issubset(row):
            raise ValueError(f"manifest row lacks {sorted(required - set(row))}")
    if len({row["request_id"] for row in rows}) != len(rows):
        raise ValueError("duplicate request_id")
    return rows


def _token_stream(seed: int, count: int) -> list[int]:
    state = seed & 0xFFFFFFFF
    out = []
    for _ in range(count):
        state = (1664525 * state + 1013904223) & 0xFFFFFFFF
        out.append(TOKEN_BASE + state % TOKEN_SPAN)
    return out


def synthetic_prompt(row: dict[str, Any]) -> list[int]:
    length = int(row["input_tokens"])
    seed = int(row.get("token_seed", 0))
    if row.get("kind") == "prefix-pool":
        prefix_n = int(row["prefix_tokens"])
        tokens = _token_stream(0xA5A50000 + int(row["prefix_id"]), prefix_n)
        tokens += _token_stream(seed, length - prefix_n)
        offset = prefix_n
    else:
        tokens = _token_stream(seed, length)
        offset = 0
    if length - offset >= 3:
        index = int(row["request_id"].rsplit("-", 1)[1])
        tokens[offset : offset + 3] = [
            TOKEN_BASE + index % 100,
            TOKEN_BASE + (index // 100) % 100,
            TOKEN_BASE + (index // 10000) % 100,
        ]
    return tokens


@dataclass
class RunContext:
    args: argparse.Namespace
    rows: list[dict[str, Any]]
    t0: float
    clean_end: float
    stop_event: asyncio.Event
    lock: asyncio.Lock
    next_index: int = 0
    in_flight: int = 0
    max_in_flight: int = 0
    exhausted: bool = False
    admission_stop_s: float | None = None

    async def next_row(self) -> dict[str, Any]:
        async with self.lock:
            if self.next_index >= len(self.rows):
                self.exhausted = True
                raise ManifestExhausted(
                    f"manifest exhausted after {self.next_index} admissions"
                )
            row = self.rows[self.next_index]
            self.next_index += 1
            return row


async def request_one(
    ctx: RunContext,
    session: aiohttp.ClientSession,
    row: dict[str, Any],
    scheduled: float,
) -> dict[str, Any]:
    loop = asyncio.get_running_loop()
    admitted = loop.time()
    ctx.in_flight += 1
    ctx.max_in_flight = max(ctx.max_in_flight, ctx.in_flight)
    status = 0
    actual_output: int | None = None
    first_token: float | None = None
    error_kind: str | None = None
    try:
        prompt: str | list[int] = (
            row["prompt"]
            if row.get("kind") == "private-trace"
            else synthetic_prompt(row)
        )
        if not isinstance(prompt, str) and len(prompt) != int(row["input_tokens"]):
            raise AssertionError("synthetic prompt length drift")
        payload = {
            "model": ctx.args.model,
            "prompt": prompt,
            "max_tokens": int(row["output_tokens"]),
            "temperature": ctx.args.temperature,
            "ignore_eos": ctx.args.ignore_eos,
            "stream": True,
            "stream_options": {"include_usage": True},
            "add_special_tokens": False,
            "seed": ctx.args.server_seed,
        }
        headers = {
            "Content-Type": "application/json",
            "x-request-id": str(row["request_id"]),
        }
        async with session.post(
            ctx.args.base_url.rstrip("/") + "/v1/completions",
            json=payload,
            headers=headers,
        ) as response:
            status = response.status
            if status != 200:
                error_kind = f"http_{status}"
            else:
                buf = b""
                async for chunk in response.content.iter_any():
                    buf += chunk
                    while b"\n" in buf:
                        line, buf = buf.split(b"\n", 1)
                        line = line.strip()
                        if not line.startswith(b"data:"):
                            continue
                        data = line[5:].strip()
                        if data == b"[DONE]":
                            continue
                        event = json.loads(data)
                        if event.get("choices") and first_token is None:
                            first_token = loop.time()
                        if event.get("usage") is not None:
                            actual_output = int(event["usage"]["completion_tokens"])
                if actual_output is None:
                    error_kind = "missing_usage"
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        error_kind = type(exc).__name__
    except Exception as exc:
        error_kind = type(exc).__name__
    finally:
        completed = loop.time()
        ctx.in_flight -= 1
    success = (
        status == 200
        and error_kind is None
        and actual_output == int(row["output_tokens"])
    )
    if status == 200 and actual_output is not None and not success:
        error_kind = "output_token_mismatch"
    return {
        "schema": SCHEMA,
        "request_id": row["request_id"],
        "scheduled_s": scheduled - ctx.t0,
        "admitted_s": admitted - ctx.t0,
        "first_token_s": None if first_token is None else first_token - ctx.t0,
        "completed_s": completed - ctx.t0,
        "input_tokens": int(row["input_tokens"]),
        "requested_output_tokens": int(row["output_tokens"]),
        "actual_output_tokens": actual_output,
        "http_status": status,
        "success": success,
        "error_kind": error_kind,
    }


async def saturation_load(
    ctx: RunContext, session: aiohttp.ClientSession
) -> list[dict[str, Any]]:
    results: list[dict[str, Any]] = []

    async def worker() -> None:
        while not ctx.stop_event.is_set():
            try:
                row = await ctx.next_row()
            except ManifestExhausted:
                ctx.stop_event.set()
                return
            results.append(
                await request_one(ctx, session, row, asyncio.get_running_loop().time())
            )

    tasks = [
        asyncio.create_task(worker()) for _ in range(ctx.args.max_concurrency)
    ]
    await asyncio.gather(*tasks)
    return results


async def finite_load(
    ctx: RunContext, session: aiohttp.ClientSession, rate: float
) -> list[dict[str, Any]]:
    sem = asyncio.Semaphore(ctx.args.max_concurrency)
    tasks: list[asyncio.Task[dict[str, Any]]] = []
    batch = 8 if str(ctx.rows[0]["arrival"]) == "burst:8" else 1
    period = batch / rate
    event_index = 0

    async def limited(row: dict[str, Any], scheduled: float) -> dict[str, Any]:
        async with sem:
            return await request_one(ctx, session, row, scheduled)

    while not ctx.stop_event.is_set():
        scheduled = ctx.t0 + event_index * period
        delay = scheduled - asyncio.get_running_loop().time()
        if delay > 0:
            try:
                await asyncio.wait_for(ctx.stop_event.wait(), timeout=delay)
                break
            except asyncio.TimeoutError:
                pass
        if ctx.stop_event.is_set():
            break
        try:
            for _ in range(batch):
                tasks.append(
                    asyncio.create_task(limited(await ctx.next_row(), scheduled))
                )
        except ManifestExhausted:
            ctx.stop_event.set()
            break
        event_index += 1
    return await asyncio.gather(*tasks) if tasks else []


async def post_profile(
    session: aiohttp.ClientSession, base_url: str, endpoint: str
) -> tuple[float, float, int]:
    loop = asyncio.get_running_loop()
    before = loop.time()
    async with session.post(base_url.rstrip("/") + endpoint) as response:
        status = response.status
        await response.read()
    return before, loop.time(), status


def _trace_loadable(path: Path) -> bool:
    try:
        opener = gzip.open if path.suffix == ".gz" else open
        with opener(path, "rt", encoding="utf-8") as f:
            parsed = json.load(f)
        return isinstance(parsed, dict) and isinstance(parsed.get("traceEvents"), list)
    except (OSError, EOFError, json.JSONDecodeError):
        return False


async def wait_new_trace(
    trace_dir: Path, before: set[Path], timeout: float
) -> Path:
    deadline = asyncio.get_running_loop().time() + timeout
    while asyncio.get_running_loop().time() < deadline:
        for path in sorted(set(trace_dir.glob("*.pt.trace.json*")) - before):
            if _trace_loadable(path):
                return path
        await asyncio.sleep(0.25)
    raise TimeoutError(f"no new loadable trace within {timeout}s")


async def timeline(
    ctx: RunContext, session: aiohttp.ClientSession
) -> list[dict[str, Any]]:
    args = ctx.args
    profiles: list[dict[str, Any]] = []
    await asyncio.sleep(max(0, ctx.clean_end - asyncio.get_running_loop().time()))
    if args.profile_after_clean:
        trace_dir = Path(args.profile_trace_dir)
        for window in range(args.num_profile_windows):
            prior = set(trace_dir.glob("*.pt.trace.json*"))
            start_before, start_after, start_status = await post_profile(
                session, args.base_url, "/start_profile"
            )
            trace = await wait_new_trace(
                trace_dir, prior, args.profile_timeout_seconds
            )
            trace_ready = asyncio.get_running_loop().time()
            stop_before, stop_after, stop_status = await post_profile(
                session, args.base_url, "/stop_profile"
            )
            profiles.append(
                {
                    "window": window + 1,
                    "start_call_s": start_before - ctx.t0,
                    "start_return_s": start_after - ctx.t0,
                    "trace_ready_s": trace_ready - ctx.t0,
                    "stop_call_s": stop_before - ctx.t0,
                    "stop_return_s": stop_after - ctx.t0,
                    "start_status": start_status,
                    "stop_status": stop_status,
                    "trace_file": trace.name,
                    "trace_sha256": sha256_file(trace),
                }
            )
            if start_status != 200 or stop_status != 200:
                raise RuntimeError("profile endpoint returned non-200")
            await asyncio.sleep(args.recovery_seconds)
    else:
        await asyncio.sleep(args.post_clean_seconds)
    ctx.admission_stop_s = asyncio.get_running_loop().time() - ctx.t0
    ctx.stop_event.set()
    return profiles


def segment_summary(
    records: list[dict[str, Any]], start: float, end: float
) -> dict[str, Any]:
    admitted = [r for r in records if start <= r["admitted_s"] < end]
    completed = [r for r in records if start <= r["completed_s"] < end]
    successes = [r for r in completed if r["success"]]
    duration = end - start
    return {
        "start_s": start,
        "end_s": end,
        "duration_s": duration,
        "admitted": len(admitted),
        "completed": len(successes),
        "failed": len(completed) - len(successes),
        "offered_rps": len(admitted) / duration,
        "completed_throughput_rps": len(successes) / duration,
        "input_tokens": sum(r["input_tokens"] for r in successes),
        "output_tokens": sum(r["actual_output_tokens"] or 0 for r in successes),
    }


async def run_load(args: argparse.Namespace) -> dict[str, Any]:
    manifest = Path(args.manifest)
    rows = load_manifest(manifest)
    arrivals = {row["arrival"] for row in rows}
    if len(arrivals) != 1:
        raise ValueError("a manifest must have one arrival class")
    if args.load_point == "saturation":
        if args.request_rate != "inf":
            raise ValueError("saturation requires --request-rate inf")
        rate = math.inf
    else:
        if not args.saturation_result:
            raise ValueError("moderate requires --saturation-result")
        sat = json.loads(Path(args.saturation_result).read_text())
        rate = args.rate_fraction * float(sat["clean"]["completed_throughput_rps"])
        if not math.isfinite(rate) or rate <= 0:
            raise ValueError("derived moderate rate must be positive and finite")

    loop = asyncio.get_running_loop()
    t0 = loop.time()
    clean_seconds = args.clean_segment_seconds * args.num_clean_segments
    ctx = RunContext(
        args=args,
        rows=rows,
        t0=t0,
        clean_end=t0 + args.warmup_seconds + clean_seconds,
        stop_event=asyncio.Event(),
        lock=asyncio.Lock(),
    )
    timeout = aiohttp.ClientTimeout(total=None, connect=30, sock_read=600)
    connector = aiohttp.TCPConnector(limit=args.max_concurrency)
    async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
        profile_task = asyncio.create_task(timeline(ctx, session))
        load_task = asyncio.create_task(
            saturation_load(ctx, session)
            if math.isinf(rate)
            else finite_load(ctx, session, rate)
        )
        try:
            profiles = await profile_task
        except Exception:
            ctx.stop_event.set()
            await load_task
            raise
        records = await load_task

    clean_start = args.warmup_seconds
    clean_end = clean_start + clean_seconds
    clean = segment_summary(records, clean_start, clean_end)
    segments = []
    for i in range(args.num_clean_segments):
        start = clean_start + i * args.clean_segment_seconds
        segments.append(
            {
                "name": chr(ord("A") + i),
                **segment_summary(
                    records, start, start + args.clean_segment_seconds
                ),
            }
        )
    successful = [r for r in records if r["success"]]
    elapsed_seconds = loop.time() - t0
    if ctx.admission_stop_s is None:
        raise RuntimeError("admission stop timestamp was not recorded")
    drain_seconds = elapsed_seconds - ctx.admission_stop_s
    result = {
        "schema": SCHEMA,
        "manifest_sha256": sha256_file(manifest),
        "manifest_rows": len(rows),
        "manifest_admitted": ctx.next_index,
        "manifest_wrapped": False,
        "manifest_exhausted": ctx.exhausted,
        "load_point": args.load_point,
        "request_rate": "inf" if math.isinf(rate) else rate,
        "rate_fraction": None if math.isinf(rate) else args.rate_fraction,
        "arrival": next(iter(arrivals)),
        "warmup_seconds": args.warmup_seconds,
        "clean_segment_seconds": args.clean_segment_seconds,
        "num_clean_segments": args.num_clean_segments,
        "elapsed_seconds": elapsed_seconds,
        "admission_stop_s": ctx.admission_stop_s,
        "drain_seconds": drain_seconds,
        "max_in_flight": ctx.max_in_flight,
        "records": len(records),
        "successful_records": len(successful),
        "failed_records": len(records) - len(successful),
        "clean": clean,
        "segments": segments,
        "profiles": profiles,
    }
    sanity = {
        "schema": SCHEMA,
        "numeric": {
            "input_tokens": numeric_sanity([r["input_tokens"] for r in records]),
            "requested_output_tokens": numeric_sanity(
                [r["requested_output_tokens"] for r in records]
            ),
            "actual_output_tokens": numeric_sanity(
                [
                    r["actual_output_tokens"]
                    for r in records
                    if r["actual_output_tokens"] is not None
                ]
            ),
            "scheduled_s": numeric_sanity([r["scheduled_s"] for r in records]),
            "admitted_s": numeric_sanity([r["admitted_s"] for r in records]),
            "completed_s": numeric_sanity([r["completed_s"] for r in records]),
        },
        "invariants": {
            "clean_duration_exact": math.isclose(clean["duration_s"], clean_seconds),
            "segment_count_exact": len(segments) == args.num_clean_segments,
            "manifest_no_wrap": ctx.next_index <= len(rows),
            "manifest_not_exhausted": not ctx.exhausted,
            "concurrency_bounded": ctx.max_in_flight <= args.max_concurrency,
            "drain_within_timeout": drain_seconds <= args.drain_timeout_seconds,
            "output_tokens_exact": all(
                r["actual_output_tokens"] == r["requested_output_tokens"]
                for r in successful
            ),
            "clean_failures_zero": clean["failed"] == 0,
            "profile_count_exact": len(profiles)
            == (args.num_profile_windows if args.profile_after_clean else 0),
            "profile_status_ok": all(
                p["start_status"] == 200 and p["stop_status"] == 200
                for p in profiles
            ),
        },
    }
    if not math.isinf(rate):
        sanity["invariants"]["moderate_offered_within_5pct"] = (
            abs(clean["offered_rps"] / rate - 1) <= 0.05
        )
    out = Path(args.result_dir)
    out.mkdir(parents=True, exist_ok=True)
    atomic_jsonl(out / "requests.jsonl", sorted(records, key=lambda r: r["admitted_s"]))
    atomic_jsonl(out / "segments.jsonl", segments)
    atomic_json(out / "result.json", result)
    atomic_json(out / "sanity.json", sanity)
    if ctx.exhausted:
        raise ManifestExhausted("manifest exhausted; result retained for diagnosis")
    failed = [name for name, ok in sanity["invariants"].items() if not ok]
    if failed:
        raise RuntimeError(f"client sanity failure: {failed}")
    return result


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)
    mat = sub.add_parser("materialize")
    mat.add_argument("--id", required=True)
    mat.add_argument("--kind", choices=("synthetic", "prefix-pool"), required=True)
    group = mat.add_mutually_exclusive_group()
    group.add_argument("--input-uniform")
    group.add_argument("--input-fixed", type=int)
    group.add_argument("--input-mixture")
    mat.add_argument("--output-fixed", type=int, required=True)
    mat.add_argument("--prefix", default="none")
    mat.add_argument("--arrival", required=True)
    mat.add_argument("--num-requests", type=int, required=True)
    mat.add_argument("--workload-seed", type=int, required=True)
    mat.add_argument("--num-prefixes", type=int, default=0)
    mat.add_argument("--prefix-len", type=int, default=0)
    mat.add_argument("--suffix-fixed", type=int, default=0)
    mat.add_argument("--out", required=True)
    private = sub.add_parser("materialize-private")
    private.add_argument("--id", required=True)
    private.add_argument("--source", required=True)
    private.add_argument("--sampling-u-max", type=float, required=True)
    private.add_argument("--max-input-tokens", type=int, required=True)
    private.add_argument("--output-cap", type=int, required=True)
    private.add_argument("--preserve-prompts", action="store_true", required=True)
    private.add_argument("--disable-shuffle", action="store_true", required=True)
    private.add_argument("--arrival", required=True)
    private.add_argument("--model", required=True)
    private.add_argument("--out", required=True)
    run = sub.add_parser("run")
    run.add_argument("--manifest", required=True)
    run.add_argument("--base-url", required=True)
    run.add_argument("--model", required=True)
    run.add_argument("--load-point", choices=("saturation", "moderate"), required=True)
    run.add_argument("--request-rate")
    run.add_argument("--saturation-result")
    run.add_argument("--rate-fraction", type=float, default=0.60)
    run.add_argument("--max-concurrency", type=int, default=256)
    run.add_argument("--ignore-eos", action="store_true")
    run.add_argument("--temperature", type=float, default=0.0)
    run.add_argument("--warmup-seconds", type=float, default=60)
    run.add_argument("--clean-segment-seconds", type=float, default=80)
    run.add_argument("--num-clean-segments", type=int, default=3)
    run.add_argument("--profile-after-clean", action="store_true")
    run.add_argument("--num-profile-windows", type=int, default=0)
    run.add_argument("--profile-warmup-iterations", type=int, default=2)
    run.add_argument("--profile-active-iterations", type=int, default=8)
    run.add_argument("--profile-trace-dir")
    run.add_argument("--profile-timeout-seconds", type=float, default=120)
    run.add_argument("--recovery-seconds", type=float, default=30)
    run.add_argument("--post-clean-seconds", type=float, default=0)
    run.add_argument("--drain-timeout-seconds", type=float, default=120)
    run.add_argument("--workload-seed", type=int, default=20260712)
    run.add_argument("--server-seed", type=int, default=20260712)
    run.add_argument("--result-dir", required=True)
    return parser


def main() -> None:
    args = build_parser().parse_args()
    if args.command == "materialize":
        materialize(args)
    elif args.command == "materialize-private":
        materialize_private(args)
    else:
        if args.profile_after_clean and not args.profile_trace_dir:
            raise ValueError("--profile-after-clean requires --profile-trace-dir")
        print(json.dumps(asyncio.run(run_load(args)), sort_keys=True))


if __name__ == "__main__":
    main()
