#!/usr/bin/env python3
"""
sweep_advanced.py — Advanced vLLM optimization sweep for RTX 3090.

Tests each optimization dimension independently to measure its impact:
  1. Attention backend: flash-attn vs flashinfer
  2. KV cache dtype: auto vs fp8
  3. Chunked prefill: on vs off
  4. Block size: 8 vs 16 vs 32
  5. GPU memory utilization: 0.85 vs 0.88 vs 0.90 vs 0.92
  6. max_num_seqs vs max_batched_tokens grid
  7. AWQ vs GPTQ vs FP8 model comparison
  8. Speculative decoding (optional, if draft model available)

Usage:
    # Baseline: test each optimization in isolation
    python scripts/sweep_advanced.py --mode baseline

    # Model shootout: compare AWQ vs GPTQ vs FP8
    python scripts/sweep_advanced.py --mode shootout

    # Full grid: exhaustively test combinations
    python scripts/sweep_advanced.py --mode grid

    # Single test with full instrumentation
    python scripts/sweep_advanced.py --mode single --config awq-flashinfer-fp8kv
"""
from __future__ import annotations

import argparse
import json
import os
import signal
import socket
import subprocess
import sys
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
from typing import Optional

import httpx

PROJECT_DIR = Path(__file__).resolve().parent.parent
REPORTS_DIR = PROJECT_DIR / "reports"
PROMPTS_FILE = PROJECT_DIR / "prompts" / "bench_prompts.jsonl"
RESULTS_CSV = REPORTS_DIR / "results_advanced.csv"
SUMMARY_MD = REPORTS_DIR / "summary_advanced.md"

# ---- GPU warmup / profiler ----
try:
    import pynvml
    HAS_PYNVML = True
except ImportError:
    HAS_PYNVML = False

# ---- vLLM availability ----
VLLM_AVAILABLE = False
try:
    import vllm
    VLLM_AVAILABLE = True
except ImportError:
    pass


@dataclass
class OptimizedConfig:
    """A single vLLM config to benchmark."""
    name: str
    model_path: str
    quant_method: str          # awq, fp8, gptq, none
    max_model_len: int = 8192
    gpu_memory_utilization: float = 0.90
    dtype: str = "auto"
    max_num_seqs: int = 8
    max_num_batched_tokens: int = 32768
    attention_backend: str = "flashinfer"   # flashinfer, flash-attn, auto
    kv_cache_dtype: str = "auto"            # auto, fp8
    enable_chunked_prefill: bool = False
    block_size: int = 16
    swap_space: int = 4
    enable_prefix_caching: bool = True
    max_seq_len_to_capture: int = 8192
    # Speculative decoding
    speculative_model: str = ""
    num_speculative_tokens: int = 5

    @property
    def label(self) -> str:
        parts = [self.quant_method, f"len{self.max_model_len}"]
        if self.attention_backend != "auto":
            parts.append(self.attention_backend.replace("-", ""))
        if self.kv_cache_dtype == "fp8":
            parts.append("fp8kv")
        if self.enable_chunked_prefill:
            parts.append("chunk")
        parts.append(f"bs{self.block_size}")
        parts.append(f"s{self.max_num_seqs}")
        return "-".join(parts)

    def to_server_flags(self, port: int = 8000) -> list[str]:
        """Convert to vLLM CLI arguments."""
        flags = [
            "--model", self.model_path,
            "--host", "0.0.0.0",
            "--port", str(port),
            "--max-model-len", str(self.max_model_len),
            "--gpu-memory-utilization", str(self.gpu_memory_utilization),
            "--dtype", self.dtype,
            "--max-num-seqs", str(self.max_num_seqs),
            "--max-num-batched-tokens", str(self.max_num_batched_tokens),
            "--tensor-parallel-size", "1",
            "--trust-remote-code",
            "--block-size", str(self.block_size),
            "--swap-space", str(self.swap_space),
            "--max-seq-len-to-capture", str(self.max_seq_len_to_capture),
        ]
        if self.enable_prefix_caching:
            flags.append("--enable-prefix-caching")
        if self.enable_chunked_prefill:
            flags.append("--enable-chunked-prefill")
        if self.attention_backend not in ("auto", ""):
            flags.extend(["--attention-backend", self.attention_backend])
        if self.kv_cache_dtype not in ("auto", ""):
            flags.extend(["--kv-cache-dtype", self.kv_cache_dtype])
        if self.speculative_model:
            flags.extend(["--speculative-model", self.speculative_model])
            flags.extend(["--num-speculative-tokens", str(self.num_speculative_tokens)])
        return flags


