#!/usr/bin/env python3
"""Custom benchmark client for vLLM / OpenAI-compatible API with multi-service load balancing and health-check failover.

Sends concurrent requests to multiple inference servers and records per-request
latency metrics (TTFT, TPOT, ITL, E2E) with fine-grained timing.

Designed for TP=2 on 8-GPU setups where 4 independent services run on
(0,1), (2,3), (4,5), (6,7).  Requests are distributed across services via
round-robin or random load balancing.  Unhealthy services are automatically
skipped with periodic retry.

Usage:
    # Single service (backward compatible)
    python3 bench_client.py \
        --host 127.0.0.1 --port 30005 \
        --model deepseek-v4-flash \
        --concurrency 64 --input-len 512 --output-len 256 --num-prompts 320 \
        --output-file results.jsonl

    # Multi-service (4 services on ports 30005-30008)
    python3 bench_client.py \
        --host 127.0.0.1 \
        --ports 30005,30006,30007,30008 \
        --model deepseek-v4-flash \
        --concurrency 64 --input-len 512 --output-len 256 --num-prompts 320 \
        --lb-strategy round_robin \
        --output-file results.jsonl

The output JSONL contains one object per request with raw timing data.
A final summary line ("completed" field) aggregates all requests.
"""

import argparse
import asyncio
import json
import random
import sys
import time
import urllib.request
from dataclasses import asdict, dataclass, field
from typing import Any

# Try aiohttp first; fall back to urllib for stdlib-only environments.
try:
    import aiohttp
    HAS_AIOHTTP = True
except ImportError:
    HAS_AIOHTTP = False


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass
class RequestResult:
    request_id: int
    input_tokens: int = 0
    output_tokens: int = 0
    first_token_time_ms: float = 0.0      # TTFT
    e2e_latency_ms: float = 0.0           # total wall time
    inter_token_latencies: list[float] = field(default_factory=list)
    success: bool = False
    error: str = ""
    target_port: int = 0                  # which service handled this request
    retry_count: int = 0                # how many times this request was retried

    @property
    def tpot_ms(self) -> float:
        """Mean TPOT = (e2e - ttft) / (output_tokens - 1) if >1 token."""
        if self.output_tokens <= 1:
            return 0.0
        return (self.e2e_latency_ms - self.first_token_time_ms) / (self.output_tokens - 1)

    @property
    def mean_itl_ms(self) -> float:
        if not self.inter_token_latencies:
            return 0.0
        return sum(self.inter_token_latencies) / len(self.inter_token_latencies)

    def to_dict(self) -> dict[str, Any]:
        d = asdict(self)
        d["tpot_ms"] = self.tpot_ms
        d["mean_itl_ms"] = self.mean_itl_ms
        return d


# ---------------------------------------------------------------------------
# Health checker
# ---------------------------------------------------------------------------

