#!/usr/bin/env python3
"""Open-loop fixed-shape workload for one real offered-load anchor."""

from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import http.client
import json
import math
import time
from pathlib import Path
from typing import Any


TARGET_PASS_RATE = 0.95


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--served-model", required=True)
    parser.add_argument("--model-path", type=Path, required=True)
    parser.add_argument("--rate", type=float, required=True)
    parser.add_argument("--requests", type=int, default=64)
    parser.add_argument("--input-tokens", type=int, default=2048)
    parser.add_argument("--output-tokens", type=int, default=1)
    parser.add_argument("--ttft-slo-ms", type=float)
    parser.add_argument("--tpot-slo-ms", type=float, default=150.0)
    parser.add_argument("--timeout-seconds", type=float, default=900.0)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1))
    return ordered[index]


def ttft_slo_ms(input_tokens: int) -> float:
    return 1000.0 + 1000.0 * input_tokens / 8000.0


def run_request(
    *,
    request_index: int,
    scheduled_at: float,
    benchmark_start: float,
    args: argparse.Namespace,
    prompt_ids: list[int],
) -> dict[str, Any]:
    delay = scheduled_at - time.perf_counter()
    if delay > 0:
        time.sleep(delay)
    admitted = time.perf_counter()
    record: dict[str, Any] = {
        "request_index": request_index,
        "scheduled_s": scheduled_at - benchmark_start,
        "admitted_s": admitted - benchmark_start,
        "admission_lag_ms": (admitted - scheduled_at) * 1000.0,
        "success": False,
    }
    connection = http.client.HTTPConnection(args.host, args.port, timeout=args.timeout_seconds)
    body = {
        "model": args.served_model,
        "prompt": prompt_ids,
        "min_tokens": args.output_tokens,
        "max_tokens": args.output_tokens,
        "ignore_eos": True,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
        "return_token_ids": True,
    }
    try:
        started = time.perf_counter()
        connection.request(
            "POST",
            "/v1/completions",
            body=json.dumps(body, separators=(",", ":")).encode(),
            headers={"Content-Type": "application/json"},
        )
        response = connection.getresponse()
        if response.status != 200:
            raise RuntimeError(
                f"HTTP {response.status}: {response.read().decode(errors='replace')}"
            )
        first_token_at = None
        last_token_at = None
        streamed_tokens = 0
        usage = None
        while True:
            raw = response.readline()
            if not raw:
                break
            line = raw.decode(errors="replace").strip()
            if not line.startswith("data:"):
                continue
            data = line[5:].strip()
            if data == "[DONE]":
                break
            payload = json.loads(data)
            if payload.get("usage"):
                usage = payload["usage"]
            emitted = 0
            for choice in payload.get("choices") or []:
                token_ids = choice.get("token_ids") or []
                emitted += len(token_ids) if token_ids else int(bool(choice.get("text")))
            if emitted:
                now = time.perf_counter()
                first_token_at = first_token_at or now
                last_token_at = now
                streamed_tokens += emitted
        finished = time.perf_counter()
        if first_token_at is None or usage is None:
            raise RuntimeError("missing streaming token or usage")
        prompt_tokens = int(usage["prompt_tokens"])
        completion_tokens = int(usage["completion_tokens"])
        if prompt_tokens != args.input_tokens or completion_tokens != args.output_tokens:
            raise RuntimeError(f"usage mismatch: {prompt_tokens}+{completion_tokens}")
        ttft = (first_token_at - started) * 1000.0
        tpot = (
            (last_token_at - first_token_at) * 1000.0 / (completion_tokens - 1)
            if completion_tokens > 1
            else None
        )
        record.update(
            {
                "success": True,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "streamed_token_count": streamed_tokens,
                "ttft_ms": ttft,
                "tpot_ms": tpot,
                "e2e_ms": (finished - started) * 1000.0,
                "slo_pass": ttft <= args.ttft_slo_ms
                and (tpot is None or tpot <= args.tpot_slo_ms),
            }
        )
    except Exception as error:  # Failed requests remain in the SLO denominator.
        record["error"] = f"{type(error).__name__}: {error}"
        record["slo_pass"] = False
    finally:
        connection.close()
    return record


def main() -> None:
    args = parse_args()
    if min(
        args.rate,
        args.requests,
        args.input_tokens,
        args.output_tokens,
        args.tpot_slo_ms,
    ) <= 0:
        raise ValueError("rate, requests, tokens, and SLO must be positive")
    if args.ttft_slo_ms is None:
        args.ttft_slo_ms = ttft_slo_ms(args.input_tokens)
    if args.ttft_slo_ms <= 0:
        raise ValueError("TTFT SLO must be positive")
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    excluded = set(tokenizer.all_special_ids)
    candidates = [
        token_id for token_id in range(tokenizer.vocab_size) if token_id not in excluded
    ]
    if len(candidates) < args.requests + 1:
        raise RuntimeError("tokenizer has too few non-special token IDs")
    base_id = candidates[0]
    prompts = [
        [candidates[index + 1], *([base_id] * (args.input_tokens - 1))]
        for index in range(args.requests)
    ]
    prompt_hash = hashlib.sha256(
        "\n".join(",".join(map(str, prompt)) for prompt in prompts).encode()
    ).hexdigest()

    benchmark_start = time.perf_counter() + 2.0
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.requests) as pool:
        futures = [
            pool.submit(
                run_request,
                request_index=index,
                scheduled_at=benchmark_start + index / args.rate,
                benchmark_start=benchmark_start,
                args=args,
                prompt_ids=prompts[index],
            )
            for index in range(args.requests)
        ]
        requests = [future.result() for future in futures]
    requests.sort(key=lambda row: int(row["request_index"]))
    completed = [row for row in requests if row["success"]]
    passed = sum(bool(row["slo_pass"]) for row in requests)
    ttfts = [float(row["ttft_ms"]) for row in completed]
    tpots = [
        float(row["tpot_ms"])
        for row in completed
        if row["tpot_ms"] is not None
    ]
    pass_rate = passed / len(requests)
    payload = {
        "schema": "qwen30-fixed-rate-anchor-v2",
        "workload": {
            "offered_request_rate": args.rate,
            "request_count": args.requests,
            "input_tokens": args.input_tokens,
            "output_tokens": args.output_tokens,
            "prefix_caching": False,
            "arrival": "open_loop_uniform",
            "last_scheduled_arrival_s": (args.requests - 1) / args.rate,
            "prompt_vector_sha256": prompt_hash,
        },
        "summary": {
            "completed": len(completed),
            "failed": len(requests) - len(completed),
            "ttft_p50_ms": percentile(ttfts, 0.50),
            "ttft_p95_ms": percentile(ttfts, 0.95),
            "ttft_max_ms": max(ttfts) if ttfts else None,
            "tpot_p50_ms": percentile(tpots, 0.50),
            "tpot_p95_ms": percentile(tpots, 0.95),
            "tpot_max_ms": max(tpots) if tpots else None,
            "admission_lag_max_ms": max(
                float(row["admission_lag_ms"]) for row in requests
            ),
            "slo": {
                "ttft_threshold_ms": args.ttft_slo_ms,
                "tpot_threshold_ms": (
                    args.tpot_slo_ms if args.output_tokens > 1 else None
                ),
                "passed": passed,
                "pass_rate": pass_rate,
                "feasible": pass_rate >= TARGET_PASS_RATE,
            },
        },
        "requests": requests,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload["summary"], sort_keys=True), flush=True)
    if len(completed) != args.requests:
        raise SystemExit(2)


if __name__ == "__main__":
    main()