#!/usr/bin/env python3
"""
bench_openai_api.py — Benchmark LLM via API OpenAI-compatible (vLLM / SGLang).

Usage:
    python scripts/bench_openai_api.py [--url http://localhost:8000] [--prompts prompts/bench_prompts.jsonl]

Measures:
    - Output tokens/s
    - Total tokens/s (input + output)
    - Time To First Token (TTFT) in ms
    - Time Per Output Token (TPOT) in ms
    - Total time, input tokens, output tokens
    - Success/failure per prompt
    - GPU stats if available (via nvidia-ml-py)

Saves results to reports/results.csv
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional

import httpx

# --- GPU monitoring (optional) ---
try:
    import pynvml
    HAS_PYNVML = True
except ImportError:
    HAS_PYNVML = False


# ============================================================
# Data structures
# ============================================================

@dataclass
class BenchResult:
    prompt_id: str
    category: str
    model: str = ""
    backend: str = ""
    config_name: str = ""
    max_model_len: int = 0
    gpu_memory_util: float = 0.0
    dtype: str = ""
    quantization: str = ""

    # Tokens
    input_tokens: int = 0
    output_tokens: int = 0

    # Timing
    ttft_ms: float = 0.0          # Time To First Token
    tpot_ms: float = 0.0          # Time Per Output Token (avg inter-token)
    total_time_s: float = 0.0
    output_tps: float = 0.0       # output tokens / total_time
    total_tps: float = 0.0        # (input + output) / total_time

    # Quality
    success: bool = True
    error: str = ""

    # GPU stats
    vram_peak_mb: float = 0.0
    gpu_util_pct: float = 0.0
    power_w: float = 0.0
    gpu_temp_c: float = 0.0

    # Prompt snippet (first 100 chars)
    prompt_snippet: str = ""


# ============================================================
# GPU Monitor
# ============================================================

class GpuMonitor:
    """Poll nvidia-smi or pynvml during benchmark."""

    def __init__(self, use_pynvml: bool = True):
        self.use_pynvml = use_pynvml and HAS_PYNVML
        self._peak_vram = 0.0
        self._gpu_utils: list[float] = []
        self._powers: list[float] = []
        self._temps: list[float] = []
        self._running = False

        if self.use_pynvml:
            pynvml.nvmlInit()
            self._handle = pynvml.nvmlDeviceGetHandleByIndex(0)

    def start(self):
        self._running = True

    def stop(self):
        self._running = False
        if self.use_pynvml:
            try:
                pynvml.nvmlShutdown()
            except Exception:
                pass

    def poll(self):
        """Call this periodically from another thread or between measurements."""
        if self.use_pynvml:
            try:
                info = pynvml.nvmlDeviceGetMemoryInfo(self._handle)
                vram_mb = info.used / (1024 * 1024)
                self._peak_vram = max(self._peak_vram, vram_mb)

                util = pynvml.nvmlDeviceGetUtilizationRates(self._handle)
                self._gpu_utils.append(util.gpu)
                self._temps.append(util.temperature if hasattr(util, 'temperature') else 0)

                # Power is slower to read, do it less often
                try:
                    power_mw = pynvml.nvmlDeviceGetPowerUsage(self._handle)
                    self._powers.append(power_mw / 1000.0)
                except Exception:
                    pass
            except Exception:
                pass

    @property
    def stats(self) -> dict:
        return {
            "vram_peak_mb": round(self._peak_vram, 1),
            "gpu_util_pct": round(sum(self._gpu_utils) / len(self._gpu_utils), 1) if self._gpu_utils else 0.0,
            "power_w": round(sum(self._powers) / len(self._powers), 1) if self._powers else 0.0,
            "gpu_temp_c": round(sum(self._temps) / len(self._temps), 1) if self._temps else 0.0,
        }


# ============================================================
# API Client
# ============================================================

class OpenAIChatClient:
    """Minimal async client for OpenAI-compatible chat completions."""

    def __init__(self, base_url: str, api_key: str = "not-needed", timeout: float = 300.0):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout
        self._client: Optional[httpx.Client] = None

    @property
    def client(self) -> httpx.Client:
        if self._client is None:
            self._client = httpx.Client(
                base_url=self.base_url,
                timeout=httpx.Timeout(self.timeout),
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
            )
        return self._client

    def check_health(self) -> tuple[bool, str]:
        """Check if the server is ready via /v1/models."""
        try:
            resp = self.client.get("/v1/models")
            if resp.status_code == 200:
                data = resp.json()
                models = [m.get("id", "") for m in data.get("data", [])]
                return True, f"OK — models: {models}"
            return False, f"HTTP {resp.status_code}: {resp.text[:200]}"
        except Exception as e:
            return False, str(e)

    def chat_completion(
        self,
        prompt: str,
        max_tokens: int = 256,
        temperature: float = 0.0,
        stream: bool = True,
    ) -> tuple[dict, float, list[float]]:
        """
        Send a chat completion request. Returns:
            (completion_dict, ttft_seconds, inter_token_times)
        If stream=True, measures TTFT and inter-token latencies.
        """
        payload = {
            "model": "default",  # vLLM/SGLang ignore this when only one model loaded
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": stream,
        }

        if not stream:
            t0 = time.perf_counter()
            resp = self.client.post("/v1/chat/completions", json=payload)
            elapsed = time.perf_counter() - t0
            if resp.status_code != 200:
                raise RuntimeError(f"API error {resp.status_code}: {resp.text[:500]}")
            data = resp.json()
            choice = data["choices"][0]
            usage = data.get("usage", {})
            return {
                "content": choice["message"]["content"],
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0),
                "finish_reason": choice.get("finish_reason", "unknown"),
            }, elapsed, []

        # --- Streaming mode ---
        t0 = time.perf_counter()
        ttft: Optional[float] = None
        inter_times: list[float] = []
        last_token_time: Optional[float] = None
        content_parts: list[str] = []
        usage_info: dict = {}
        finish_reason = "unknown"

        with self.client.stream("POST", "/v1/chat/completions", json=payload) as resp:
            if resp.status_code != 200:
                raise RuntimeError(f"API error {resp.status_code}: {resp.read().decode()[:500]}")

            for line in resp.iter_lines():
                if not line or not line.startswith("data: "):
                    continue
                data_str = line[6:]  # strip "data: "
                if data_str.strip() == "[DONE]":
                    break

                try:
                    chunk = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                now = time.perf_counter()

                # First token
                if ttft is None:
                    ttft = now - t0
                else:
                    if last_token_time is not None:
                        inter_times.append(now - last_token_time)
                last_token_time = now

                choices = chunk.get("choices", [])
                if choices:
                    delta = choices[0].get("delta", {})
                    content = delta.get("content", "")
                    if content:
                        content_parts.append(content)
                    if choices[0].get("finish_reason"):
                        finish_reason = choices[0]["finish_reason"]

                # Usage sometimes in final chunk
                if "usage" in chunk and chunk["usage"]:
                    usage_info = chunk["usage"]

        elapsed = time.perf_counter() - t0
        ttft_s = ttft if ttft is not None else elapsed

        return {
            "content": "".join(content_parts),
            "input_tokens": usage_info.get("prompt_tokens", 0),
            "output_tokens": usage_info.get("completion_tokens", len(content_parts)),
            "finish_reason": finish_reason,
        }, ttft_s, inter_times


# ============================================================
# Benchmark runner
# ============================================================

def load_prompts(prompts_path: str) -> list[dict]:
    """Load prompts from JSONL file."""
    prompts = []
    with open(prompts_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                prompts.append(json.loads(line))
    return prompts


def run_benchmark(
    client: OpenAIChatClient,
    prompts: list[dict],
    gpu_monitor: Optional[GpuMonitor],
    config_info: dict,
    repeat: int = 1,
    warmup: bool = True,
) -> list[BenchResult]:
    """Run benchmark on all prompts and return results."""
    results: list[BenchResult] = []

    # Warmup
    if warmup:
        print("  🔥 Warmup request...")
        try:
            client.chat_completion("Hello, respond with 'OK'.", max_tokens=10, stream=True)
            print("  ✅ Warmup OK")
        except Exception as e:
            print(f"  ⚠️  Warmup failed: {e}")

    total = len(prompts) * repeat

    for rep in range(repeat):
        for i, p in enumerate(prompts, 1):
            prompt_text = p["prompt"]
            prompt_id = f"{p['id']}_r{rep}" if repeat > 1 else p["id"]
            max_tokens = p.get("max_tokens", 256)

            result = BenchResult(
                prompt_id=prompt_id,
                category=p.get("category", "unknown"),
                prompt_snippet=prompt_text[:100].replace("\n", " "),
                **config_info,
            )

            print(f"  [{i}/{total}] {prompt_id} ({result.category})...", end=" ", flush=True)

            if gpu_monitor:
                gpu_monitor.poll()

            try:
                completion, ttft_s, inter_times = client.chat_completion(
                    prompt_text,
                    max_tokens=max_tokens,
                    temperature=0.0,
                    stream=True,
                )
                elapsed = time.perf_counter()  # total time will be updated below
                # We need the actual total time from the request
                # Re-run with timing — use the completion struct

                result.input_tokens = completion.get("input_tokens", 0)
                result.output_tokens = completion.get("output_tokens", 0)
                result.ttft_ms = ttft_s * 1000

                if inter_times:
                    result.tpot_ms = (sum(inter_times) / len(inter_times)) * 1000

                # For total time, we use ttft + sum of inter_times + small overhead
                total_stream_time = ttft_s + sum(inter_times) if inter_times else ttft_s
                result.total_time_s = total_stream_time

                if result.total_time_s > 0:
                    result.output_tps = result.output_tokens / result.total_time_s
                    total_tokens = result.input_tokens + result.output_tokens
                    result.total_tps = total_tokens / result.total_time_s

                if gpu_monitor:
                    gpu_monitor.poll()
                    gpu_stats = gpu_monitor.stats
                    result.vram_peak_mb = gpu_stats["vram_peak_mb"]
                    result.gpu_util_pct = gpu_stats["gpu_util_pct"]
                    result.power_w = gpu_stats["power_w"]
                    result.gpu_temp_c = gpu_stats["gpu_temp_c"]

                print(f"✅ {result.output_tps:.1f} tok/s (out={result.output_tokens}, TTFT={result.ttft_ms:.0f}ms)")

            except Exception as e:
                result.success = False
                result.error = str(e)[:500]
                print(f"❌ {str(e)[:100]}")

            results.append(result)

    return results


# ============================================================
# CSV persistence
# ============================================================

CSV_COLUMNS = [
    "prompt_id", "category", "model", "backend", "config_name",
    "max_model_len", "gpu_memory_util", "dtype", "quantization",
    "input_tokens", "output_tokens",
    "ttft_ms", "tpot_ms", "total_time_s", "output_tps", "total_tps",
    "success", "error",
    "vram_peak_mb", "gpu_util_pct", "power_w", "gpu_temp_c",
    "prompt_snippet",
]


def save_results(results: list[BenchResult], csv_path: str):
    """Append results to CSV file."""
    file_exists = os.path.exists(csv_path)
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)

    with open(csv_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS, extrasaction='ignore')
        if not file_exists or os.path.getsize(csv_path) == 0:
            writer.writeheader()
        for r in results:
            writer.writerow(asdict(r))

    print(f"\n📊 {len(results)} results saved to {csv_path}")


def print_summary(results: list[BenchResult]):
    """Print a summary table."""
    successful = [r for r in results if r.success]
    failed = [r for r in results if not r.success]

    print("\n" + "=" * 60)
    print(" BENCHMARK SUMMARY")
    print("=" * 60)
    print(f"  Total requests:  {len(results)}")
    print(f"  Successful:      {len(successful)}")
    print(f"  Failed:          {len(failed)}")
    print()

    if successful:
        output_tps_vals = [r.output_tps for r in successful]
        total_tps_vals = [r.total_tps for r in successful]
        ttft_vals = [r.ttft_ms for r in successful]
        tpot_vals = [r.tpot_ms for r in successful if r.tpot_ms > 0]

        print(f"  Output TPS:      {min(output_tps_vals):.1f} / {sum(output_tps_vals)/len(output_tps_vals):.1f} / {max(output_tps_vals):.1f}  (min/avg/max)")
        print(f"  Total TPS:       {min(total_tps_vals):.1f} / {sum(total_tps_vals)/len(total_tps_vals):.1f} / {max(total_tps_vals):.1f}  (min/avg/max)")
        print(f"  TTFT (ms):       {min(ttft_vals):.0f} / {sum(ttft_vals)/len(ttft_vals):.0f} / {max(ttft_vals):.0f}  (min/avg/max)")
        if tpot_vals:
            print(f"  TPOT (ms):       {min(tpot_vals):.1f} / {sum(tpot_vals)/len(tpot_vals):.1f} / {max(tpot_vals):.1f}  (min/avg/max)")

        # Per category
        print()
        for cat in sorted(set(r.category for r in successful)):
            cat_results = [r for r in successful if r.category == cat]
            avg_tps = sum(r.output_tps for r in cat_results) / len(cat_results)
            avg_ttft = sum(r.ttft_ms for r in cat_results) / len(cat_results)
            print(f"  {cat:10s}: avg {avg_tps:.1f} tok/s, TTFT {avg_ttft:.0f}ms  ({len(cat_results)} prompts)")

    if failed:
        print(f"\n  ❌ Failed requests:")
        for r in failed:
            print(f"     - {r.prompt_id}: {r.error[:120]}")

    # GPU stats
    gpu_results = [r for r in successful if r.vram_peak_mb > 0]
    if gpu_results:
        peak_vram = max(r.vram_peak_mb for r in gpu_results)
        avg_util = sum(r.gpu_util_pct for r in gpu_results) / len(gpu_results)
        avg_power = sum(r.power_w for r in gpu_results) / len(gpu_results)
        print(f"\n  GPU Stats:")
        print(f"  VRAM peak:       {peak_vram:.0f} MiB / 24576 MiB")
        print(f"  GPU util avg:    {avg_util:.1f}%")
        print(f"  Power avg:       {avg_power:.1f} W")

    print("=" * 60)


# ============================================================
# CLI
# ============================================================

def main():
    parser = argparse.ArgumentParser(description="Benchmark LLM via OpenAI-compatible API")
    parser.add_argument("--url", default="http://localhost:8000", help="API base URL")
    parser.add_argument("--prompts", default=None, help="Path to JSONL prompts file")
    parser.add_argument("--output", default=None, help="Output CSV path")
    parser.add_argument("--repeat", type=int, default=1, help="Repeat each prompt N times")
    parser.add_argument("--no-warmup", action="store_true", help="Skip warmup request")
    parser.add_argument("--no-gpu", action="store_true", help="Disable GPU monitoring")
    parser.add_argument("--timeout", type=float, default=300.0, help="Request timeout (seconds)")
    parser.add_argument("--model", default="", help="Model name (metadata only)")
    parser.add_argument("--backend", default="vllm", help="Backend name (vllm/sglang, metadata only)")
    parser.add_argument("--config", default="", help="Config name (metadata only)")
    parser.add_argument("--max-model-len", type=int, default=0, help="Max model len (metadata)")
    parser.add_argument("--gpu-mem-util", type=float, default=0.0, help="GPU memory util (metadata)")
    parser.add_argument("--dtype", default="", help="Dtype (metadata)")
    parser.add_argument("--quantization", default="", help="Quantization (metadata)")

    args = parser.parse_args()

    # Default paths relative to project
    project_dir = Path(__file__).resolve().parent.parent
    if args.prompts is None:
        args.prompts = str(project_dir / "prompts" / "bench_prompts.jsonl")
    if args.output is None:
        args.output = str(project_dir / "reports" / "results.csv")

    # Config info for metadata
    config_info = {
        "model": args.model,
        "backend": args.backend,
        "config_name": args.config,
        "max_model_len": args.max_model_len,
        "gpu_memory_util": args.gpu_mem_util,
        "dtype": args.dtype,
        "quantization": args.quantization,
    }

    print("=" * 60)
    print(" Qwen SpeedLab — OpenAI API Benchmark")
    print("=" * 60)
    print(f"  URL:      {args.url}")
    print(f"  Prompts:  {args.prompts}")
    print(f"  Output:   {args.output}")
    print(f"  Repeat:   {args.repeat}")
    print(f"  Timeout:  {args.timeout}s")
    print(f"  Config:   {config_info}")
    print("=" * 60)
    print()

    # Load prompts
    if not os.path.exists(args.prompts):
        print(f"❌ Prompts file not found: {args.prompts}")
        sys.exit(1)

    prompts = load_prompts(args.prompts)
    print(f"📄 Loaded {len(prompts)} prompts")

    # Create API client
    client = OpenAIChatClient(base_url=args.url, timeout=args.timeout)

    # Health check
    print("\n🔍 Checking server health...")
    ok, msg = client.check_health()
    if not ok:
        print(f"❌ Server not ready: {msg}")
        print("   Is the server running? Start with: bash scripts/serve_vllm.sh")
        sys.exit(1)
    print(f"✅ Server ready: {msg}")

    # GPU monitor
    gpu_monitor = None
    if not args.no_gpu:
        gpu_monitor = GpuMonitor(use_pynvml=True)
        gpu_monitor.start()
        print("📊 GPU monitoring enabled (pynvml)")

    # Run benchmark
    print(f"\n🚀 Running benchmark ({len(prompts) * args.repeat} requests)...\n")
    t_start = time.perf_counter()
    results = run_benchmark(
        client=client,
        prompts=prompts,
        gpu_monitor=gpu_monitor,
        config_info=config_info,
        repeat=args.repeat,
        warmup=not args.no_warmup,
    )
    total_elapsed = time.perf_counter() - t_start

    if gpu_monitor:
        gpu_monitor.stop()

    # Save
    save_results(results, args.output)

    # Summary
    print_summary(results)
    print(f"\n⏱️  Total benchmark time: {total_elapsed:.1f}s")

    # Close client
    client.client.close()


if __name__ == "__main__":
    main()