# ============================================================
# Config presets for each sweep mode
# ============================================================

# Models — best to worst for RTX 3090
MODELS = {
    "awq":  "QuantTrio/Qwen3.6-27B-AWQ",
    "gptq": "groxaxo/Qwen3.6-27B-GPTQ-Pro-4bit",
    "fp8":  "Qwen/Qwen3.6-27B-FP8",
    "base": "Qwen/Qwen3.6-27B",
}


def baseline_configs() -> list[OptimizedConfig]:
    """
    A/B test each optimization dimension independently.
    All configs use AWQ as the base model (best for Ampere).
    """
    BASE = dict(
        model_path=MODELS["awq"], quant_method="awq",
        max_model_len=8192, gpu_memory_utilization=0.90,
        max_num_seqs=8, max_num_batched_tokens=32768,
    )
    configs = []

    # ---- 1. Attention backend shootout ----
    for attn in ["flashinfer", "flash-attn"]:
        c = OptimizedConfig(name=f"attn-{attn}", attention_backend=attn, **BASE)
        configs.append(c)

    # ---- 2. FP8 KV cache (VRAM saving) ----
    for kv in ["auto", "fp8"]:
        c = OptimizedConfig(name=f"kv-{kv}", kv_cache_dtype=kv, **BASE)
        configs.append(c)

    # ---- 3. Chunked prefill ----
    for cp in [False, True]:
        tag = "chunk-on" if cp else "chunk-off"
        c = OptimizedConfig(name=tag, enable_chunked_prefill=cp, **BASE)
        configs.append(c)

    # ---- 4. Block size sweep ----
    for bs in [8, 16, 32]:
        c = OptimizedConfig(name=f"blksz-{bs}", block_size=bs, **BASE)
        configs.append(c)

    # ---- 5. GPU memory utilization sweep ----
    for mem in [0.85, 0.88, 0.90, 0.92]:
        c = OptimizedConfig(name=f"gpumem-{int(mem*100)}", gpu_memory_utilization=mem, **BASE)
        configs.append(c)

    # ---- 6. Batch size sweep ----
    for seqs, batched in [(4, 16384), (8, 32768), (12, 49152)]:
        c = OptimizedConfig(
            name=f"batch-s{seqs}-b{batched}",
            max_num_seqs=seqs, max_num_batched_tokens=batched, **BASE,
        )
        configs.append(c)

    # ---- 7. Combined best guess ----
    c = OptimizedConfig(
        name="combined-best",
        attention_backend="flashinfer",
        kv_cache_dtype="fp8",
        enable_chunked_prefill=True,
        block_size=16,
        max_num_seqs=8,
        max_num_batched_tokens=32768,
        **BASE,
    )
    configs.append(c)

    return configs


def model_shootout_configs() -> list[OptimizedConfig]:
    """Compare AWQ vs GPTQ vs FP8 on equal footing."""
    configs = []

    COMMON = dict(
        max_model_len=8192, gpu_memory_utilization=0.90,
        max_num_seqs=8, max_num_batched_tokens=32768,
        attention_backend="flashinfer", kv_cache_dtype="fp8",
        enable_chunked_prefill=True,
    )

    # AWQ — expected best
    configs.append(OptimizedConfig(
        name="shootout-awq", model_path=MODELS["awq"], quant_method="awq", **COMMON,
    ))

    # GPTQ Marlin — potentially faster on Ampere
    configs.append(OptimizedConfig(
        name="shootout-gptq", model_path=MODELS["gptq"], quant_method="gptq", **COMMON,
    ))

    # FP8 — higher precision, tighter on VRAM
    configs.append(OptimizedConfig(
        name="shootout-fp8", model_path=MODELS["fp8"], quant_method="fp8",
        max_num_seqs=6, max_num_batched_tokens=16384, **COMMON,
    ))

    # FP8 + FP8 KV cache
    configs.append(OptimizedConfig(
        name="shootout-fp8-fp8kv", model_path=MODELS["fp8"], quant_method="fp8",
        max_num_seqs=8, max_num_batched_tokens=32768, **COMMON,
    ))

    return configs


