#!/usr/bin/env python3
"""Continuous open-loop driver for the cache-size sweep.

The intent is to fill the prefix-cache from 0 up to GPU ceiling within a
single vLLM lifetime, so the per-step `cache_size` field (added by
`apply_step_timing_v2.py`) sweeps through every value the engine can
hold. Offline analysis bins by `cache_size` to recover the per-step
overhead curve.

Workload:
  - Open-loop Poisson at fixed rate (default 1.5 req/s).
  - Random per-request content (UUID + hash, calibrated to ~4096
    tokens). Zero prefix-cache hits ⇒ cache strictly grows until
    LRU eviction kicks in.
  - max_tokens / min_tokens = 256, temperature=0, ignore_eos=True.
  - Duration default 8 min.

Writes per-request metrics to `requests.jsonl`. The patch emits
per-step rows to `engine_step.jsonl` (set via AGENTIC_STEP_LOG_PATH)
and worker rows to `worker_step.jsonl.r0` (set via
CT_WORKER_STEP_LOG_PATH).

Usage:
  run_cache_sweep.py \\
      --url http://127.0.0.1:8000/v1/chat/completions \\
      --model /path/to/qwen \\
      --rate 1.5 --duration 480 \\
      --output-dir results/<date>/<config>
"""

import argparse
import asyncio
import hashlib
import json
import random
import time
import uuid
from dataclasses import asdict, dataclass
from pathlib import Path

import httpx


@dataclass
class ReqMetric:
    req_id: str
    rate_target: float
    input_tokens_target: int
    output_tokens_target: int
    t_send_ns: int
    t_first_token_ns: int | None = None
    t_last_token_ns: int | None = None
    prompt_tokens: int = 0
    completion_tokens: int = 0
    inflight_at_send: int = 0
    error: str | None = None


def make_random_prompt(target_tokens: int, rng: random.Random) -> str:
    """Same calibration as the bench_loop.py used elsewhere:
    'Block N: <32-hex>' tokenizes to ~35 tokens on Qwen3-Coder.

    Deterministic given `rng` — same RNG state produces the same prompt.
    Used so two runs with the same --seed get bit-identical request streams.
    """
    n_parts = max(1, target_tokens // 35)
    # 32 random hex chars from rng (uuid.uuid4 would use os.urandom, unseedable)
    seed = "".join(f"{rng.randrange(16):x}" for _ in range(32))
    parts = []
    for i in range(n_parts):
        h = hashlib.md5(f"{seed}_{i}".encode()).hexdigest()
        parts.append(f"Block {i}: {h}")
    return " ".join(parts)


async def send_one(client, url, model, prompt, inp_tokens, out_tokens,
                   rate, inflight, inflight_cap, fh, rid):
    if inflight[0] >= inflight_cap:
        m = ReqMetric(req_id=rid, rate_target=rate,
                      input_tokens_target=inp_tokens,
                      output_tokens_target=out_tokens,
                      t_send_ns=time.perf_counter_ns(),
                      inflight_at_send=inflight[0],
                      error="dropped_inflight_cap")
        fh.write(json.dumps(asdict(m)) + "\n")
        return

    inflight[0] += 1
    m = ReqMetric(req_id=rid, rate_target=rate,
                  input_tokens_target=inp_tokens,
                  output_tokens_target=out_tokens,
                  t_send_ns=time.perf_counter_ns(),
                  inflight_at_send=inflight[0])
    try:
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": out_tokens,
            "min_tokens": out_tokens,
            "temperature": 0,
            "ignore_eos": True,
            "stream": True,
            "stream_options": {"include_usage": True},
        }
        async with client.stream("POST", url, json=payload, timeout=600.0) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                data = line[6:]
                if data.strip() == "[DONE]":
                    break
                try:
                    chunk = json.loads(data)
                except json.JSONDecodeError:
                    continue
                usage = chunk.get("usage")
                if usage:
                    m.prompt_tokens = usage.get("prompt_tokens", m.prompt_tokens)
                    m.completion_tokens = usage.get(
                        "completion_tokens", m.completion_tokens)
                choices = chunk.get("choices") or []
                if not choices:
                    continue
                delta = choices[0].get("delta", {})
                if "role" in delta:
                    continue
                now = time.perf_counter_ns()
                if m.t_first_token_ns is None:
                    m.t_first_token_ns = now
                m.t_last_token_ns = now
    except Exception as e:
        m.error = f"{type(e).__name__}: {e}"
    finally:
        inflight[0] -= 1
    fh.write(json.dumps(asdict(m)) + "\n")


