"""Benchmark the inference harness for load testing.

This script measures latency, throughput, and resource usage under various
load conditions to validate SLA targets.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import statistics
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path

import httpx


@dataclass
class BenchmarkResult:
    """Results from a benchmark run.

    Attributes:
        total_requests: Total number of requests sent.
        successful_requests: Number of successful requests.
        failed_requests: Number of failed requests.
        duration_seconds: Total benchmark duration.
        latencies_ms: List of individual request latencies.
        errors: List of error messages.
    """

    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    duration_seconds: float = 0.0
    latencies_ms: list[float] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def qps(self) -> float:
        """Queries per second."""
        if self.duration_seconds == 0:
            return 0.0
        return self.successful_requests / self.duration_seconds

    @property
    def p50_ms(self) -> float:
        """50th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.quantiles(self.latencies_ms, n=100)[49]

    @property
    def p95_ms(self) -> float:
        """95th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.quantiles(self.latencies_ms, n=100)[94]

    @property
    def p99_ms(self) -> float:
        """99th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.quantiles(self.latencies_ms, n=100)[98]

    @property
    def mean_ms(self) -> float:
        """Mean latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.mean(self.latencies_ms)

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "total_requests": self.total_requests,
            "successful_requests": self.successful_requests,
            "failed_requests": self.failed_requests,
            "duration_seconds": self.duration_seconds,
            "qps": self.qps,
            "latency_ms": {
                "mean": self.mean_ms,
                "p50": self.p50_ms,
                "p95": self.p95_ms,
                "p99": self.p99_ms,
                "min": min(self.latencies_ms) if self.latencies_ms else 0.0,
                "max": max(self.latencies_ms) if self.latencies_ms else 0.0,
            },
            "error_count": len(self.errors),
        }


async def send_request(
    client: httpx.AsyncClient,
    url: str,
    query: str,
) -> tuple[float, str | None]:
    """Send a single request and measure latency.

    Args:
        client: HTTP client.
        url: Target URL.
        query: Query text.

    Returns:
        Tuple of (latency_ms, error_message or None).
    """
    payload = {
        "messages": [{"role": "user", "content": query}],
        "model": "router",
    }

    start = time.perf_counter()
    try:
        response = await client.post(url, json=payload)
        elapsed_ms = (time.perf_counter() - start) * 1000

        if response.status_code != 200:
            return elapsed_ms, f"HTTP {response.status_code}: {response.text[:100]}"

        return elapsed_ms, None
    except Exception as e:
        elapsed_ms = (time.perf_counter() - start) * 1000
        return elapsed_ms, str(e)


async def run_benchmark(
    base_url: str,
    target_qps: int,
    duration_seconds: int,
    warmup_seconds: int = 5,
) -> BenchmarkResult:
    """Run the benchmark at a target QPS.

    Args:
        base_url: Base URL of the server.
        target_qps: Target queries per second.
        duration_seconds: Duration of the benchmark.
        warmup_seconds: Warmup period before measuring.

    Returns:
        BenchmarkResult with metrics.
    """
    url = f"{base_url}/v1/chat/completions"
    result = BenchmarkResult()

    # Generate sample queries
    queries = [
        "What is the capital of France?",
        "Explain quantum computing in simple terms.",
        "Write a Python function to sort a list.",
        "What are the benefits of exercise?",
        "How does photosynthesis work?",
        "Summarize the plot of Romeo and Juliet.",
        "What is machine learning?",
        "Explain the theory of relativity.",
        "How do I make chocolate chip cookies?",
        "What causes climate change?",
    ]

    interval = 1.0 / target_qps if target_qps > 0 else 0.01
    query_idx = 0

    async with httpx.AsyncClient(timeout=30.0) as client:
        # Warmup phase
        print(f"Warming up for {warmup_seconds} seconds...")
        warmup_end = time.perf_counter() + warmup_seconds
        while time.perf_counter() < warmup_end:
            query = queries[query_idx % len(queries)]
            await send_request(client, url, query)
            query_idx += 1
            await asyncio.sleep(interval)

        # Benchmark phase
        print(f"Running benchmark for {duration_seconds} seconds at {target_qps} QPS...")
        benchmark_start = time.perf_counter()
        benchmark_end = benchmark_start + duration_seconds

        tasks: list[asyncio.Task] = []

        async def rate_limited_sender():
            nonlocal query_idx
            while time.perf_counter() < benchmark_end:
                query = queries[query_idx % len(queries)]
                query_idx += 1

                # Create task for this request
                task = asyncio.create_task(send_request(client, url, query))
                tasks.append(task)

                # Wait for next interval
                await asyncio.sleep(interval)

        # Run sender
        await rate_limited_sender()

        # Wait for remaining tasks
        print("Waiting for remaining requests to complete...")
        results = await asyncio.gather(*tasks, return_exceptions=True)

        result.duration_seconds = time.perf_counter() - benchmark_start
        result.total_requests = len(results)

        for r in results:
            if isinstance(r, Exception):
                result.failed_requests += 1
                result.errors.append(str(r))
            else:
                latency_ms, error = r
                if error:
                    result.failed_requests += 1
                    result.errors.append(error)
                else:
                    result.successful_requests += 1
                    result.latencies_ms.append(latency_ms)

    return result


def print_results(result: BenchmarkResult) -> None:
    """Print benchmark results in a formatted way."""
    print("\n" + "=" * 60)
    print("BENCHMARK RESULTS")
    print("=" * 60)
    print(f"Total requests:      {result.total_requests}")
    print(f"Successful:          {result.successful_requests}")
    print(f"Failed:              {result.failed_requests}")
    print(f"Duration:            {result.duration_seconds:.2f}s")
    print(f"Achieved QPS:        {result.qps:.2f}")
    print()
    print("Latency (ms):")
    print(f"  Mean:              {result.mean_ms:.2f}")
    print(f"  P50:               {result.p50_ms:.2f}")
    print(f"  P95:               {result.p95_ms:.2f}")
    print(f"  P99:               {result.p99_ms:.2f}")
    if result.latencies_ms:
        print(f"  Min:               {min(result.latencies_ms):.2f}")
        print(f"  Max:               {max(result.latencies_ms):.2f}")
    print("=" * 60)

    # Check SLA targets
    print("\nSLA CHECK:")
    p99_target = 100  # ms
    if result.p99_ms <= p99_target:
        print(f"  [PASS] P99 latency ({result.p99_ms:.2f}ms) <= {p99_target}ms")
    else:
        print(f"  [FAIL] P99 latency ({result.p99_ms:.2f}ms) > {p99_target}ms")

    if result.failed_requests == 0:
        print("  [PASS] No failed requests")
    else:
        print(f"  [WARN] {result.failed_requests} failed requests")


def main() -> int:
    """CLI entry point."""
    parser = argparse.ArgumentParser(description="Benchmark the inference harness")
    parser.add_argument(
        "--url",
        default="http://localhost:8082",
        help="Base URL of the server (default: http://localhost:8082)",
    )
    parser.add_argument(
        "--qps",
        type=int,
        default=10,
        help="Target queries per second (default: 10)",
    )
    parser.add_argument(
        "--duration",
        type=int,
        default=30,
        help="Benchmark duration in seconds (default: 30)",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=5,
        help="Warmup duration in seconds (default: 5)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Output file for JSON results",
    )

    args = parser.parse_args()

    try:
        result = asyncio.run(
            run_benchmark(
                base_url=args.url,
                target_qps=args.qps,
                duration_seconds=args.duration,
                warmup_seconds=args.warmup,
            )
        )

        print_results(result)

        if args.output:
            with open(args.output, "w") as f:
                json.dump(result.to_dict(), f, indent=2)
            print(f"\nResults saved to {args.output}")

        return 0 if result.failed_requests == 0 else 1

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
