#!/usr/bin/env python3
"""
bench_vllm.py — Benchmark vLLM directly via its Python API (LLM class).

This bypasses the HTTP server and measures:
    - Prefill throughput (tokens/s)
    - Decode throughput (tokens/s)
    - End-to-end latency
    - Peak VRAM usage

Usage:
    python scripts/bench_vllm.py --model Qwen/Qwen3.6-27B-FP8 --max-model-len 8192

Note:
    Requires vLLM installed. This script loads the model in-process,
    so it will consume VRAM immediately. Make sure no other process
    is using the GPU.

    For HTTP-based benchmarking (against a running server), use
    bench_openai_api.py instead.
"""
from __future__ import annotations

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

# GPU monitoring
try:
    import pynvml
    HAS_PYNVML = True
except ImportError:
    HAS_PYNVML = False


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

@dataclass
class DirectBenchResult:
    prompt_id: str
    category: str
    model: str
    config_name: str
    max_model_len: int
    gpu_memory_util: float
    dtype: str
    quantization: str
    input_tokens: int
    output_tokens: int
    prefill_time_s: float     # Time to process input (first token)
    decode_time_s: float      # Time to generate output tokens
    total_time_s: float
    prefill_tps: float        # input tokens / prefill_time
    decode_tps: float         # output tokens / decode_time
    output_tps: float         # output tokens / total_time
    vram_peak_mb: float
    gpu_util_pct: float
    success: bool
    error: str


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

class GpuMonitor:
    def __init__(self):
        self.use_pynvml = HAS_PYNVML
        self._peak_vram = 0.0
        self._gpu_utils: list[float] = []

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

    def poll(self):
        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)
            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,
        }

    def close(self):
        if self.use_pynvml:
            try:
                pynvml.nvmlShutdown()
            except Exception:
                pass


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

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


def run_direct_benchmark(
    llm,                  # vLLM LLM instance
    tokenizer,            # Tokenizer
    prompts: list[dict],
    gpu_monitor: GpuMonitor,
    config_info: dict,
    warmup: bool = True,
) -> list[DirectBenchResult]:
    """Run direct vLLM benchmark (no HTTP)."""
    results: list[DirectBenchResult] = []

    # Warmup
    if warmup:
        print("  🔥 Warmup...")
        sampling_params = llm.__class__.__module__.split(".")[0]  # won't work, use hardcoded import
        try:
            from vllm import SamplingParams
            llm.generate(["Hello"], SamplingParams(max_tokens=10, temperature=0))
            print("  ✅ Warmup OK")
        except Exception as e:
            print(f"  ⚠️  Warmup failed: {e}")

    from vllm import SamplingParams

    for i, p in enumerate(prompts, 1):
        prompt_text = p["prompt"]
        prompt_id = p["id"]
        max_tokens = p.get("max_tokens", 256)

        result = DirectBenchResult(
            prompt_id=prompt_id,
            category=p.get("category", "unknown"),
            **config_info,
        )

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

        gpu_monitor.poll()

        try:
            # Tokenize to get input count
            input_ids = tokenizer.encode(prompt_text)
            result.input_tokens = len(input_ids)

            sampling_params = SamplingParams(
                max_tokens=max_tokens,
                temperature=0.0,
                ignore_eos=False,
            )

            t0 = time.perf_counter()
            outputs = llm.generate([prompt_text], sampling_params, use_tqdm=False)
            elapsed = time.perf_counter() - t0

            gpu_monitor.poll()

            if outputs:
                out = outputs[0]
                result.output_tokens = len(out.outputs[0].token_ids)

                # vLLM provides metrics on the output
                metrics = out.metrics
                if hasattr(metrics, 'first_token_time') and metrics.first_token_time is not None:
                    # first_token_time is relative to arrival time
                    result.prefill_time_s = metrics.first_token_time - getattr(metrics, 'arrival_time', 0)
                else:
                    # Estimate: use a fraction of total time based on input/output ratio
                    result.prefill_time_s = elapsed * 0.3  # rough heuristic

                result.decode_time_s = elapsed - result.prefill_time_s
                result.total_time_s = elapsed

                if result.prefill_time_s > 0:
                    result.prefill_tps = result.input_tokens / result.prefill_time_s
                if result.decode_time_s > 0:
                    result.decode_tps = result.output_tokens / result.decode_time_s
                if elapsed > 0:
                    result.output_tps = result.output_tokens / elapsed

                print(f"✅ {result.output_tps:.1f} tok/s (prefill={result.prefill_tps:.0f}, decode={result.decode_tps:.0f})")

            else:
                result.success = False
                result.error = "No output generated"
                print("❌ No output")

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

        results.append(result)

    return results


# ============================================================
# CSV
# ============================================================

