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

from __future__ import annotations

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


TPOT_SLOS_MS = (40.0, 120.0, 150.0, 180.0)
TTFT_SLO_MS = 1256.0


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=128)
    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 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 last_token_at is None or usage is None:
            raise RuntimeError("missing streaming tokens 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)
        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,
            }
        )
    except Exception as error:  # Preserve failed requests as SLO failures.
        record["error"] = f"{type(error).__name__}: {error}"
    finally:
        connection.close()
    return record


def main() -> None:
    args = parse_args()
    if args.rate <= 0 or args.requests <= 0:
        raise ValueError("rate and requests 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"]]
    slos = {}
    for limit in TPOT_SLOS_MS:
        passed = sum(
            row["success"]
            and float(row["ttft_ms"]) <= TTFT_SLO_MS
            and float(row["tpot_ms"]) <= limit
            for row in requests
        )
        slos[f"tpot_{int(limit)}ms"] = {
            "passed": passed,
            "pass_rate": passed / len(requests),
            "feasible": passed / len(requests) >= 0.95,
        }
    ttfts = [float(row["ttft_ms"]) for row in completed]
    tpots = [float(row["tpot_ms"]) for row in completed]
    payload = {
        "schema": "qwen235b-t0-rate-anchor-v1",
        "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),
            "tpot_p50_ms": percentile(tpots, 0.50),
            "tpot_p95_ms": percentile(tpots, 0.95),
            "admission_lag_max_ms": max(float(row["admission_lag_ms"]) for row in requests),
            "slos": slos,
        },
        "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()