def grid_configs() -> list[OptimizedConfig]:
    """Full grid search over the most impactful dimensions. ~20 configs."""
    configs = []

    for model_key, model_path in [("awq", MODELS["awq"]), ("gptq", MODELS["gptq"])]:
        for attn in ["flashinfer", "flash-attn"]:
            for kv in ["fp8", "auto"]:
                for bs in [16, 32]:
                    for seqs in [8, 12]:
                        c = OptimizedConfig(
                            name=f"grid-{model_key}-{attn}-{kv}kv-bs{bs}-s{seqs}",
                            model_path=model_path,
                            quant_method=model_key,
                            attention_backend=attn,
                            kv_cache_dtype=kv,
                            block_size=bs,
                            max_num_seqs=seqs,
                            max_num_batched_tokens=seqs * 4096,
                            max_model_len=8192,
                            gpu_memory_utilization=0.90,
                            enable_chunked_prefill=True,
                        )
                        configs.append(c)

    return configs


# ============================================================
# Server lifecycle
# ============================================================

class VllmServer:
    """Manage a vLLM server process."""

    def __init__(self, config: OptimizedConfig, port: int = 8000):
        self.config = config
        self.port = port
        self.process: Optional[subprocess.Popen] = None
        self.logfile = Path(f"/tmp/vllm_adv_{config.name}.log")

    def start(self):
        cmd = [sys.executable, "-m", "vllm.entrypoints.openai.api_server"]
        cmd.extend(self.config.to_server_flags(port=self.port))

        print(f"  {' '.join(cmd)}")

        self.logfile.parent.mkdir(parents=True, exist_ok=True)
        fout = open(str(self.logfile), "w")
        self.process = subprocess.Popen(
            cmd, stdout=fout, stderr=subprocess.STDOUT,
            start_new_session=True,
            env={**os.environ, "VLLM_ATTENTION_BACKEND": self.config.attention_backend},
        )
        return True

    def wait_ready(self, timeout: int = 600) -> tuple[bool, str]:
        client = httpx.Client(base_url=f"http://localhost:{self.port}", timeout=httpx.Timeout(8.0))
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.process and self.process.poll() is not None:
                tail = self._log_tail(60)
                return False, f"CRASH rc={self.process.returncode}\n{tail}"
            try:
                resp = client.get("/v1/models")
                if resp.status_code == 200:
                    client.close()
                    return True, "OK"
            except Exception:
                pass
            time.sleep(5)
        client.close()
        return False, f"TIMEOUT ({timeout}s)\n{self._log_tail(30)}"

    def stop(self):
        if not self.process:
            return
        print("  ⏹️  Stopping server...")
        try:
            os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
            self.process.wait(timeout=30)
        except subprocess.TimeoutExpired:
            os.killpg(os.getpgid(self.process.pid), signal.SIGKILL)
            self.process.wait(timeout=10)
        except Exception:
            pass
        self.process = None
        time.sleep(5)

    def _log_tail(self, lines: int = 40) -> str:
        if not self.logfile.exists():
            return "(no log)"
        try:
            with open(self.logfile, "r") as f:
                return "".join(f.readlines()[-lines:])
        except Exception:
            return "(error)"


# ============================================================
# Benchmark integration
# ============================================================

@dataclass
class SweepResult:
    config: OptimizedConfig
    success: bool
    error: str = ""
    avg_output_tps: float = 0.0
    avg_total_tps: float = 0.0
    avg_ttft_ms: float = 0.0
    avg_tpot_ms: float = 0.0
    peak_vram_mb: float = 0.0
    avg_gpu_util_pct: float = 0.0
    avg_power_w: float = 0.0
    ok: int = 0
    fail: int = 0
    duration_s: float = 0.0