class HealthChecker:
    """Periodically check service health and maintain a healthy-port list."""

    def __init__(
        self,
        host: str,
        ports: list[int],
        check_interval: float = 5.0,
        max_failures: int = 2,
        recovery_interval: float = 10.0,
    ):
        self.host = host
        self.all_ports = ports
        self.check_interval = check_interval
        self.max_failures = max_failures
        self.recovery_interval = recovery_interval

        # port -> consecutive failure count
        self._failures: dict[int, int] = {p: 0 for p in ports}
        # port -> last healthy timestamp
        self._last_healthy: dict[int, float] = {p: time.monotonic() for p in ports}
        self._lock = asyncio.Lock()
        self._task: asyncio.Task | None = None

    async def _check_one(self, port: int, session: aiohttp.ClientSession | None) -> bool:
        url = f"http://{self.host}:{port}/health"
        try:
            if HAS_AIOHTTP and session is not None:
                async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
                    return resp.status == 200
            else:
                req = urllib.request.Request(url, method="GET")
                with urllib.request.urlopen(req, timeout=5) as resp:
                    return resp.status == 200
        except Exception:
            return False

    async def _check_all(self, session: aiohttp.ClientSession | None) -> None:
        checks = [self._check_one(p, session) for p in self.all_ports]
        results = await asyncio.gather(*checks, return_exceptions=True)
        now = time.monotonic()
        async with self._lock:
            for port, ok in zip(self.all_ports, results):
                if isinstance(ok, Exception):
                    ok = False
                if ok:
                    self._failures[port] = 0
                    self._last_healthy[port] = now
                else:
                    self._failures[port] += 1

    async def start(self, session: aiohttp.ClientSession | None) -> None:
        async def _loop() -> None:
            while True:
                await asyncio.sleep(self.check_interval)
                await self._check_all(session)
        self._task = asyncio.create_task(_loop())

    async def stop(self) -> None:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass

    @property
    async def healthy_ports(self) -> list[int]:
        now = time.monotonic()
        async with self._lock:
            healthy = []
            for p in self.all_ports:
                if self._failures[p] < self.max_failures:
                    healthy.append(p)
                elif now - self._last_healthy[p] > self.recovery_interval:
                    # Give it another chance after recovery_interval.
                    self._failures[p] = max(0, self._failures[p] - 1)
                    healthy.append(p)
            return healthy

    async def is_healthy(self, port: int) -> bool:
        ports = await self.healthy_ports
        return port in ports


# ---------------------------------------------------------------------------
# Load balancer with failover
# ---------------------------------------------------------------------------

class LoadBalancer:
    """Distribute requests across healthy backend ports with auto-failover."""

    def __init__(
        self,
        ports: list[int],
        health_checker: HealthChecker,
        strategy: str = "round_robin",
    ):
        self.all_ports = ports
        self.health_checker = health_checker
        self.strategy = strategy
        self._idx = 0
        self._lock = asyncio.Lock()

    def _next_round_robin(self, healthy: list[int]) -> int:
        if not healthy:
            return self.all_ports[self._idx % len(self.all_ports)]
        for _ in range(len(self.all_ports)):
            port = self.all_ports[self._idx % len(self.all_ports)]
            self._idx += 1
            if port in healthy:
                return port
        # Fallback: all unhealthy, pick first healthy anyway.
        return healthy[0] if healthy else self.all_ports[0]

    def _next_random(self, healthy: list[int]) -> int:
        if healthy:
            return random.choice(healthy)
        return random.choice(self.all_ports)

    async def next_port(self) -> int:
        healthy = await self.health_checker.healthy_ports
        async with self._lock:
            if self.strategy == "round_robin":
                return self._next_round_robin(healthy)
            elif self.strategy in ("random", "least_conn"):
                return self._next_random(healthy)
            else:
                return self._next_round_robin(healthy)

    @property
    def port_count(self) -> int:
        return len(self.all_ports)


# ---------------------------------------------------------------------------
# Token helpers
# ---------------------------------------------------------------------------

def make_prompt(token_len: int, vocab_size: int = 32000) -> str:
    """Generate a random-ish prompt of approximately `token_len` tokens."""
    words = ["the", "quick", "brown", "fox", "jumps", "over", "lazy",
             "dog", "hello", "world", "deep", "seek", "model", "test",
             "benchmark", "performance", "latency", "throughput", "token"]
    needed = max(token_len, 1)
    tokens: list[str] = []
    while len(tokens) < needed:
        tokens.extend(words)
    return " ".join(tokens[:needed])


# ---------------------------------------------------------------------------
# aiohttp-based async request (preferred)
# ---------------------------------------------------------------------------

