"""Lightweight benchmark for TTS/ASR queueing and API throughput.

This benchmark uses the FastAPI app with a mocked upstream speech API so it
measures this project's queueing, request handling, and SSE delivery cost
without requiring a real external model endpoint.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import statistics
import time
from pathlib import Path

import httpx

BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in __import__("sys").path:
    __import__("sys").path.insert(0, str(BACKEND_DIR))

import main  # noqa: E402
import tts_asr  # noqa: E402
from job_system import reset_job_manager  # noqa: E402


def _wav_bytes(duration_ms: int = 320) -> bytes:
    sample_rate = 16000
    frames = max(1, int(sample_rate * duration_ms / 1000))
    data = b"".join((i % 32768).to_bytes(2, "little", signed=False) for i in range(frames))
    data_size = len(data)
    return (
        b"RIFF" + (36 + data_size).to_bytes(4, "little")
        + b"WAVE"
        + b"fmt " + (16).to_bytes(4, "little")
        + (1).to_bytes(2, "little")
        + (1).to_bytes(2, "little")
        + sample_rate.to_bytes(4, "little")
        + sample_rate.to_bytes(4, "little")
        + (2).to_bytes(2, "little")
        + (16).to_bytes(2, "little")
        + b"data" + data_size.to_bytes(4, "little")
        + data
    )


def _parse_sse_done(text: str) -> dict:
    for chunk in reversed([item for item in text.split("\n\n") if item.strip()]):
        event = ""
        data = ""
        for line in chunk.splitlines():
            if line.startswith("event:"):
                event = line.split(":", 1)[1].strip()
            elif line.startswith("data:"):
                data = line.split(":", 1)[1].strip()
        if event == "done" and data:
            payload = json.loads(data)
            result = dict(payload.get("result") or {})
            for key in ("queue_ms", "run_ms", "total_ms", "queued_count", "running_count", "busy_level", "busy_ratio"):
                if key in payload:
                    result[key] = payload[key]
            return result
    raise RuntimeError("done event not found")


def _percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    if len(values) == 1:
        return values[0]
    index = (len(values) - 1) * q
    lower = int(index)
    upper = min(lower + 1, len(values) - 1)
    if lower == upper:
        return values[lower]
    weight = index - lower
    return values[lower] * (1 - weight) + values[upper] * weight


async def _build_mock_client(tts_delay_ms: int, asr_delay_ms: int) -> httpx.AsyncClient:
    async def transport(request: httpx.Request):
        if request.url.path.endswith("/audio/speech"):
            await asyncio.sleep(tts_delay_ms / 1000.0)
            return httpx.Response(200, content=_wav_bytes(420), headers={"x-request-id": "bench-tts"}, request=request)
        await asyncio.sleep(asr_delay_ms / 1000.0)
        return httpx.Response(200, json={"text": "benchmark transcript", "language": "zh"}, headers={"x-request-id": "bench-asr"}, request=request)

    return httpx.AsyncClient(
        base_url="https://benchmark.example/v1/",
        transport=httpx.MockTransport(transport),
    )


async def _run_case(case_name: str, concurrency: int, request_count: int, audio_b64: str | None = None) -> dict:
    results: list[dict] = []
    latencies: list[float] = []

    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=main.app),
        base_url="http://testserver",
        timeout=120.0,
    ) as client:
        semaphore = asyncio.Semaphore(concurrency)

        async def fire(index: int) -> None:
            async with semaphore:
                started = time.perf_counter()
                if case_name == "tts":
                    response = await client.post(
                        "/v1/tts-asr/tts",
                        json={"text": f"第 {index} 条基准文本", "speaker": "Vivian", "format": "wav"},
                    )
                else:
                    response = await client.post(
                        "/v1/tts-asr/asr",
                        json={"audio_base64": audio_b64, "language": "zh-CN"},
                    )
                response.raise_for_status()
                payload = _parse_sse_done(response.text)
                latencies.append((time.perf_counter() - started) * 1000.0)
                results.append(payload)

        await asyncio.gather(*(fire(index) for index in range(request_count)))

    queue_values = sorted(float(item.get("queue_ms", 0) or 0) for item in results)
    run_values = sorted(float(item.get("run_ms", 0) or 0) for item in results)
    total_values = sorted(float(item.get("total_ms", 0) or 0) for item in results)
    latency_values = sorted(latencies)
    elapsed_sum_ms = sum(latency_values)
    return {
        "case": case_name,
        "requests": request_count,
        "concurrency": concurrency,
        "avg_latency_ms": round(statistics.fmean(latency_values), 2),
        "p95_latency_ms": round(_percentile(latency_values, 0.95), 2),
        "avg_queue_ms": round(statistics.fmean(queue_values), 2),
        "p95_queue_ms": round(_percentile(queue_values, 0.95), 2),
        "avg_run_ms": round(statistics.fmean(run_values), 2),
        "p95_run_ms": round(_percentile(run_values, 0.95), 2),
        "avg_total_ms": round(statistics.fmean(total_values), 2),
        "p95_total_ms": round(_percentile(total_values, 0.95), 2),
        "throughput_rps_estimate": round((request_count * 1000.0) / max(latency_values[-1], elapsed_sum_ms / max(request_count, 1)), 2),
    }


async def main_async(args) -> None:
    os.environ["JOB_BACKEND"] = "memory"
    os.environ["JOB_TTS_CONCURRENCY"] = str(args.tts_workers)
    os.environ["JOB_TTS_MAX_QUEUE"] = str(max(args.tts_requests, args.tts_workers))
    os.environ["JOB_ASR_CONCURRENCY"] = str(args.asr_workers)
    os.environ["JOB_ASR_MAX_QUEUE"] = str(max(args.asr_requests, args.asr_workers))
    reset_job_manager()

    mock_client = await _build_mock_client(args.tts_delay_ms, args.asr_delay_ms)
    tts_asr._httpx_client = mock_client
    try:
        audio_b64 = base64.b64encode(_wav_bytes(args.audio_duration_ms)).decode("utf-8")
        tts_stats = await _run_case("tts", args.tts_concurrency, args.tts_requests)
        asr_stats = await _run_case("asr", args.asr_concurrency, args.asr_requests, audio_b64=audio_b64)
    finally:
        await mock_client.aclose()
        tts_asr._httpx_client = None
        reset_job_manager()

    print(
        json.dumps(
            {
                "benchmark_date": time.strftime("%Y-%m-%d %H:%M:%S"),
                "assumptions": {
                    "upstream_tts_delay_ms": args.tts_delay_ms,
                    "upstream_asr_delay_ms": args.asr_delay_ms,
                    "job_backend": "memory",
                },
                "tts": tts_stats,
                "asr": asr_stats,
                "recommended_defaults": {
                    "JOB_TTS_CONCURRENCY": args.tts_workers,
                    "JOB_TTS_MAX_QUEUE": max(16, args.tts_workers * 4),
                    "JOB_ASR_CONCURRENCY": args.asr_workers,
                    "JOB_ASR_MAX_QUEUE": max(8, args.asr_workers * 4),
                    "TTS_ASR_MAX_CONNECTIONS": max(24, (args.tts_workers + args.asr_workers) * 4),
                    "TTS_ASR_MAX_KEEPALIVE_CONNECTIONS": max(12, (args.tts_workers + args.asr_workers) * 2),
                },
            },
            ensure_ascii=False,
            indent=2,
        )
    )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tts-delay-ms", type=int, default=120)
    parser.add_argument("--asr-delay-ms", type=int, default=280)
    parser.add_argument("--tts-workers", type=int, default=4)
    parser.add_argument("--asr-workers", type=int, default=2)
    parser.add_argument("--tts-concurrency", type=int, default=8)
    parser.add_argument("--asr-concurrency", type=int, default=4)
    parser.add_argument("--tts-requests", type=int, default=32)
    parser.add_argument("--asr-requests", type=int, default=16)
    parser.add_argument("--audio-duration-ms", type=int, default=320)
    return parser.parse_args()


if __name__ == "__main__":
    asyncio.run(main_async(parse_args()))