def parse_csv_for_config(csv_path: str, config_name: str) -> Optional[dict]:
    """Aggregate results from CSV for a config."""
    if not os.path.exists(csv_path):
        return None
    import csv
    try:
        with open(csv_path, "r") as f:
            rows = [r for r in csv.DictReader(f) if r.get("config_name") == config_name]
    except Exception:
        return None

    if not rows:
        # Fallback: take last 30 entries
        with open(csv_path, "r") as f:
            rows = list(csv.DictReader(f))[-30:]

    ok = [r for r in rows if r.get("success") == "True"]
    if not ok:
        return {"ok": 0, "fail": len(rows)}

    def avg(key):
        vals = [float(r[key]) for r in ok if r.get(key)]
        return sum(vals) / len(vals) if vals else 0.0

    def pct(key):
        vals = [float(r[key]) for r in ok if r.get(key) and float(r.get(key)) > 0]
        return sum(vals) / len(vals) if vals else 0.0

    def maxv(key):
        vals = [float(r[key]) for r in ok if r.get(key)]
        return max(vals) if vals else 0.0

    return {
        "ok": len(ok),
        "fail": len(rows) - len(ok),
        "avg_output_tps": avg("output_tps"),
        "avg_total_tps": avg("total_tps"),
        "avg_ttft_ms": avg("ttft_ms"),
        "avg_tpot_ms": avg("tpot_ms"),
        "peak_vram_mb": maxv("vram_peak_mb"),
        "avg_gpu_util_pct": pct("gpu_util_pct"),
        "avg_power_w": pct("power_w"),
    }


def run_benchmark_subprocess(config: OptimizedConfig, port: int = 8000) -> tuple[bool, str]:
    """Run bench_openai_api.py as a subprocess."""
    bench_script = PROJECT_DIR / "scripts" / "bench_openai_api.py"
    cmd = [
        sys.executable, str(bench_script),
        "--url", f"http://localhost:{port}",
        "--prompts", str(PROMPTS_FILE),
        "--output", str(RESULTS_CSV),
        "--model", config.model_path,
        "--backend", "vllm",
        "--config", config.name,
        "--max-model-len", str(config.max_model_len),
        "--gpu-mem-util", str(config.gpu_memory_utilization),
        "--dtype", config.dtype,
        "--quantization", config.quant_method,
        "--repeat", "2",
    ]
    print(f"  🔬 Running benchmark...")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            # Extract key lines
            for line in result.stdout.split("\n"):
                if any(kw in line for kw in ["Output TPS:", "TTFT", "SUMMARY", "Success"]):
                    print(f"    {line.strip()}")
            return True, ""
        else:
            err = result.stderr[-500:] if result.stderr else result.stdout[-500:]
            return False, err
    except subprocess.TimeoutExpired:
        return False, "Benchmark subprocess timed out"
    except Exception as e:
        return False, str(e)


# ============================================================
# Orchestrator
# ============================================================

def sweep(configs: list[OptimizedConfig], port: int = 8000) -> list[SweepResult]:
    results: list[SweepResult] = []
    REPORTS_DIR.mkdir(parents=True, exist_ok=True)

    # Backup old CSV
    if RESULTS_CSV.exists():
        RESULTS_CSV.rename(RESULTS_CSV.with_suffix(".csv.bak"))

    print(f"\n{'='*70}")
    print(f" ADVANCED OPTIMIZATION SWEEP — {len(configs)} configs")
    print(f"{'='*70}")
    for i, c in enumerate(configs, 1):
        print(f"  [{i:2d}] {c.name:50s}  ({c.label[:70]})")
    print(f"{'='*70}\n")

    for idx, config in enumerate(configs, 1):
        result = SweepResult(config=config, success=False)
        t0 = time.time()

        print(f"\n{'#'*70}")
        print(f" [{idx}/{len(configs)}] {config.name}")
        print(f"  attn={config.attention_backend}  kv={config.kv_cache_dtype}  "
              f"chunk={config.enable_chunked_prefill}  bs={config.block_size}  "
              f"seqs={config.max_num_seqs}")
        print(f"{'#'*70}")

        # Start server
        server = VllmServer(config, port=port)
        try:
            server.start()
        except Exception as e:
            result.error = f"START FAIL: {e}"
            results.append(result)
            continue

        # Wait for ready
        ok, msg = server.wait_ready()
        if not ok:
            result.error = msg[:500]
            result.duration_s = time.time() - t0
            server.stop()
            results.append(result)
            print(f"  ❌ Failed to start: {msg[:200]}")
            continue

        print(f"  ✅ Server ready")

        # Run benchmark
        bench_ok, bench_err = run_benchmark_subprocess(config, port=port)

        # Parse results
        metrics = parse_csv_for_config(str(RESULTS_CSV), config.name)
        if metrics:
            result.success = metrics.get("ok", 0) > 0
            result.avg_output_tps = metrics.get("avg_output_tps", 0)
            result.avg_total_tps = metrics.get("avg_total_tps", 0)
            result.avg_ttft_ms = metrics.get("avg_ttft_ms", 0)
            result.avg_tpot_ms = metrics.get("avg_tpot_ms", 0)
            result.peak_vram_mb = metrics.get("peak_vram_mb", 0)
            result.avg_gpu_util_pct = metrics.get("avg_gpu_util_pct", 0)
            result.avg_power_w = metrics.get("avg_power_w", 0)
            result.ok = metrics.get("ok", 0)
            result.fail = metrics.get("fail", 0)

        if not bench_ok:
            result.error = bench_err[:500]
            if not result.success:
                result.error = bench_err[:500] or result.error

        result.duration_s = time.time() - t0
        server.stop()

        status = "✅" if result.success else "❌"
        print(f"  {status} TPS={result.avg_output_tps:.1f}  TTFT={result.avg_ttft_ms:.0f}ms  "
              f"VRAM={result.peak_vram_mb:.0f}MB  [{result.ok}/{result.ok + result.fail}]  "
              f"({result.duration_s:.0f}s)")

        results.append(result)

    return results


