#!/usr/bin/env python3
"""Replay a private exact-trace anchor without emitting prompt text."""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import math
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("--requests-file", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--served-model",
        default=None,
        help="Optional server-side model alias used only for HTTP routing.",
    )
    parser.add_argument("--tpot-slo-ms", type=float, default=150.0)
    parser.add_argument("--timeout-seconds", type=float, default=1800.0)
    return parser.parse_args()


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


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


def row_vector_sha256(rows: list[dict[str, Any]]) -> str:
    digest = hashlib.sha256()
    for row in rows:
        digest.update(
            json.dumps(
                [
                    row["source_index"],
                    row["arrived_at"],
                    row["input_length"],
                    row["output_length"],
                    row["session_id"],
                    row["runtime_block_ids"],
                ],
                separators=(",", ":"),
            ).encode()
        )
        digest.update(b"\n")
    return digest.hexdigest()


async def request_one(
    session: aiohttp.ClientSession,
    row: dict[str, Any],
    *,
    scheduled_at: float,
    benchmark_start: float,
    tpot_slo_ms: float,
    served_model: str | None,
) -> dict[str, Any]:
    loop = asyncio.get_running_loop()
    delay = scheduled_at - loop.time()
    if delay > 0:
        await asyncio.sleep(delay)
    admitted_at = loop.time()
    record: dict[str, Any] = {
        "source_index": int(row["source_index"]),
        "session_id": int(row["session_id"]),
        "scheduled_s": scheduled_at - benchmark_start,
        "admitted_s": admitted_at - benchmark_start,
        "admission_lag_ms": (admitted_at - scheduled_at) * 1000.0,
        "input_tokens": int(row["input_length"]),
        "requested_output_tokens": int(row["output_length"]),
        "success": False,
    }
    body = dict(row["body"])
    if served_model is not None:
        body["model"] = served_model
    body.update(
        {
            "temperature": 0,
            "stream": True,
            "stream_options": {"include_usage": True},
            "return_token_ids": True,
        }
    )
    try:
        started = loop.time()
        async with session.post("/v1/completions", json=body) as response:
            if response.status != 200:
                detail = (await response.text())[:1000]
                raise RuntimeError(f"HTTP {response.status}: {detail}")
            first_token_at = None
            last_token_at = None
            streamed_tokens = 0
            usage = None
            while True:
                raw = await response.content.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 = loop.time()
                    first_token_at = first_token_at or now
                    last_token_at = now
                    streamed_tokens += emitted
        finished = loop.time()
        if first_token_at is None or last_token_at is None or usage is None:
            raise RuntimeError("missing streaming token or usage")
        actual_input = int(usage["prompt_tokens"])
        actual_output = int(usage["completion_tokens"])
        if actual_input != int(row["input_length"]) or actual_output != int(
            row["output_length"]
        ):
            raise RuntimeError(f"usage mismatch: {actual_input}+{actual_output}")
        ttft = (first_token_at - started) * 1000.0
        tpot = (
            (last_token_at - first_token_at) * 1000.0 / (actual_output - 1)
            if actual_output > 1
            else None
        )
        record.update(
            {
                "success": True,
                "actual_input_tokens": actual_input,
                "actual_output_tokens": actual_output,
                "streamed_token_count": streamed_tokens,
                "ttft_ms": ttft,
                "tpot_ms": tpot,
                "e2e_ms": (finished - started) * 1000.0,
                "slo_pass": ttft <= ttft_slo_ms(actual_input)
                and (tpot is None or tpot <= tpot_slo_ms),
            }
        )
    except Exception as error:
        record.update(
            {
                "error": f"{type(error).__name__}: {error}",
                "slo_pass": False,
            }
        )
    return record


async def replay(args: argparse.Namespace, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    import aiohttp

    timeout = aiohttp.ClientTimeout(total=args.timeout_seconds)
    connector = aiohttp.TCPConnector(limit=0, ttl_dns_cache=300, force_close=True)
    benchmark_start = asyncio.get_running_loop().time() + 2.0
    async with aiohttp.ClientSession(
        base_url=f"http://{args.host}:{args.port}",
        timeout=timeout,
        connector=connector,
    ) as session:
        tasks = [
            asyncio.create_task(
                request_one(
                    session,
                    row,
                    scheduled_at=benchmark_start + float(row["arrived_at"]),
                    benchmark_start=benchmark_start,
                    tpot_slo_ms=args.tpot_slo_ms,
                    served_model=args.served_model,
                )
            )
            for row in rows
        ]
        return await asyncio.gather(*tasks)


def main() -> None:
    args = parse_args()
    if args.tpot_slo_ms <= 0 or args.timeout_seconds <= 0:
        raise ValueError("SLO and timeout must be positive")
    rows = [json.loads(line) for line in args.requests_file.open() if line.strip()]
    if not rows:
        raise ValueError("requests file is empty")
    arrivals = [float(row["arrived_at"]) for row in rows]
    if any(right < left for left, right in zip(arrivals, arrivals[1:])):
        raise ValueError("request arrival order drift")
    requests = asyncio.run(replay(args, rows))
    requests.sort(key=lambda row: int(row["source_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-exact-trace-anchor-v1",
        "contract": {
            "requests_file": str(args.requests_file.resolve()),
            "requests_file_sha256": hashlib.sha256(
                args.requests_file.read_bytes()
            ).hexdigest(),
            "row_vector_sha256": row_vector_sha256(rows),
            "requests": len(rows),
            "first_arrival_s": arrivals[0],
            "last_arrival_s": arrivals[-1],
            "arrival": "original_trace_timestamp_and_order",
            "input_output_prompt": "exact_source_values",
            "served_model_alias": args.served_model,
            "http_connection_reuse": False,
            "ttft_slo": "1000ms + 1000ms * input_tokens / 8000",
            "tpot_slo_ms": args.tpot_slo_ms,
            "target_pass_rate": TARGET_PASS_RATE,
            "prompt_text_emitted": False,
        },
        "summary": {
            "completed": len(completed),
            "failed": len(requests) - len(completed),
            "passed": passed,
            "pass_rate": pass_rate,
            "feasible": pass_rate >= TARGET_PASS_RATE,
            "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_p95_ms": percentile(
                [float(row["admission_lag_ms"]) for row in requests], 0.95
            ),
        },
        "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) != len(requests):
        raise SystemExit(2)


if __name__ == "__main__":
    main()