async def main_async(args):
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    req_path = out_dir / "requests.jsonl"

    inflight = [0]
    pending: list[asyncio.Task] = []
    interval_mean = 1.0 / args.rate

    # Shared seed across configs gives bit-identical arrival times AND prompt
    # content. arrival_rng feeds expovariate, content_rng feeds make_random_prompt,
    # rid_rng feeds the per-request id. All three derive from --seed so two runs
    # with the same seed are bit-identical from the producer side; only server
    # response timing differs (which is what we want to measure).
    if args.seed is not None:
        seed = args.seed
    else:
        seed = int(time.time_ns()) & 0xFFFFFFFF
    print(f"[bench] seed={seed} rate={args.rate} shape=({args.input_tokens},"
          f"{args.output_tokens}) duration={args.duration}s output={out_dir}")
    arrival_rng = random.Random(seed)
    content_rng = random.Random(seed ^ 0xC0FFEE)
    rid_rng = random.Random(seed ^ 0xDEADBEEF)

    fh = open(req_path, "a", buffering=1)
    t0 = time.perf_counter()
    last_print = t0

    async with httpx.AsyncClient(timeout=httpx.Timeout(600.0)) as client:
        # producer — prompts generated here in order so async scheduling can't
        # reorder the content RNG draw across requests.
        async def producer():
            while time.perf_counter() - t0 < args.duration:
                prompt = make_random_prompt(args.input_tokens, content_rng)
                rid = f"{rid_rng.randrange(1 << 64):016x}"
                pending.append(asyncio.create_task(
                    send_one(client, args.url, args.model, prompt,
                             args.input_tokens, args.output_tokens,
                             args.rate, inflight, args.inflight_cap, fh, rid)
                ))
                await asyncio.sleep(arrival_rng.expovariate(1.0 / interval_mean))

        # heartbeat
        async def heartbeat():
            nonlocal last_print
            while time.perf_counter() - t0 < args.duration + 1:
                now = time.perf_counter()
                if now - last_print >= 30:
                    print(f"  t+{int(now - t0):4d}s inflight={inflight[0]} "
                          f"pending={len(pending)}")
                    last_print = now
                await asyncio.sleep(2.0)

        prod = asyncio.create_task(producer())
        hb = asyncio.create_task(heartbeat())
        await prod
        hb.cancel()
        try:
            await hb
        except asyncio.CancelledError:
            pass

        # final drain
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)

    fh.close()

    # Quick summary so the orchestrator can sanity-check the cell ran.
    n_lines = sum(1 for _ in open(req_path))
    with open(out_dir / "run_summary.json", "w") as f:
        json.dump({
            "rate": args.rate,
            "input_tokens": args.input_tokens,
            "output_tokens": args.output_tokens,
            "duration_target_s": args.duration,
            "duration_actual_s": time.perf_counter() - t0,
            "n_requests": n_lines,
        }, f, indent=2)
    print(f"[done] {n_lines} requests in {time.perf_counter() - t0:.1f}s")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--rate", type=float, default=1.5)
    ap.add_argument("--input-tokens", type=int, default=4096)
    ap.add_argument("--output-tokens", type=int, default=256)
    ap.add_argument("--duration", type=float, default=480.0,
                    help="Total run duration in seconds (default 8 min)")
    ap.add_argument("--seed", type=int, default=None,
                    help="Master seed; same value across configs gives "
                         "bit-identical Poisson arrivals and prompt content. "
                         "Default: time-based (different each run).")
    ap.add_argument("--inflight-cap", type=int, default=256)
    ap.add_argument("--output-dir", required=True)
    args = ap.parse_args()
    asyncio.run(main_async(args))


if __name__ == "__main__":
    main()