# ============================================================
# Report generation
# ============================================================

def generate_report(results: list[SweepResult], mode: str):
    ok = [r for r in results if r.success]
    fail = [r for r in results if not r.success]
    ok_sorted = sorted(ok, key=lambda r: r.avg_output_tps, reverse=True)

    lines = [
        f"# Qwen SpeedLab — Advanced Optimization Report",
        f"",
        f"**Date:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        f"**GPU:** NVIDIA GeForce RTX 3090 (24 GB, SM 8.6)",
        f"**Mode:** {mode}",
        f"**Configs:** {len(results)} tested — {len(ok)} passed, {len(fail)} failed",
        f"",
        "---",
        f"",
    ]

    if ok_sorted:
        best = ok_sorted[0]
        lines += [
            f"## 🏆 Best Configuration",
            f"",
            f"```",
            f"  Config name:      {best.config.name}",
            f"  Model:            {best.config.model_path}",
            f"  Attention:        {best.config.attention_backend}",
            f"  KV cache dtype:   {best.config.kv_cache_dtype}",
            f"  Chunked prefill:  {best.config.enable_chunked_prefill}",
            f"  Block size:       {best.config.block_size}",
            f"  Max seqs:         {best.config.max_num_seqs}",
            f"  Max batched tok:  {best.config.max_num_batched_tokens}",
            f"  GPU mem util:     {best.config.gpu_memory_utilization}",
            f"```",
            f"",
            f"| Metric | Value |",
            f"|--------|-------|",
            f"| **Output TPS** | **{best.avg_output_tps:.1f} tok/s** |",
            f"| Total TPS | {best.avg_total_tps:.1f} tok/s |",
            f"| TTFT | {best.avg_ttft_ms:.0f} ms |",
            f"| TPOT | {best.avg_tpot_ms:.1f} ms |",
            f"| Peak VRAM | {best.peak_vram_mb:.0f} MB |",
            f"| GPU util | {best.avg_gpu_util_pct:.1f}% |",
            f"| Power | {best.avg_power_w:.1f} W |",
            f"| Prompt success | {best.ok}/{best.ok + best.fail} |",
            f"| Test duration | {best.duration_s:.0f}s |",
            f"",
        ]

    # ---- Optimization dimension analysis ----
    lines += [
        "## 📊 Optimization Impact Analysis",
        "",
        "Each table shows the effect of toggling one optimization dimension.",
        "",
    ]

    # Group by dimension
    dimensions = {
        "attention_backend": "Attention Backend",
        "kv_cache_dtype": "KV Cache Dtype",
        "enable_chunked_prefill": "Chunked Prefill",
        "block_size": "Block Size",
        "gpu_memory_utilization": "GPU Memory Utilization",
        "max_num_seqs": "Max Concurrent Sequences",
    }

    for dim, title in dimensions.items():
        groups: dict[str, list[SweepResult]] = {}
        for r in ok:
            key = str(getattr(r.config, dim, "?"))
            groups.setdefault(key, []).append(r)

        if len(groups) <= 1:
            continue

        lines += [f"### {title}", ""]
        lines += ["| Value | Avg TPS | TTFT ms | VRAM MB | Tests |",
                   "|-------|---------|---------|---------|-------|"]
        for key in sorted(groups.keys()):
            grp = groups[key]
            avg_tps = sum(r.avg_output_tps for r in grp) / len(grp)
            avg_ttft = sum(r.avg_ttft_ms for r in grp) / len(grp)
            avg_vram = sum(r.peak_vram_mb for r in grp) / len(grp)
            n = sum(r.ok for r in grp)
            lines.append(f"| `{key}` | {avg_tps:.1f} | {avg_ttft:.0f} | {avg_vram:.0f} | {n} |")
        lines.append("")

    # ---- Full results table ----
    lines += [
        "## 📋 All Configurations",
        "",
        "| # | Name | Attn | KV | Chunk | BS | Seqs | TPS | TTFT | VRAM | Status |",
        "|---|------|------|----|-------|----|------|-----|------|------|--------|",
    ]
    for i, r in enumerate(results, 1):
        if r.success:
            tps = f"{r.avg_output_tps:.1f}"
            ttft = f"{r.avg_ttft_ms:.0f}"
            vram = f"{r.peak_vram_mb:.0f}"
            status = "✅"
        else:
            tps = ttft = vram = "-"
            status = "❌"
        lines.append(
            f"| {i} | {r.config.name} | {r.config.attention_backend} | {r.config.kv_cache_dtype} | "
            f"{r.config.enable_chunked_prefill} | {r.config.block_size} | {r.config.max_num_seqs} | "
            f"{tps} | {ttft} | {vram} | {status} |"
        )
    lines.append("")

    # ---- Failed ----
    if fail:
        lines += ["## ❌ Failed Configurations", ""]
        for r in fail:
            lines.append(f"- **{r.config.name}**: `{r.error[:200]}`")
        lines.append("")

    # ---- Recommendations ----
    lines += [
        "## 💡 Recommendations for RTX 3090",
        "",
    ]

    # Analyze what worked best
    if ok_sorted:
        flashinfer_ok = [r for r in ok if r.config.attention_backend == "flashinfer"]
        flashinfer_tps = sum(r.avg_output_tps for r in flashinfer_ok) / len(flashinfer_ok) if flashinfer_ok else 0
        flashattn_ok = [r for r in ok if r.config.attention_backend == "flash-attn"]
        flashattn_tps = sum(r.avg_output_tps for r in flashattn_ok) / len(flashattn_ok) if flashattn_ok else 0

        fp8kv_ok = [r for r in ok if r.config.kv_cache_dtype == "fp8"]
        fp8kv_tps = sum(r.avg_output_tps for r in fp8kv_ok) / len(fp8kv_ok) if fp8kv_ok else 0
        auto_ok = [r for r in ok if r.config.kv_cache_dtype == "auto"]
        auto_tps = sum(r.avg_output_tps for r in auto_ok) / len(auto_ok) if auto_ok else 0

        lines.append("1. **Attention backend**:")
        if flashinfer_tps > flashattn_tps:
            lines.append(f"   FlashInfer is faster ({flashinfer_tps:.1f} vs {flashattn_tps:.1f} tok/s). Use `--attention-backend flashinfer`.")
        else:
            lines.append(f"   FlashAttention is comparable ({flashattn_tps:.1f} vs {flashinfer_tps:.1f} tok/s). Default is fine.")

        lines.append("")
        lines.append("2. **KV Cache dtype**:")
        if fp8kv_tps > auto_tps * 0.95:
            lines.append(f"   FP8 KV cache saves VRAM with minimal quality loss ({fp8kv_tps:.1f} vs {auto_tps:.1f} tok/s). Use `--kv-cache-dtype fp8`.")
        else:
            lines.append(f"   FP8 KV cache costs speed ({fp8kv_tps:.1f} vs {auto_tps:.1f} tok/s). Keep auto unless you need the VRAM.")

        best_quant = ok_sorted[0].config.quant_method
        lines.append("")
        lines.append(f"3. **Best quantization**: `{best_quant}` — based on measured throughput.")

        best_bs = max(set(r.config.block_size for r in ok), key=lambda bs: sum(r.avg_output_tps for r in ok if r.config.block_size == bs) / max(1, len([r for r in ok if r.config.block_size == bs])))
        lines.append(f"4. **Optimal block size**: `{best_bs}`")

        best_seqs = max(set(r.config.max_num_seqs for r in ok), key=lambda s: sum(r.avg_output_tps for r in ok if r.config.max_num_seqs == s) / max(1, len([r for r in ok if r.config.max_num_seqs == s])))
        lines.append(f"5. **Optimal max_num_seqs**: `{best_seqs}`")

    lines += [
        "",
        "6. **Always enable**: prefix caching, trust-remote-code",
        "7. **Chunked prefill**: helps throughput under concurrent load",
        "8. **Swap space**: 4GB CPU swap acts as a safety net against OOM",
        "",
        "---",
        f"*Report generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*",
        "",
    ]

    SUMMARY_MD.parent.mkdir(parents=True, exist_ok=True)
    with open(SUMMARY_MD, "w") as f:
        f.write("\n".join(lines))
    print(f"\n📄 Report: {SUMMARY_MD}")


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