async def send_request_aiohttp(
    session: aiohttp.ClientSession,
    host: str,
    port: int,
    model: str,
    request_id: int,
    prompt: str,
    max_tokens: int,
    temperature: float,
) -> RequestResult:
    url = f"http://{host}:{port}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": True,
    }
    headers = {"Content-Type": "application/json"}

    result = RequestResult(request_id=request_id, target_port=port)
    start_time = time.perf_counter()
    first_token_received = False
    last_token_time = 0.0

    try:
        async with session.post(url, json=payload, headers=headers) as resp:
            if resp.status != 200:
                text = await resp.text()
                result.error = f"HTTP {resp.status}: {text[:200]}"
                return result

            async for line in resp.content:
                line = line.decode("utf-8").strip()
                if not line or not line.startswith("data: "):
                    continue
                data_str = line[len("data: "):]
                if data_str == "[DONE]":
                    break
                try:
                    chunk = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                choices = chunk.get("choices", [])
                if not choices:
                    continue

                delta = choices[0].get("delta", {})
                content = delta.get("content", "")
                if content:
                    now = time.perf_counter()
                    if not first_token_received:
                        result.first_token_time_ms = (now - start_time) * 1000
                        first_token_received = True
                    else:
                        result.inter_token_latencies.append((now - last_token_time) * 1000)
                    last_token_time = now
                    result.output_tokens += 1

            result.e2e_latency_ms = (time.perf_counter() - start_time) * 1000
            result.input_tokens = len(prompt.split())  # approximate
            result.success = True

    except asyncio.TimeoutError:
        result.error = "timeout"
    except Exception as e:
        result.error = str(e)[:200]

    return result


# ---------------------------------------------------------------------------
# urllib-based synchronous request (fallback, no aiohttp)
# ---------------------------------------------------------------------------

def send_request_urllib(
    host: str,
    port: int,
    model: str,
    request_id: int,
    prompt: str,
    max_tokens: int,
    temperature: float,
) -> RequestResult:
    """Blocking fallback using stdlib urllib.  Used when aiohttp is absent."""
    url = f"http://{host}:{port}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False,
    }
    data = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json", "Content-Length": str(len(data))}

    result = RequestResult(request_id=request_id, target_port=port)
    start_time = time.perf_counter()

    try:
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=600) as resp:
            body = json.loads(resp.read().decode("utf-8"))

        choices = body.get("choices", [])
        if not choices:
            result.error = "no choices in response"
            return result

        content = choices[0].get("message", {}).get("content", "")
        result.output_tokens = len(content.split()) if content else 0
        result.e2e_latency_ms = (time.perf_counter() - start_time) * 1000
        result.first_token_time_ms = result.e2e_latency_ms
        result.input_tokens = len(prompt.split())
        result.success = True

    except Exception as e:
        result.error = str(e)[:200]

    return result


# ---------------------------------------------------------------------------
# Unified send_request wrapper with retry & failover
# ---------------------------------------------------------------------------

async def send_request_with_retry(
    session: aiohttp.ClientSession | None,
    host: str,
    lb: LoadBalancer,
    model: str,
    request_id: int,
    prompt: str,
    max_tokens: int,
    temperature: float,
    max_retries: int = 2,
) -> RequestResult:
    """Send a request, retrying on failure with a different port each time."""
    result = RequestResult(request_id=request_id)
    for attempt in range(max_retries + 1):
        port = await lb.next_port()
        if HAS_AIOHTTP and session is not None:
            result = await send_request_aiohttp(
                session, host, port, model, request_id, prompt, max_tokens, temperature
            )
        else:
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None, send_request_urllib,
                host, port, model, request_id, prompt, max_tokens, temperature
            )
        result.retry_count = attempt
        if result.success:
            return result
        # If failed, try another port on next iteration (unless last attempt).
        if attempt < max_retries:
            await asyncio.sleep(0.5 * (attempt + 1))  # exponential backoff-ish
    return result


# ---------------------------------------------------------------------------
# Benchmark orchestration
# ---------------------------------------------------------------------------