DIRECT_CSV_COLUMNS = [
    "prompt_id", "category", "model", "config_name",
    "max_model_len", "gpu_memory_util", "dtype", "quantization",
    "input_tokens", "output_tokens",
    "prefill_time_s", "decode_time_s", "total_time_s",
    "prefill_tps", "decode_tps", "output_tps",
    "vram_peak_mb", "gpu_util_pct",
    "success", "error",
]


def save_results(results: list[DirectBenchResult], csv_path: str):
    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=DIRECT_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[DirectBenchResult]):
    successful = [r for r in results if r.success]
    if not successful:
        print("\n❌ No successful results")
        return

    decode_tps_vals = [r.decode_tps for r in successful if r.decode_tps > 0]
    output_tps_vals = [r.output_tps for r in successful]

    print("\n" + "=" * 60)
    print(" DIRECT VLLM BENCHMARK SUMMARY")
    print("=" * 60)
    print(f"  Successful: {len(successful)}/{len(results)}")
    if decode_tps_vals:
        print(f"  Decode TPS:  {min(decode_tps_vals):.1f} / {sum(decode_tps_vals)/len(decode_tps_vals):.1f} / {max(decode_tps_vals):.1f}  (min/avg/max)")
    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)")

    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)
        print(f"\n  VRAM peak:   {peak_vram:.0f} MiB / 24576 MiB")
    print("=" * 60)


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

def main():
    parser = argparse.ArgumentParser(description="Benchmark vLLM directly (in-process)")
    parser.add_argument("--model", default="Qwen/Qwen3.6-27B-FP8", help="Model name or path")
    parser.add_argument("--max-model-len", type=int, default=8192, help="Max model length")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.90, help="GPU memory utilization")
    parser.add_argument("--dtype", default="auto", help="Data type (auto, float16, bfloat16)")
    parser.add_argument("--quantization", default="", help="Quantization method")
    parser.add_argument("--prompts", default=None, help="JSONL prompts file")
    parser.add_argument("--output", default=None, help="Output CSV path")
    parser.add_argument("--config", default="direct", help="Config name for logging")
    parser.add_argument("--max-num-seqs", type=int, default=4, help="Max number of sequences")
    parser.add_argument("--trust-remote-code", action="store_true", default=True)

    args = parser.parse_args()

    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_direct.csv")

    print("=" * 60)
    print(" Qwen SpeedLab — Direct vLLM Benchmark")
    print("=" * 60)
    print(f"  Model:       {args.model}")
    print(f"  Max len:     {args.max_model_len}")
    print(f"  GPU mem:     {args.gpu_memory_utilization}")
    print(f"  Dtype:       {args.dtype}")
    print(f"  Prompts:     {args.prompts}")
    print("=" * 60)
    print()

    # Lazy import vLLM (heavy)
    print("📦 Importing vLLM...")
    try:
        from vllm import LLM, SamplingParams
    except ImportError:
        print("❌ vLLM not installed. Run: pip install vllm")
        sys.exit(1)

    # 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")

    # GPU monitor
    gpu_monitor = GpuMonitor()
    gpu_monitor.poll()

    # Load model
    print(f"\n🚀 Loading model: {args.model}")
    print(f"   max_model_len={args.max_model_len}, gpu_memory_utilization={args.gpu_memory_utilization}")
    t0 = time.perf_counter()

    try:
        llm = LLM(
            model=args.model,
            max_model_len=args.max_model_len,
            gpu_memory_utilization=args.gpu_memory_utilization,
            dtype=args.dtype,
            quantization=args.quantization if args.quantization else None,
            trust_remote_code=args.trust_remote_code,
            max_num_seqs=args.max_num_seqs,
            tensor_parallel_size=1,
            enable_prefix_caching=True,
        )
        load_time = time.perf_counter() - t0
        print(f"✅ Model loaded in {load_time:.1f}s")

        # Get tokenizer
        tokenizer = llm.get_tokenizer()

        gpu_monitor.poll()

    except Exception as e:
        print(f"❌ Failed to load model: {e}")
        gpu_monitor.close()
        sys.exit(1)

    # Run benchmark
    config_info = {
        "model": args.model,
        "config_name": args.config,
        "max_model_len": args.max_model_len,
        "gpu_memory_util": args.gpu_memory_utilization,
        "dtype": args.dtype,
        "quantization": args.quantization,
    }

    print(f"\n🚀 Running benchmark ({len(prompts)} prompts)...\n")
    results = run_direct_benchmark(
        llm=llm,
        tokenizer=tokenizer,
        prompts=prompts,
        gpu_monitor=gpu_monitor,
        config_info=config_info,
        warmup=True,
    )

    gpu_monitor.close()

    # Save
    save_results(results, args.output)

    # Summary
    print_summary(results)


if __name__ == "__main__":
    main()