def main():
    parser = argparse.ArgumentParser(description="Advanced vLLM optimization sweep for RTX 3090")
    parser.add_argument("--mode", default="baseline",
                        choices=["baseline", "shootout", "grid", "single", "list"],
                        help="baseline=A/B each opt, shootout=AWQ vs GPTQ vs FP8, grid=full combos")
    parser.add_argument("--config", type=str, help="Single config name (for --mode single)")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--dry-run", action="store_true", help="List configs, don't run")

    args = parser.parse_args()

    # Config selection
    if args.mode == "baseline":
        configs = baseline_configs()
    elif args.mode == "shootout":
        configs = model_shootout_configs()
    elif args.mode == "grid":
        configs = grid_configs()
    elif args.mode == "single":
        all_configs = baseline_configs() + model_shootout_configs()
        configs = [c for c in all_configs if c.name == args.config]
        if not configs:
            print(f"Unknown config: {args.config}")
            print(f"Available: {sorted(set(c.name for c in all_configs))}")
            sys.exit(1)
    else:  # list
        for c in baseline_configs() + model_shootout_configs():
            print(f"  {c.name:45s}  {c.label}")
        return

    print(f"\n{'='*70}")
    print(f" ADVANCED SWEEP — {args.mode} mode — {len(configs)} config(s)")
    print(f"{'='*70}")

    if args.dry_run:
        for c in configs:
            print(f"  {c.name}")
            print(f"    model={c.model_path}")
            print(f"    attn={c.attention_backend} kv={c.kv_cache_dtype} chunk={c.enable_chunked_prefill}")
            print(f"    bs={c.block_size} seqs={c.max_num_seqs} batched={c.max_num_batched_tokens}")
        return

    # Preflight
    print("🔍 Preflight checks...")
    # GPU
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        )
        free_mb = int(result.stdout.strip())
        print(f"  GPU free: {free_mb} MiB")
        if free_mb < 20000:
            print("  ⚠️  Less than 20GB free — old processes?")
    except Exception as e:
        print(f"  ⚠️  GPU check failed: {e}")

    # Port
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    if sock.connect_ex(('localhost', args.port)) == 0:
        print(f"  ⚠️  Port {args.port} already in use!")
    sock.close()

    # HF_TOKEN
    if not os.environ.get("HF_TOKEN"):
        print("  ⚠️  HF_TOKEN not set")
    print("  ✅ Ready\n")

    # Run
    results = sweep(configs, port=args.port)

    # Report
    generate_report(results, args.mode)

    # Final summary
    ok = [r for r in results if r.success]
    print(f"\n{'='*70}")
    print(f" SWEEP COMPLETE")
    print(f"  Configs: {len(results)}  |  Passed: {len(ok)}  |  Failed: {len(results) - len(ok)}")
    if ok:
        best = max(ok, key=lambda r: r.avg_output_tps)
        print(f"  🏆 Best: {best.config.name} — {best.avg_output_tps:.1f} tok/s")
        print(f"     attn={best.config.attention_backend}  kv={best.config.kv_cache_dtype}  "
              f"bs={best.config.block_size}  seqs={best.config.max_num_seqs}")
    print(f"{'='*70}")


if __name__ == "__main__":
    main()