#!/usr/bin/env python3
"""Issue exact fixed-shape completion requests and record streaming latency."""

from __future__ import annotations

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


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("--input-tokens", type=int, default=2048)
    parser.add_argument("--output-tokens", type=int, default=128)
    parser.add_argument("--concurrency", type=int, required=True)
    parser.add_argument("--requests", type=int, required=True)
    parser.add_argument("--timeout-seconds", type=float, default=600.0)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def run_request(
    *,
    request_index: int,
    args: argparse.Namespace,
    prompt_token_id: int,
    start_barrier: threading.Barrier,
) -> dict[str, Any]:
    body = {
        "model": args.served_model,
        "prompt": [prompt_token_id] * args.input_tokens,
        "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,
    }
    encoded = json.dumps(body, separators=(",", ":")).encode()
    connection = http.client.HTTPConnection(
        args.host, args.port, timeout=args.timeout_seconds
    )
    start_barrier.wait()
    started = time.perf_counter()
    connection.request(
        "POST",
        "/v1/completions",
        body=encoded,
        headers={"Content-Type": "application/json"},
    )
    response = connection.getresponse()
    if response.status != 200:
        detail = response.read().decode(errors="replace")
        raise RuntimeError(f"request {request_index} failed: HTTP {response.status}: {detail}")

    first_token_at: float | None = None
    last_token_at: float | None = None
    streamed_token_count = 0
    usage: dict[str, Any] | None = 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 []
            if token_ids:
                emitted += len(token_ids)
            elif choice.get("text"):
                emitted += 1
        if emitted:
            now = time.perf_counter()
            if first_token_at is None:
                first_token_at = now
            last_token_at = now
            streamed_token_count += emitted
    finished = time.perf_counter()
    connection.close()

    if first_token_at is None or last_token_at is None or usage is None:
        raise RuntimeError(
            f"request {request_index} missing streaming token or usage metadata"
        )
    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"request {request_index} usage mismatch: prompt={prompt_tokens}, "
            f"completion={completion_tokens}"
        )

    ttft_ms = (first_token_at - started) * 1000.0
    tpot_ms = (
        (last_token_at - first_token_at) * 1000.0 / (completion_tokens - 1)
        if completion_tokens > 1
        else 0.0
    )
    return {
        "request_index": request_index,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "streamed_token_count": streamed_token_count,
        "ttft_ms": ttft_ms,
        "tpot_ms": tpot_ms,
        "e2e_ms": (finished - started) * 1000.0,
        "ttft_slo_ms": 1000.0 + args.input_tokens / 8.0,
        "tpot_slo_ms": 40.0,
        "joint_slo_pass": ttft_ms <= 1000.0 + args.input_tokens / 8.0
        and tpot_ms <= 40.0,
    }


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * fraction)))
    return ordered[index]


def main() -> None:
    args = parse_args()
    if args.concurrency <= 0 or args.requests < args.concurrency:
        raise ValueError("requests must be at least concurrency, and both must be positive")
    if args.input_tokens <= 0 or args.output_tokens <= 0:
        raise ValueError("token lengths must be positive")

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    candidate_ids = tokenizer.encode(" hello", add_special_tokens=False)
    if not candidate_ids:
        raise RuntimeError("tokenizer returned no prompt token id")
    prompt_token_id = int(candidate_ids[0])

    results: list[dict[str, Any]] = []
    for batch_start in range(0, args.requests, args.concurrency):
        batch_count = min(args.concurrency, args.requests - batch_start)
        barrier = threading.Barrier(batch_count)
        with concurrent.futures.ThreadPoolExecutor(max_workers=batch_count) as pool:
            futures = [
                pool.submit(
                    run_request,
                    request_index=batch_start + offset,
                    args=args,
                    prompt_token_id=prompt_token_id,
                    start_barrier=barrier,
                )
                for offset in range(batch_count)
            ]
            results.extend(future.result() for future in futures)

    ttfts = [float(row["ttft_ms"]) for row in results]
    tpots = [float(row["tpot_ms"]) for row in results]
    payload = {
        "schema": "qwen235b-t0-smoke-v1",
        "workload": {
            "input_tokens": args.input_tokens,
            "output_tokens": args.output_tokens,
            "uniform_qps": None,
            "prefix_caching": False,
            "concurrency": args.concurrency,
            "request_count": args.requests,
            "prompt_token_id": prompt_token_id,
        },
        "summary": {
            "completed_requests": len(results),
            "joint_slo_pass_count": sum(bool(row["joint_slo_pass"]) for row in results),
            "ttft_mean_ms": statistics.fmean(ttfts),
            "ttft_p95_ms": percentile(ttfts, 0.95),
            "tpot_mean_ms": statistics.fmean(tpots),
            "tpot_p95_ms": percentile(tpots, 0.95),
        },
        "requests": sorted(results, key=lambda row: int(row["request_index"])),
    }
    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 __name__ == "__main__":
    main()