async def run_benchmark(
    host: str,
    ports: list[int],
    lb: LoadBalancer,
    health_checker: HealthChecker,
    model: str,
    concurrency: int,
    input_len: int,
    output_len: int,
    num_prompts: int,
    temperature: float,
    output_file: str,
) -> dict[str, Any]:
    """Run the benchmark and write a JSONL file."""
    prompts = [make_prompt(input_len) for _ in range(num_prompts)]
    results: list[RequestResult] = []
    semaphore = asyncio.Semaphore(concurrency)

    start = time.perf_counter()

    if HAS_AIOHTTP:
        connector = aiohttp.TCPConnector(limit=concurrency * 2)
        timeout = aiohttp.ClientTimeout(total=600, sock_read=300)
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            await health_checker.start(session)

            async def bounded_send(idx: int, prompt: str) -> RequestResult:
                async with semaphore:
                    return await send_request_with_retry(
                        session, host, lb, model, idx, prompt, output_len, temperature
                    )

            tasks = [bounded_send(i, p) for i, p in enumerate(prompts)]
            results = await asyncio.gather(*tasks)
            await health_checker.stop()
    else:
        await health_checker.start(None)
        from concurrent.futures import ThreadPoolExecutor

        async def bounded_send(idx: int, prompt: str) -> RequestResult:
            async with semaphore:
                return await send_request_with_retry(
                    None, host, lb, model, idx, prompt, output_len, temperature
                )

        tasks = [bounded_send(i, p) for i, p in enumerate(prompts)]
        results = await asyncio.gather(*tasks)
        await health_checker.stop()

    duration = time.perf_counter() - start

    # Write JSONL
    with open(output_file, "w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r.to_dict(), ensure_ascii=False) + "\n")

    # Compute summary
    successful = [r for r in results if r.success]
    failed = len(results) - len(successful)
    retried = [r for r in results if r.retry_count > 0]

    def percentile(values: list[float], p: float) -> float:
        if not values:
            return 0.0
        s = sorted(values)
        k = (len(s) - 1) * p / 100.0
        f = int(k)
        c = min(f + 1, len(s) - 1)
        return s[f] + (k - f) * (s[c] - s[f])

    ttfts = [r.first_token_time_ms for r in successful]
    tpots = [r.tpot_ms for r in successful if r.output_tokens > 1]
    e2es = [r.e2e_latency_ms for r in successful]
    itls = [r.mean_itl_ms for r in successful if r.inter_token_latencies]

    total_input_tokens = sum(r.input_tokens for r in successful)
    total_output_tokens = sum(r.output_tokens for r in successful)

    # Per-port breakdown
    per_port_stats: dict[int, dict[str, Any]] = {}
    for port in ports:
        port_results = [r for r in successful if r.target_port == port]
        port_failed = [r for r in results if r.target_port == port and not r.success]
        if port_results or port_failed:
            port_ttfts = [r.first_token_time_ms for r in port_results]
            port_e2es = [r.e2e_latency_ms for r in port_results]
            per_port_stats[port] = {
                "count": len(port_results),
                "failed": len(port_failed),
                "mean_ttft_ms": sum(port_ttfts) / len(port_ttfts) if port_ttfts else 0.0,
                "p95_ttft_ms": percentile(port_ttfts, 95) if port_ttfts else 0.0,
                "mean_e2e_ms": sum(port_e2es) / len(port_e2es) if port_e2es else 0.0,
            }

    # Health-check stats
    healthy_at_end = await health_checker.healthy_ports

    summary = {
        "completed": len(successful),
        "failed": failed,
        "retried": len(retried),
        "duration": duration,
        "request_throughput": len(successful) / duration if duration > 0 else 0.0,
        "input_throughput": total_input_tokens / duration if duration > 0 else 0.0,
        "output_throughput": total_output_tokens / duration if duration > 0 else 0.0,
        "total_throughput": (total_input_tokens + total_output_tokens) / duration if duration > 0 else 0.0,
        "total_input_tokens": total_input_tokens,
        "total_output_tokens": total_output_tokens,
        "mean_ttft_ms": sum(ttfts) / len(ttfts) if ttfts else 0.0,
        "median_ttft_ms": percentile(ttfts, 50),
        "p90_ttft_ms": percentile(ttfts, 90),
        "p95_ttft_ms": percentile(ttfts, 95),
        "p99_ttft_ms": percentile(ttfts, 99),
        "mean_tpot_ms": sum(tpots) / len(tpots) if tpots else 0.0,
        "median_tpot_ms": percentile(tpots, 50),
        "p90_tpot_ms": percentile(tpots, 90),
        "p95_tpot_ms": percentile(tpots, 95),
        "p99_tpot_ms": percentile(tpots, 99),
        "mean_e2e_latency_ms": sum(e2es) / len(e2es) if e2es else 0.0,
        "median_e2e_latency_ms": percentile(e2es, 50),
        "p90_e2e_latency_ms": percentile(e2es, 90),
        "p95_e2e_latency_ms": percentile(e2es, 95),
        "p99_e2e_latency_ms": percentile(e2es, 99),
        "mean_itl_ms": sum(itls) / len(itls) if itls else 0.0,
        "median_itl_ms": percentile(itls, 50),
        "p90_itl_ms": percentile(itls, 90),
        "p95_itl_ms": percentile(itls, 95),
        "p99_itl_ms": percentile(itls, 99),
        "per_port_stats": per_port_stats,
        "healthy_ports_at_end": healthy_at_end,
        "lb_strategy": lb.strategy,
        "num_services": lb.port_count,
    }

    # Append summary as the last line
    with open(output_file, "a", encoding="utf-8") as f:
        f.write(json.dumps(summary, ensure_ascii=False) + "\n")

    return summary


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description="Custom vLLM benchmark client with multi-service LB and health-check failover")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=30005,
                        help="Single service port (backward compatible)")
    parser.add_argument("--ports", default="",
                        help="Comma-separated list of ports for multi-service, e.g. 30005,30006,30007,30008")
    parser.add_argument("--model", default="deepseek-v4-flash")
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--input-len", type=int, default=512)
    parser.add_argument("--output-len", type=int, default=256)
    parser.add_argument("--num-prompts", type=int, default=10)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--lb-strategy", default="round_robin",
                        choices=["round_robin", "random", "least_conn"],
                        help="Load balancing strategy across services")
    parser.add_argument("--health-check-interval", type=float, default=5.0,
                        help="Seconds between health checks (default: 5)")
    parser.add_argument("--health-max-failures", type=int, default=2,
                        help="Consecutive failures before marking a port unhealthy (default: 2)")
    parser.add_argument("--health-recovery-interval", type=float, default=10.0,
                        help="Seconds before retrying an unhealthy port (default: 10)")
    parser.add_argument("--request-max-retries", type=int, default=2,
                        help="Max retries per request on failure (default: 2)")
    parser.add_argument("--output-file", required=True)
    args = parser.parse_args()

    # Determine ports: --ports takes precedence, else fall back to --port.
    if args.ports:
        ports = [int(p.strip()) for p in args.ports.split(",")]
    else:
        ports = [args.port]

    health_checker = HealthChecker(
        host=args.host,
        ports=ports,
        check_interval=args.health_check_interval,
        max_failures=args.health_max_failures,
        recovery_interval=args.health_recovery_interval,
    )
    lb = LoadBalancer(ports, health_checker=health_checker, strategy=args.lb_strategy)

    summary = asyncio.run(run_benchmark(
        host=args.host,
        ports=ports,
        lb=lb,
        health_checker=health_checker,
        model=args.model,
        concurrency=args.concurrency,
        input_len=args.input_len,
        output_len=args.output_len,
        num_prompts=args.num_prompts,
        temperature=args.temperature,
        output_file=args.output_file,
    ))

    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()