#!/usr/bin/env python3
"""
bench-llm — Local LLM Benchmarking Tool for LM Studio
=====================================================

Benchmarks local models loaded in LM Studio (localhost:1234, OpenAI-compatible API)
across three dimensions: Speed, Ability, and Agent Fitness.

Usage:
    ./bench-llm <model-id>              # run all tests
    ./bench-llm <model-id> --speed-only # just speed
    ./bench-llm <model-id> --quick      # fast smoke test
    ./bench-llm --list                  # list available models

Output:
    ~/benchmarks/<model-id>-<date>.json  — detailed JSON results
    Terminal summary table

Requirements:
    pip install openai
"""

import argparse
import json
import os
import re
import statistics
import subprocess
import sys
import textwrap
import time
from datetime import datetime, timezone
from pathlib import Path

try:
    import requests
except ImportError:
    requests = None

try:
    from openai import OpenAI
except ImportError:
    print("ERROR: 'openai' package not installed. Run: pip install openai")
    sys.exit(1)

# ──────────────────────────────────────────────────────────────────────────────
# CONFIGURATION
# ──────────────────────────────────────────────────────────────────────────────

API_BASE = "http://localhost:1234/v1"
API_KEY = "lm-studio"  # LM Studio accepts any string
BENCHMARKS_DIR = Path.home() / "benchmarks"
WORKSPACE = Path.home() / ".openclaw" / "workspace"
TIMEOUT_SECONDS = 300  # 5 min per test
WARMUP_PROMPT = "Say hello."
MAX_OUTPUT_TOKENS = 512  # max tokens for generation tests
SPEED_RUNS = 3  # runs per speed test

# ──────────────────────────────────────────────────────────────────────────────
# MODEL FILE RESOLUTION
# ──────────────────────────────────────────────────────────────────────────────

def fetch_lmstudio_models():
    """Fetch model metadata from LM Studio's internal /api/v0/models endpoint."""
    url = "http://localhost:1234/api/v0/models"
    try:
        resp = requests.get(url, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        return data.get("data", [])
    except Exception as e:
        print(f"  ⚠  Warning: Could not fetch LM Studio model metadata: {e}")
        return []


def find_gguf_files():
    """Find all non-mmproj GGUF files under ~/.lmstudio.
    Returns list of {'path': str, 'size_bytes': int}.
    """
    gguf_dir = Path.home() / ".lmstudio"
    if not gguf_dir.exists():
        return []
    try:
        result = subprocess.run(
            ["find", str(gguf_dir), "-name", "*.gguf"],
            capture_output=True, text=True, timeout=15
        )
        paths = [p.strip() for p in result.stdout.strip().split("\n") if p.strip()]
    except Exception:
        return []

    files = []
    for p in paths:
        if "mmproj" in p.lower():
            continue
        try:
            size = os.path.getsize(p)
        except OSError:
            size = 0
        files.append({"path": p, "size_bytes": size})
    return files


def _normalize_publisher(pub):
    """Normalize publisher names for fuzzy matching."""
    if not pub:
        return ""
    return pub.lower().replace("_", "-").replace(" ", "-").strip()


def _normalize_arch(arch):
    """Normalize architecture names for fuzzy matching."""
    if not arch:
        return ""
    return arch.lower().replace("_", "-").replace(" ", "-").strip()


def _human_size(bytes_val):
    """Convert bytes to human-readable string."""
    if bytes_val is None or bytes_val <= 0:
        return "0 B"
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if bytes_val < 1024:
            return f"{bytes_val:.1f} {unit}" if unit != "B" else f"{bytes_val} B"
        bytes_val /= 1024
    return f"{bytes_val:.1f} PB"


def resolve_model_to_gguf(model_id, models_meta=None, gguf_files=None):
    """Resolve a model ID to its backing GGUF file(s).

    Returns dict with:
        status: READY | UNAVAILABLE | UNKNOWN
        gguf_path: str or None
        gguf_size_bytes: int or None
        gguf_size_human: str
        arch: str or None
        quant: str or None
        publisher: str or None
    """
    if models_meta is None:
        models_meta = fetch_lmstudio_models()
    if gguf_files is None:
        gguf_files = find_gguf_files()

    # Find model metadata
    model_meta = None
    for m in models_meta:
        if m.get("id") == model_id:
            model_meta = m
            break

    if model_meta is None:
        return {
            "status": "UNKNOWN",
            "gguf_path": None,
            "gguf_size_bytes": None,
            "gguf_size_human": "N/A",
            "arch": None,
            "quant": None,
            "publisher": None,
        }

    arch = model_meta.get("arch", "")
    publisher = model_meta.get("publisher", "")
    quant = model_meta.get("quantization", "")

    # Try to match GGUF files by arch + publisher
    matched = []
    norm_pub = _normalize_publisher(publisher)
    norm_arch = _normalize_arch(arch)

    for gf in gguf_files:
        fname_lower = os.path.basename(gf["path"]).lower()
        dir_lower = str(gf["path"]).lower()
        pub_dir = normalize_publisher_from_path(gf["path"])

        # Score matching: publisher match is strongest signal
        score = 0
        if norm_pub and pub_dir and norm_pub in pub_dir:
            score += 40
        elif norm_pub and norm_pub in dir_lower:
            score += 30
        if norm_arch and norm_arch in dir_lower:
            score += 20
        if norm_arch and any(token in fname_lower for token in norm_arch.split("-")):
            score += 10
        if quant and quant.lower().replace("_","").replace(".","") in fname_lower.replace("_","").replace(".",""):
            score += 15

        if score >= 25:
            matched.append((score, gf))

    matched.sort(key=lambda x: x[0], reverse=True)

    if matched:
        gf = matched[0][1]
        return {
            "status": "READY",
            "gguf_path": gf["path"],
            "gguf_size_bytes": gf["size_bytes"],
            "gguf_size_human": _human_size(gf["size_bytes"]),
            "arch": arch,
            "quant": quant,
            "publisher": publisher,
        }
    else:
        return {
            "status": "UNAVAILABLE",
            "gguf_path": None,
            "gguf_size_bytes": None,
            "gguf_size_human": "N/A",
            "arch": arch,
            "quant": quant,
            "publisher": publisher,
        }


def normalize_publisher_from_path(path):
    """Extract publisher from GGUF path: ~/.lmstudio/models/<publisher>/..."""
    parts = Path(path).parts
    try:
        models_idx = parts.index("models")
        if models_idx + 1 < len(parts):
            return _normalize_publisher(parts[models_idx + 1])
    except ValueError:
        pass
    return ""


def preflight_check(model_id, output_json=False):
    """Run pre-flight check for a model and print results."""
    models_meta = fetch_lmstudio_models()
    gguf_files = find_gguf_files()
    info = resolve_model_to_gguf(model_id, models_meta, gguf_files)

    if output_json:
        print(json.dumps(info, indent=2, default=str))
        return info

    status_icon = {"READY": "✅", "UNAVAILABLE": "❌", "UNKNOWN": "❓"}
    icon = status_icon.get(info["status"], "❓")

    print(f"\n{icon} Pre-flight Check: {model_id}")
    print(f"   Status:       {info['status']}")
    print(f"   Architecture: {info['arch'] or 'N/A'}")
    print(f"   Quantization: {info['quant'] or 'N/A'}")
    print(f"   Publisher:    {info['publisher'] or 'N/A'}")
    if info["gguf_path"]:
        print(f"   GGUF Path:    {info['gguf_path']}")
        print(f"   GGUF Size:    {info['gguf_size_human']} ({info['gguf_size_bytes']:,} bytes)")
    else:
        print(f"   GGUF Path:    <not found on disk>")
        print(f"   GGUF Size:    N/A")
    print()
    return info


# ──────────────────────────────────────────────────────────────────────────────
# SPEED SANITY VALIDATION
# ──────────────────────────────────────────────────────────────────────────────

# Rough performance ceilings for 7900 XTX-level hardware (conservative estimates)
# T/s by model size proxy: file size in GB, with a generous upper bound
_PERFORMANCE_CEILING = {
    # (min_size_gb, max_size_gb): max_plausible_tps
    (0, 4):      200,   # <4B class
    (4, 8):      150,   # 4-8B class
    (8, 16):     110,   # 8-13B class
    (16, 24):    80,    # 13-20B class
    (24, 36):    70,    # 24-30B class (e.g., 27B at Q4/Q5)
    (36, 50):    50,    # 30-40B class
    (50, 80):    35,    # 50-70B class
    (80, 150):   20,    # 80-120B class (e.g., MoE dense-equivalent)
    (150, 9999): 10,    # 120B+ — likely MoE or sharded
}

MIN_PLAUSIBLE_TTFT_SECONDS = 0.05


def validate_speed_plausibility(speed_results, model_gguf_info):
    """Check speed benchmark results for physical plausibility.

    Returns dict with warnings, confidence rating, and individual checks.
    """
    warnings = []
    checks = []

    gguf_size_bytes = model_gguf_info.get("gguf_size_bytes")
    if not gguf_size_bytes or gguf_size_bytes <= 0:
        return {
            "confidence": "LOW",
            "warnings": ["No GGUF file size available; cannot validate plausibility"],
            "checks": [],
        }

    size_gb = gguf_size_bytes / (1024 ** 3)

    # Determine performance ceiling
    max_plausible_tps = None
    for (lo, hi), ceiling in _PERFORMANCE_CEILING.items():
        if lo <= size_gb < hi:
            max_plausible_tps = ceiling
            break
    if max_plausible_tps is None:
        max_plausible_tps = 10

    all_tps_values = []
    all_ttft_values = []

    # Collect all T/s values
    for prompt_key, speed_data in speed_results.items():
        if speed_data.get("error"):
            continue
        runs = speed_data.get("runs", [])
        for run in runs:
            tps = run.get("tokens_per_sec")
            ttft = run.get("ttft_s")
            if tps is not None and tps > 0:
                all_tps_values.append(tps)
            if ttft is not None and ttft > 0:
                all_ttft_values.append(ttft)

    if not all_tps_values:
        return {
            "confidence": "LOW",
            "warnings": ["No valid T/s measurements"],
            "checks": [],
        }

    mean_tps = statistics.mean(all_tps_values)
    mean_ttft = statistics.mean(all_ttft_values) if all_ttft_values else None

    # Check TPS plausibility
    if mean_tps > max_plausible_tps * 1.3:
        warnings.append(
            f"Mean T/s ({mean_tps:.1f}) exceeds plausible ceiling ({max_plausible_tps}) "
            f"for {size_gb:.1f} GB model — possible measurement error or speculative decoding contamination"
        )
        checks.append({"check": "tps_ceiling", "passed": False, "detail": warnings[-1]})
    elif mean_tps > max_plausible_tps:
        warnings.append(
            f"Mean T/s ({mean_tps:.1f}) is above typical ceiling ({max_plausible_tps}) "
            f"for {size_gb:.1f} GB model"
        )
        checks.append({"check": "tps_ceiling", "passed": False, "detail": warnings[-1]})
    else:
        checks.append({"check": "tps_ceiling", "passed": True})

    # Check TTFT plausibility
    if mean_ttft is not None and mean_ttft < MIN_PLAUSIBLE_TTFT_SECONDS:
        warnings.append(
            f"Mean TTFT ({mean_ttft:.4f}s) is suspiciously low (<{MIN_PLAUSIBLE_TTFT_SECONDS}s) "
            f"for {size_gb:.1f} GB model — may indicate pre-warmed cache or measurement error"
        )
        checks.append({"check": "ttft_plausible", "passed": False, "detail": warnings[-1]})
    elif mean_ttft is not None:
        checks.append({"check": "ttft_plausible", "passed": True})
    else:
        checks.append({"check": "ttft_plausible", "passed": True, "detail": "No TTFT data"})

    # Determine confidence
    all_checks_passed = all(c.get("passed", True) for c in checks)
    if not warnings:
        confidence = "HIGH"
    elif len(warnings) == 1 and "above typical" in warnings[0]:
        confidence = "MEDIUM"
    else:
        confidence = "LOW"

    return {
        "confidence": confidence,
        "warnings": warnings,
        "checks": checks,
        "model_size_gb": round(size_gb, 2),
        "max_plausible_tps": max_plausible_tps,
        "mean_tps": round(mean_tps, 1) if mean_tps is not None else None,
        "mean_ttft": round(mean_ttft, 4) if mean_ttft is not None else None,
    }


# ──────────────────────────────────────────────────────────────────────────────
# SPEED TEST PROMPTS (short, medium, long)
# ──────────────────────────────────────────────────────────────────────────────

SPEED_PROMPTS = {
    "short_50": (
        "Explain what a variable is in programming. Be concise."
        # ~50 chars of instruction
    ),
    "medium_200": (
        "Explain the concept of recursion in computer science. "
        "Include a simple example in Python. Be thorough but concise. "
        "Cover base case, recursive case, and stack implications."
        # ~200 chars
    ),
    "long_800": (
        "Write a detailed explanation of the following programming paradigms: "
        "object-oriented programming, functional programming, and procedural programming. "
        "For each paradigm, describe its core philosophy, key concepts, typical use cases, "
        "advantages, disadvantages, and provide a short code example in a language "
        "commonly associated with that paradigm. Compare and contrast the three approaches. "
        "Discuss how modern languages often blend multiple paradigms. Be thorough."
        # ~800 chars
    ),
}

# ──────────────────────────────────────────────────────────────────────────────
# ABILITY TEST DEFINITIONS
# ──────────────────────────────────────────────────────────────────────────────

def check_coding_test(response_text: str) -> dict:
    """Test: Write median_of_list(numbers) that passes unit tests."""
    # Extract Python function from response
    code_match = re.search(r'```(?:python)?\s*(.*?)```', response_text, re.DOTALL)
    if code_match:
        code = code_match.group(1).strip()
    else:
        # Try to find a def statement and grab everything after it
        def_match = re.search(r'(def\s+median_of_list.*?)(?:\n\n|\Z)', response_text, re.DOTALL)
        if def_match:
            code = def_match.group(1).strip()
        else:
            code = response_text.strip()

    # Execute in isolated namespace
    namespace = {}
    try:
        exec(code, namespace)
    except Exception as e:
        return {"passed": False, "error": f"Compilation error: {e}", "details": []}

    if 'median_of_list' not in namespace:
        return {"passed": False, "error": "Function 'median_of_list' not found", "details": []}

    func = namespace['median_of_list']
    test_cases = [
        ("[1,2,3] -> 2",             [1, 2, 3],          2),
        ("[1,2,3,4] -> 2.5",         [1, 2, 3, 4],       2.5),
        ("[] -> None",               [],                  None),
        ("[5] -> 5",                 [5],                 5),
        ("[1,100,2,99,3] -> 3",      [1, 100, 2, 99, 3], 3),
        ("[-5,-1,-3] -> -3",         [-5, -1, -3],       -3),
    ]

    details = []
    all_passed = True
    for desc, inp, expected in test_cases:
        try:
            result = func(inp.copy() if isinstance(inp, list) else inp)
            if result == expected:
                details.append(f"  PASS: {desc}")
            else:
                details.append(f"  FAIL: {desc} (got {result})")
                all_passed = False
        except Exception as e:
            details.append(f"  FAIL: {desc} (exception: {e})")
            all_passed = False

    return {"passed": all_passed, "details": details, "error": None if all_passed else "Some tests failed"}


def check_logic_test(response_text: str) -> dict:
    """Test: Pet-ownership logic puzzle — answer: Alice=Fish, Bob=Bird, Carol=Cat, Dan=Dog."""
    # Extract the answer from the response
    text = response_text.lower()

    # Look for assignment patterns
    alice_pet = re.findall(r'alice\s*[=:>-]+\s*(\w+)', text)
    bob_pet   = re.findall(r'bob\s*[=:>-]+\s*(\w+)', text)
    carol_pet = re.findall(r'carol\s*[=:>-]+\s*(\w+)', text)
    dan_pet   = re.findall(r'dan\s*[=:>-]+\s*(\w+)', text)

    # Also try "Alice owns/has a X" patterns
    if not alice_pet:
        alice_pet = re.findall(r'alice[^.]*?(?:owns?|has?)\s*(?:a|an)?\s*(\w+)', text)
    if not bob_pet:
        bob_pet = re.findall(r'bob[^.]*?(?:owns?|has?)\s*(?:a|an)?\s*(\w+)', text)
    if not carol_pet:
        carol_pet = re.findall(r'carol[^.]*?(?:owns?|has?)\s*(?:a|an)?\s*(\w+)', text)
    if not dan_pet:
        dan_pet = re.findall(r'dan[^.]*?(?:owns?|has?)\s*(?:a|an)?\s*(\w+)', text)

    expected = {
        "alice": "fish",
        "bob": "bird",
        "carol": "cat",
        "dan": "dog",
    }

    def get_pet(name, matches):
        for m in matches:
            pet = m.strip().lower().rstrip('s')  # handle plural
            if pet in ("fish", "bird", "cat", "dog"):
                return pet
        return None

    found = {
        "alice": get_pet("alice", alice_pet),
        "bob": get_pet("bob", bob_pet),
        "carol": get_pet("carol", carol_pet),
        "dan": get_pet("dan", dan_pet),
    }

    correct = sum(1 for k, v in expected.items() if found[k] == v)
    total = len(expected)
    passed = correct == total

    details = [f"  Expected: {expected}"]
    details.append(f"  Found:    {found}")
    details.append(f"  Correct:  {correct}/{total}")

    # Fallback: if regex didn't find all, check for direct statement
    if correct < total:
        # Check for "Alice has fish" etc anywhere
        for person, pet in expected.items():
            if found[person] != pet:
                if re.search(rf'{person}\s+[^.]*?\b{pet}\b', text, re.IGNORECASE):
                    correct += 1
                    details.append(f"  Fallback found: {person} = {pet}")

    passed = correct == total
    return {
        "passed": passed,
        "correct": correct,
        "total": total,
        "expected": expected,
        "found": found,
        "details": details,
    }


def check_json_test(response_text: str) -> dict:
    """Test: Structured JSON output with specific schema."""
    # Extract JSON from response (in code block or raw)
    json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', response_text, re.DOTALL)
    if not json_match:
        json_match = re.search(r'(\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\})', response_text, re.DOTALL)

    details = []

    if not json_match:
        return {"passed": False, "error": "No JSON object found in response", "details": details}

    try:
        data = json.loads(json_match.group(1))
    except json.JSONDecodeError as e:
        return {"passed": False, "error": f"Invalid JSON: {e}", "details": details}

    # Validate schema
    checks = []
    checks.append(("Has 'name' (string)", isinstance(data.get("name"), str)))
    checks.append(("Has 'age' (int)", isinstance(data.get("age"), int)))
    checks.append(("Has 'skills' (list of strings)", 
                   isinstance(data.get("skills"), list) and 
                   all(isinstance(s, str) for s in data.get("skills", []))))
    checks.append(("Has 'active' (bool)", isinstance(data.get("active"), bool)))

    for desc, ok in checks:
        details.append(f"  {'PASS' if ok else 'FAIL'}: {desc}")

    all_passed = all(ok for _, ok in checks)
    return {
        "passed": all_passed,
        "checks": {desc: ok for desc, ok in checks},
        "parsed": data,
        "details": details,
    }


def check_instruction_test(response_text: str) -> dict:
    """Test: Write numbers 1-10, mark evens with *, sum evens."""
    details = []

    # Check: numbers 1-10 present (in order)
    numbers_present = all(str(i) in response_text for i in range(1, 11))
    details.append(f"  {'PASS' if numbers_present else 'FAIL'}: Numbers 1-10 present")

    # Check: evens marked with *
    evens_marked = all(
        bool(re.search(rf'\b{num}\s*\*', response_text)) 
        for num in [2, 4, 6, 8, 10]
    )
    details.append(f"  {'PASS' if evens_marked else 'FAIL'}: Even numbers marked with *")

    # Check: sum of evens = 30
    sum_match = re.search(r'SUM\s*[=:>-]*\s*(\d+)', response_text, re.IGNORECASE)
    if sum_match:
        sum_val = int(sum_match.group(1))
        sum_correct = sum_val == 30
        details.append(f"  {'PASS' if sum_correct else 'FAIL'}: SUM={sum_val} (expected 30)")
    else:
        # Try to find just 30 near "sum"
        if re.search(r'sum[^0-9]*30', response_text, re.IGNORECASE):
            sum_correct = True
            details.append("  PASS: SUM=30 found (loose match)")
        else:
            sum_correct = False
            details.append("  FAIL: SUM value not found")

    # Check format
    has_numbers_section = bool(re.search(r'NUMBERS?\s*[=:>-]', response_text, re.IGNORECASE))
    has_sum_section = bool(re.search(r'SUM\s*[=:>-]', response_text, re.IGNORECASE))

    passed = numbers_present and evens_marked and sum_correct

    return {
        "passed": passed,
        "checks": {
            "numbers_1_10": numbers_present,
            "evens_marked": evens_marked,
            "sum_correct": sum_correct,
            "has_numbers_section": has_numbers_section,
            "has_sum_section": has_sum_section,
        },
        "details": details,
    }


def check_summarization_test(response_text: str, source_text: str) -> dict:
    """Test: Summarize a paragraph, check key facts preserved."""
    # Source paragraph about a fictional technology
    key_facts = [
        "helium",
        "quantum",
        "2027",       # or "twenty twenty-seven"
        "silicon",
        "terahertz",  # or "THz"
        "celsius",    # or "C"
    ]

    details = []
    found = []
    text_lower = response_text.lower()

    for fact in key_facts:
        # Check for the fact or common variations
        patterns = [fact]
        if fact == "helium":
            patterns.append("he")
        if fact == "quantum":
            patterns.append("qubit")
        if fact == "2027":
            patterns.extend(["twenty.twenty.seven", "27"])
        if fact == "silicon":
            patterns.append("si")
        if fact == "terahertz":
            patterns.append("thz")
        if fact == "celsius":
            patterns.append("°c")

        matched = any(
            re.search(p, text_lower) for p in patterns
        )
        found.append(matched)
        details.append(f"  {'PASS' if matched else 'FAIL'}: Contains '{fact}'")

    score = sum(found)
    passed = score >= 4  # At least 4 of 6 key facts

    return {
        "passed": passed,
        "key_facts_total": len(key_facts),
        "key_facts_found": score,
        "details": details,
    }


# ──────────────────────────────────────────────────────────────────────────────
# AGENT FITNESS TEST DEFINITIONS
# ──────────────────────────────────────────────────────────────────────────────

BULK_WORK_PROMPT = """You are Doxy, a bulk-work agent. Your task:

1. Search /var/log/syslog for lines containing "ERROR" or "CRITICAL" from the last 7 days
2. For each matching line, extract the timestamp, service name, and error message
3. Group errors by service name
4. For each service with more than 5 errors, compile a summary report
5. Output the report as JSON with this schema:
   {
     "services": [
       {
         "name": "<service_name>",
         "error_count": <int>,
         "first_seen": "<ISO timestamp>",
         "last_seen": "<ISO timestamp>",
         "sample_errors": ["<msg1>", "<msg2>", "<msg3>"]
       }
     ],
     "total_errors": <int>,
     "scan_period": "last 7 days"
   }

Reply with: (1) ACKNOWLEDGMENT that you understand the task, (2) your APPROACH outline, (3) any BLOCKERS you identify, (4) the OUTPUT in the exact format requested."""


def check_agent_acknowledgment(response_text: str) -> dict:
    """Check if the model acknowledges the task, outlines approach, identifies blockers."""
    text = response_text.lower()
    details = []

    ack = bool(re.search(r'(acknowledg|underst(?:oo|a)nd|received|got it|I\'ll|I will)', text, re.IGNORECASE))
    details.append(f"  {'PASS' if ack else 'FAIL'}: Acknowledges task")

    approach = bool(re.search(r'(approach|plan|steps?|method|strategy|outline)', text, re.IGNORECASE))
    details.append(f"  {'PASS' if approach else 'FAIL'}: Outlines approach")

    blockers = bool(re.search(r'(blocker|limitation|constraint|can\'t|cannot|unable|note|however|but|challenge)', text, re.IGNORECASE))
    details.append(f"  {'PASS' if blockers else 'FAIL'}: Identifies blockers/limitations")

    passed = ack and approach  # Blockers is a bonus, not required
    return {
        "passed": passed,
        "checks": {"acknowledged": ack, "approach_outlined": approach, "blockers_identified": blockers},
        "details": details,
    }


def check_format_adherence(response_text: str) -> dict:
    """Check if the bulk-work output matches the requested JSON schema."""
    details = []

    # Extract JSON from response
    json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', response_text, re.DOTALL)
    if not json_match:
        json_match = re.search(r'(\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\})', response_text, re.DOTALL)

    if not json_match:
        return {"passed": False, "error": "No JSON found in response", "details": details}

    try:
        data = json.loads(json_match.group(1))
    except json.JSONDecodeError as e:
        return {"passed": False, "error": f"Invalid JSON: {e}", "details": details}

    checks = []
    checks.append(("Has 'services' array", isinstance(data.get("services"), list)))
    checks.append(("Has 'total_errors' (int)", isinstance(data.get("total_errors"), int)))
    checks.append(("Has 'scan_period' (string)", isinstance(data.get("scan_period"), str)))

    if isinstance(data.get("services"), list):
        svc_ok = all(
            isinstance(s, dict) and 
            "name" in s and 
            "error_count" in s
            for s in data["services"]
        )
        checks.append(("Services have required fields", svc_ok))

    for desc, ok in checks:
        details.append(f"  {'PASS' if ok else 'FAIL'}: {desc}")

    passed = all(ok for _, ok in checks)
    return {
        "passed": passed,
        "checks": {desc: ok for desc, ok in checks},
        "details": details,
    }


def check_conciseness(response_text: str, input_chars: int) -> dict:
    """Measure output-to-input ratio, penalize verbosity."""
    output_chars = len(response_text)
    ratio = output_chars / max(input_chars, 1)

    # For a bulk-work task, output can be larger (it's producing a report)
    # But we're measuring the ratio as data — extremely verbose models (ratio > 15)
    # might waste tokens in agent workflows
    details = []
    details.append(f"  Input chars:  {input_chars}")
    details.append(f"  Output chars: {output_chars}")
    details.append(f"  Ratio:        {ratio:.2f}")
    
    if ratio > 20:
        rating = "VERY_VERBOSE"
        details.append("  Rating: VERY_VERBOSE (>20x)")
    elif ratio > 10:
        rating = "VERBOSE"
        details.append("  Rating: VERBOSE (>10x)")
    elif ratio > 5:
        rating = "NORMAL"
        details.append("  Rating: NORMAL (5-10x)")
    else:
        rating = "CONCISE"
        details.append("  Rating: CONCISE (<5x)")

    passed = ratio <= 20  # Only fail if extremely verbose
    return {
        "passed": passed,
        "input_chars": input_chars,
        "output_chars": output_chars,
        "ratio": round(ratio, 2),
        "rating": rating,
        "details": details,
    }


# ──────────────────────────────────────────────────────────────────────────────
# TEST DEFINITIONS REGISTRY
# ──────────────────────────────────────────────────────────────────────────────

ABILITY_TESTS = [
    {
        "id": "coding",
        "name": "Code Generation (median_of_list)",
        "prompt": textwrap.dedent("""\
            Write a Python function called `median_of_list(numbers)` that returns the median
            of a list of numbers. Requirements:
            - If the list is empty, return None
            - If the list has an odd number of elements, return the middle element when sorted
            - If the list has an even number of elements, return the average of the two middle elements
            - The function should NOT modify the input list
            Output ONLY the function code in a Python code block. No explanation needed."""),
        "system": "You are a skilled Python programmer. Write clean, correct code. Output only code.",
        "temperature": 0,
        "max_tokens": 512,
        "checker": check_coding_test,
    },
    {
        "id": "logic",
        "name": "Logic Puzzle (Pet Ownership)",
        "prompt": textwrap.dedent("""\
            Solve this logic puzzle:
            
            Four friends — Alice, Bob, Carol, and Dan — each own exactly one pet.
            The pets are: a fish, a bird, a cat, and a dog. No two friends have the same pet.
            
            Clues:
            1. Alice's pet lives in water.
            2. Bob's pet can sing beautifully.
            3. Carol's favorite wild animal is a tiger, and her pet is closely related.
            4. Dan's pet is known as "man's best friend."
            
            Who owns which pet? Output your answer in this exact format:
            Alice=<pet>, Bob=<pet>, Carol=<pet>, Dan=<pet>"""),
        "system": "You are a logical reasoning expert. Think step by step, then give the final answer.",
        "temperature": 0,
        "max_tokens": 512,
        "checker": check_logic_test,
    },
    {
        "id": "json",
        "name": "JSON Compliance",
        "prompt": textwrap.dedent("""\
            Return a JSON object with these exact keys:
            - name: a string (use any name)
            - age: an integer (use any age)
            - skills: an array of strings (at least 2 skills)
            - active: a boolean
            
            Output ONLY the JSON object, nothing else. No markdown, no explanation."""),
        "system": "You are an API that returns only valid JSON. Never include any text outside the JSON.",
        "temperature": 0,
        "max_tokens": 384,
        "checker": check_json_test,
    },
    {
        "id": "instruction",
        "name": "Instruction Following",
        "prompt": textwrap.dedent("""\
            Write the numbers 1 through 10. Mark each even number with an asterisk (*).
            Then write the sum of all even numbers.
            
            Format your output EXACTLY like this:
            NUMBERS: 1, 2*, 3, 4*, 5, 6*, 7, 8*, 9, 10*
            SUM: 30"""),
        "system": "Follow instructions exactly. Output only the requested format.",
        "temperature": 0,
        "max_tokens": 256,
        "checker": check_instruction_test,
    },
    {
        "id": "summarization",
        "name": "Summarization (Key Facts)",
        "prompt": textwrap.dedent("""\
            Summarize the following paragraph in 2-3 sentences. Preserve all key technical facts:
            
            "Helium-based quantum computing took a major leap forward in 2027 when researchers at
            the Zurich Institute demonstrated a 1000-qubit processor operating at 0.01 Kelvin using
            superfluid helium-4 as the quantum bus. Unlike traditional silicon-based approaches that
            require massive dilution refrigerators, the helium platform achieves coherence times exceeding
            10 seconds at a fraction of the cooling cost. The team projects that by 2030, room-temperature
            quantum operations could be feasible using helium-vacancy centers in diamond, potentially
            bringing quantum computing to consumer devices. Early benchmarks show the helium processor
            solving optimization problems 10,000x faster than classical supercomputers at 2.5 terahertz
            clock speeds, all while consuming less power than a household microwave."""),
        "system": "You are a precise summarizer. Include all key facts in your summary.",
        "temperature": 0,
        "max_tokens": 512,
        "checker": lambda r, _s=None: check_summarization_test(r, _s or ""),
    },
]

AGENT_FITNESS_TESTS = [
    {
        "id": "acknowledgment",
        "name": "Task Acknowledgment",
        "prompt": BULK_WORK_PROMPT,
        "system": "You are Doxy, a capable bulk-work agent. You analyze tasks, identify blockers, and produce structured output.",
        "temperature": 0,
        "max_tokens": 1024,
        "checker": check_agent_acknowledgment,
    },
    {
        "id": "format",
        "name": "Format Adherence",
        "prompt": BULK_WORK_PROMPT,
        "system": "You are Doxy, a capable bulk-work agent. Always produce output in the exact format requested.",
        "temperature": 0,
        "max_tokens": 1024,
        "checker": check_format_adherence,
    },
    {
        "id": "conciseness",
        "name": "Conciseness (Output/Input Ratio)",
        "prompt": BULK_WORK_PROMPT,
        "system": "You are Doxy, a concise and efficient bulk-work agent. Be thorough but avoid unnecessary verbosity.",
        "temperature": 0,
        "max_tokens": 1024,
        "checker": lambda r, inp=len(BULK_WORK_PROMPT): check_conciseness(r, inp),
    },
]

# Note: The bulk-work task is shared across agent fitness tests,
# but we run them separately with different system prompts to test
# distinct agent qualities. We run it once for all by making the
# fitness tests use the same prompt and checking different aspects.


# ──────────────────────────────────────────────────────────────────────────────
# API HELPERS
# ──────────────────────────────────────────────────────────────────────────────

def get_client():
    """Create an OpenAI-compatible client for LM Studio."""
    return OpenAI(base_url=API_BASE, api_key=API_KEY, timeout=TIMEOUT_SECONDS)


def list_models():
    """List available models from LM Studio."""
    try:
        client = get_client()
        models = client.models.list()
        return [m.id for m in models.data]
    except Exception as e:
        print(f"ERROR: Cannot connect to LM Studio at {API_BASE}")
        print(f"  {e}")
        return []


def unload_models():
    """Kill llama-server processes to free VRAM, then wait for LM Studio to show empty."""
    print("  🧹 Unloading models...")
    subprocess.run(["pkill", "-f", "llama-server"], capture_output=True)
    time.sleep(3)
    
    # Verify no models are loaded
    for attempt in range(5):
        try:
            client = get_client()
            models = client.models.list()
            if len(models.data) == 0:
                print("  ✓ VRAM cleared, no models loaded")
                return True
        except Exception:
            # If API is down entirely, that's also "no models"
            print("  ✓ llama-server appears stopped")
            return True
        time.sleep(2)
    
    print("  ⚠  Warning: Models may still be loaded after cleanup attempt")
    return False


def stream_chat(client, model_id, messages, max_tokens=256, temperature=0, timeout=TIMEOUT_SECONDS):
    """Stream a chat completion and return timing data + full response.

    Handles both normal and reasoning models (which put output in reasoning_content).
    TTFT is measured from first token of ANY kind.
    Generation TPS is calculated from content tokens (or reasoning tokens if no content).
    """
    response_data = {
        "content": "",
        "reasoning_content": "",
        "ttft_s": None,              # Time to first token (any type)
        "ttft_reasoning_s": None,    # Time to first reasoning token
        "ttft_content_s": None,      # Time to first content token
        "total_time_s": None,
        "content_token_times": [],   # Timestamps of content token arrivals
        "reasoning_token_times": [], # Timestamps of reasoning token arrivals
        "finish_reason": None,
        "usage": None,
        "error": None,
        "is_reasoning_model": False,
    }

    start_time = time.perf_counter()
    first_any_token_time = None
    
    try:
        stream = client.chat.completions.create(
            model=model_id,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
            stream_options={"include_usage": True},
        )

        for chunk in stream:
            if chunk.usage:
                response_data["usage"] = {
                    "prompt_tokens": chunk.usage.prompt_tokens,
                    "completion_tokens": chunk.usage.completion_tokens,
                    "total_tokens": chunk.usage.total_tokens,
                }

            if not chunk.choices:
                continue

            delta = chunk.choices[0].delta
            now = time.perf_counter()

            # Track first token of any type
            is_first_token = (first_any_token_time is None)
            
            # Capture reasoning content if present
            got_reasoning = False
            if hasattr(delta, 'reasoning_content') and delta.reasoning_content:
                got_reasoning = True
                if is_first_token:
                    first_any_token_time = now
                    response_data["ttft_s"] = now - start_time
                    response_data["ttft_reasoning_s"] = now - start_time
                    response_data["is_reasoning_model"] = True
                response_data["reasoning_token_times"].append(now)
                response_data["reasoning_content"] += delta.reasoning_content

            # Capture content tokens
            if delta.content:
                if is_first_token:
                    first_any_token_time = now
                    response_data["ttft_s"] = now - start_time
                    response_data["ttft_content_s"] = now - start_time
                elif response_data["ttft_content_s"] is None:
                    # First content token came after reasoning tokens
                    response_data["ttft_content_s"] = now - start_time
                response_data["content_token_times"].append(now)
                response_data["content"] += delta.content

            if chunk.choices[0].finish_reason:
                response_data["finish_reason"] = chunk.choices[0].finish_reason

    except Exception as e:
        response_data["error"] = str(e)
        return response_data

    end_time = time.perf_counter()
    response_data["total_time_s"] = end_time - start_time

    # Fallback TTFT if no tokens arrived (unlikely but safe)
    if response_data["ttft_s"] is None:
        response_data["ttft_s"] = response_data["total_time_s"]

    return response_data


def get_effective_response(resp_data):
    """Return the effective response text: content if available, else reasoning_content."""
    content = resp_data.get("content", "").strip()
    if content:
        return content
    return resp_data.get("reasoning_content", "").strip()


def calc_tps(resp_data):
    """Calculate tokens-per-second for generation phase.
    
    Uses content tokens if available, otherwise reasoning tokens.
    Handles edge case where all tokens arrive in one burst (ttft ≈ total_time).
    """
    comp_tokens = resp_data.get("usage", {}).get("completion_tokens") or 0
    total_time = resp_data.get("total_time_s") or 0
    ttft = resp_data.get("ttft_s") or 0

    # Generation time = total time - ttft (or use total_time if ttft ≈ total_time)
    gen_time = total_time - ttft if total_time > ttft + 0.001 else total_time
    
    tps = comp_tokens / gen_time if gen_time > 0 else 0
    return tps


def run_speed_test(client, model_id, prompt, label, runs=SPEED_RUNS):
    """Run a speed benchmark with multiple iterations."""
    results = []
    
    for i in range(runs):
        messages = [
            {"role": "system", "content": "Be thorough and detailed."},
            {"role": "user", "content": prompt},
        ]
        
        resp = stream_chat(client, model_id, messages, max_tokens=MAX_OUTPUT_TOKENS, temperature=0)
        
        if resp["error"]:
            return {"error": resp["error"], "label": label, "runs": results}
        
        # Calculate tokens/sec
        tps = calc_tps(resp)
        completion_tokens = resp.get("usage", {}).get("completion_tokens") or 0
        prefill_tokens = resp.get("usage", {}).get("prompt_tokens") or 0
        prefill_tps = prefill_tokens / resp["ttft_s"] if resp["ttft_s"] and resp["ttft_s"] > 0 else 0
        
        results.append({
            "run": i + 1,
            "ttft_s": round(resp["ttft_s"], 3) if resp["ttft_s"] else None,
            "total_time_s": round(resp["total_time_s"], 3),
            "tokens_per_sec": round(tps, 1),
            "prefill_tps": round(prefill_tps, 1),
            "prompt_tokens": prefill_tokens,
            "completion_tokens": completion_tokens,
            "finish_reason": resp["finish_reason"],
            "error": resp["error"],
        })

    # Calculate stats
    ttfts = [r["ttft_s"] for r in results if r["ttft_s"] is not None]
    tpss = [r["tokens_per_sec"] for r in results]
    total_times = [r["total_time_s"] for r in results]

    def stats(values):
        if not values:
            return {"mean": None, "median": None, "min": None, "max": None}
        return {
            "mean": round(statistics.mean(values), 3),
            "median": round(statistics.median(values), 3),
            "min": round(min(values), 3),
            "max": round(max(values), 3),
        }

    return {
        "label": label,
        "prompt_chars": len(prompt),
        "stats": {
            "ttft_s": stats(ttfts),
            "tokens_per_sec": stats(tpss),
            "total_time_s": stats(total_times),
        },
        "runs": results,
        "error": None,
    }


def run_ability_test(client, model_id, test_def):
    """Run a single ability test."""
    messages = [
        {"role": "system", "content": test_def["system"]},
        {"role": "user", "content": test_def["prompt"]},
    ]

    print(f"    Running: {test_def['name']}...", end=" ", flush=True)
    resp = stream_chat(
        client, model_id, messages,
        max_tokens=test_def.get("max_tokens", 256),
        temperature=test_def.get("temperature", 0),
    )

    if resp["error"]:
        print(f"ERROR: {resp['error']}")
        return {
            "test": test_def["name"],
            "passed": False,
            "error": resp["error"],
            "response": "",
            "timing": {"total_time_s": resp["total_time_s"], "ttft_s": resp["ttft_s"]},
        }

    # Run the checker on effective response (content or reasoning)
    effective = get_effective_response(resp)
    result = test_def["checker"](effective)
    status = "PASS" if result.get("passed") else "FAIL"
    print(status)

    return {
        "test": test_def["name"],
        "passed": result.get("passed", False),
        "error": result.get("error"),
        "details": result.get("details", []),
        "response": effective,
        "reasoning_content": resp["reasoning_content"],
        "is_reasoning_model": resp.get("is_reasoning_model", False),
        "timing": {
            "total_time_s": round(resp["total_time_s"], 3) if resp["total_time_s"] else None,
            "ttft_s": round(resp["ttft_s"], 3) if resp["ttft_s"] else None,
        },
        "usage": resp.get("usage"),
    }


def run_agent_test(client, model_id, test_def):
    """Run a single agent fitness test."""
    messages = [
        {"role": "system", "content": test_def["system"]},
        {"role": "user", "content": test_def["prompt"]},
    ]

    print(f"    Running: {test_def['name']}...", end=" ", flush=True)
    resp = stream_chat(
        client, model_id, messages,
        max_tokens=test_def.get("max_tokens", 768),
        temperature=test_def.get("temperature", 0),
    )

    if resp["error"]:
        print(f"ERROR: {resp['error']}")
        return {
            "test": test_def["name"],
            "passed": False,
            "error": resp["error"],
            "response": "",
        }

    effective = get_effective_response(resp)
    result = test_def["checker"](effective)
    status = "PASS" if result.get("passed") else "FAIL"
    print(status)

    return {
        "test": test_def["name"],
        "passed": result.get("passed", False),
        "error": result.get("error"),
        "details": result.get("details", []),
        "response": effective,
        "is_reasoning_model": resp.get("is_reasoning_model", False),
        "metrics": {k: v for k, v in result.items() if k not in ("passed", "error", "details")},
        "timing": {
            "total_time_s": round(resp["total_time_s"], 3) if resp["total_time_s"] else None,
            "ttft_s": round(resp["ttft_s"], 3) if resp["ttft_s"] else None,
        },
        "usage": resp.get("usage"),
    }


# ──────────────────────────────────────────────────────────────────────────────
# MAIN BENCHMARK RUNNER
# ──────────────────────────────────────────────────────────────────────────────

def run_benchmark(model_id, speed_only=False, quick=False, ability_only=False, agent_only=False, json_only=False):
    """Run the full benchmark suite."""
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d")

    # Resolve model to GGUF file
    models_meta = fetch_lmstudio_models()
    gguf_files = find_gguf_files()
    gguf_info = resolve_model_to_gguf(model_id, models_meta, gguf_files)

    result = {
        "model_id": model_id,
        "timestamp": timestamp,
        "api_base": API_BASE,
        "benchmark_version": "2.0.0",
        "gguf_info": gguf_info,
        "speed": {},
        "ability": {},
        "agent_fitness": {},
        "summary": {},
    }

    # Print header
    if not json_only:
        print(f"\n{'='*60}")
        print(f"  bench-llm — Benchmarking: {model_id}")
        print(f"  GGUF Status:   {gguf_info['status']}")
        if gguf_info["gguf_path"]:
            print(f"  GGUF Path:     {gguf_info['gguf_path']}")
            print(f"  GGUF Size:     {gguf_info['gguf_size_human']}")
        print(f"  Arch:          {gguf_info['arch'] or 'N/A'}")
        print(f"  Quant:         {gguf_info['quant'] or 'N/A'}")
        print(f"{'='*60}\n")

    # If model has no backing GGUF and isn't loaded, fail fast
    if gguf_info["status"] == "UNAVAILABLE":
        result["error"] = f"Model '{model_id}' has no GGUF file on disk"
        result["summary"] = {"error": "GGUF not found on disk"}
        if json_only:
            print(json.dumps(result, indent=2, default=str))
        else:
            print(f"  ❌ Model '{model_id}' has no backing GGUF file on disk. Skipping.\n")
        return result
    if gguf_info["status"] == "UNKNOWN":
        result["error"] = f"Model '{model_id}' not found in LM Studio"
        result["summary"] = {"error": "Model not registered in LM Studio"}
        if json_only:
            print(json.dumps(result, indent=2, default=str))
        else:
            print(f"  ❓ Model '{model_id}' not found in LM Studio. Skipping.\n")
        return result

    # Unload/clean VRAM
    unload_models()

    client = get_client()

    # Verify model is accessible
    available = list_models()
    if model_id not in available:
        print(f"\n  ⚠  Model '{model_id}' not found in available models.")
        print(f"  Available: {', '.join(available)}")
        result["error"] = f"Model '{model_id}' not available"
        return result

    # ── WARMUP ──
    print("  🔥 Warming up model...", end=" ", flush=True)
    try:
        warmup_resp = stream_chat(
            client, model_id,
            [{"role": "user", "content": WARMUP_PROMPT}],
            max_tokens=16, temperature=0,
        )
        if warmup_resp["error"]:
            print(f"FAILED: {warmup_resp['error']}")
            result["error"] = f"Warmup failed: {warmup_resp['error']}"
            return result
        print(f"OK ({warmup_resp['ttft_s']:.2f}s TTFT)")
    except Exception as e:
        print(f"FAILED: {e}")
        result["error"] = f"Warmup failed: {e}"
        return result

    # ── SPEED TESTS ──
    if not ability_only and not agent_only:
        speed_runs = SPEED_RUNS if not quick else 1
        print("\n  ⚡ SPEED BENCHMARKS")
        print("  " + "-" * 40)
        
        all_speed_results = {}
        for key, prompt in SPEED_PROMPTS.items():
            print(f"    {key} ({len(prompt)} chars, {speed_runs} runs)...")

            speed_result = run_speed_test(client, model_id, prompt, key, runs=speed_runs)
            all_speed_results[key] = speed_result
            
            if speed_result.get("error"):
                print(f"      ERROR: {speed_result['error']}")
            else:
                stats = speed_result.get("stats", {})
                tps_stats = stats.get("tokens_per_sec", {})
                ttft_stats = stats.get("ttft_s", {})
                print(f"      TPS:  mean={tps_stats.get('mean', 'N/A')} median={tps_stats.get('median', 'N/A')}")
                print(f"      TTFT: mean={ttft_stats.get('mean', 'N/A')}s median={ttft_stats.get('median', 'N/A')}s")

        result["speed"] = all_speed_results

        # ── Speed sanity validation ──
        if not json_only:
            print("\n  🔍 Validating speed plausibility...")
        plausibility = validate_speed_plausibility(all_speed_results, gguf_info)
        result["plausibility"] = plausibility
        if not json_only:
            print(f"    Confidence: {plausibility['confidence']}")
            for w in plausibility.get("warnings", []):
                print(f"    ⚠  {w}")

    # ── ABILITY TESTS ──
    if not speed_only and not agent_only:
        print("\n  🧠 ABILITY TESTS")
        print("  " + "-" * 40)
        
        ability_results = {}
        for test_def in ABILITY_TESTS:
            test_result = run_ability_test(client, model_id, test_def)
            ability_results[test_def["id"]] = test_result
        
        passed = sum(1 for t in ability_results.values() if t.get("passed"))
        total = len(ability_results)
        print(f"\n    Ability Score: {passed}/{total}")

        result["ability"] = {
            "tests": ability_results,
            "score": f"{passed}/{total}",
            "pass_rate": round(passed / total * 100, 1) if total > 0 else 0,
        }

    # ── AGENT FITNESS TESTS ──
    if not speed_only and not ability_only:
        print("\n  🤖 AGENT FITNESS TESTS")
        print("  " + "-" * 40)

        agent_results = {}
        for test_def in AGENT_FITNESS_TESTS:
            test_result = run_agent_test(client, model_id, test_def)
            agent_results[test_def["id"]] = test_result
        
        passed = sum(1 for t in agent_results.values() if t.get("passed"))
        total = len(agent_results)
        print(f"\n    Agent Fitness Score: {passed}/{total}")

        result["agent_fitness"] = {
            "tests": agent_results,
            "score": f"{passed}/{total}",
            "pass_rate": round(passed / total * 100, 1) if total > 0 else 0,
        }

    # ── COMPILE SUMMARY ──
    result["summary"] = compile_summary(result)

    # ── SAVE RESULTS ──
    BENCHMARKS_DIR.mkdir(parents=True, exist_ok=True)
    safe_model_id = re.sub(r'[^\w\-.]', '_', model_id)
    output_path = BENCHMARKS_DIR / f"{safe_model_id}-{date_str}.json"

    with open(output_path, 'w') as f:
        json.dump(result, f, indent=2, default=str)

    if json_only:
        print(json.dumps(result, indent=2, default=str))
    else:
        print(f"\n  📄 Results saved: {output_path}")

        # ── PRINT SUMMARY ──
        print_summary(result)

    return result


def compile_summary(result):
    """Create a summary from benchmark results."""
    summary = {}

    # GGUF info
    if result.get("gguf_info"):
        summary["gguf_file_size"] = result["gguf_info"].get("gguf_size_human", "N/A")
        summary["arch"] = result["gguf_info"].get("arch", "N/A")
        summary["quant"] = result["gguf_info"].get("quant", "N/A")

    # Plausibility
    if result.get("plausibility"):
        summary["confidence"] = result["plausibility"].get("confidence", "N/A")

    # Speed summary
    if result.get("speed"):
        speed_summary = {}
        for key, data in result["speed"].items():
            if data.get("error"):
                speed_summary[key] = {"error": data["error"]}
            else:
                stats = data.get("stats", {})
                speed_summary[key] = {
                    "tps_mean": stats.get("tokens_per_sec", {}).get("mean"),
                    "tps_median": stats.get("tokens_per_sec", {}).get("median"),
                    "ttft_mean_s": stats.get("ttft_s", {}).get("mean"),
                    "ttft_median_s": stats.get("ttft_s", {}).get("median"),
                }
        summary["speed"] = speed_summary

    # Ability summary
    if result.get("ability"):
        summary["ability"] = result["ability"]["score"]
        summary["ability_score"] = result["ability"].get("pass_rate", 0)

    # Agent fitness summary
    if result.get("agent_fitness"):
        summary["agent_fitness"] = result["agent_fitness"]["score"]

    # Overall
    if result.get("speed") and result.get("ability"):
        # Use short prompt TPS as the headline number
        short_stats = result["speed"].get("short_50", {}).get("stats", {})
        tps = short_stats.get("tokens_per_sec", {}).get("mean", 0)
        ttft = short_stats.get("ttft_s", {}).get("mean", 0)
        ability_pct = result["ability"].get("pass_rate", 0) if result.get("ability") else 0
        agent_pct = result["agent_fitness"].get("pass_rate", 0) if result.get("agent_fitness") else 0
        summary["headline"] = f"{tps} T/s, {ttft}s TTFT, Ab:{ability_pct}%, Ag:{agent_pct}%"

    return summary


def print_summary(result):
    """Print a formatted terminal summary."""
    print(f"\n{'='*80}")
    print(f"  BENCHMARK RESULTS: {result['model_id']}")
    print(f"{'='*80}")

    # Enhanced summary table
    gguf_info = result.get("gguf_info", {})
    print(f"\n  📊 SUMMARY TABLE")
    print(f"  {'─'*70}")
    print(f"  {'Model ID':<35} {result['model_id'][:35]}")
    print(f"  {'GGUF Size':<35} {gguf_info.get('gguf_size_human', 'N/A')}")
    print(f"  {'Architecture':<35} {gguf_info.get('arch', 'N/A')}")
    print(f"  {'Quantization':<35} {gguf_info.get('quant', 'N/A')}")

    # Summary stats from short prompt
    if result.get("speed"):
        short_stats = result["speed"].get("short_50", {}).get("stats", {})
        tps_mean = short_stats.get("tokens_per_sec", {}).get("mean")
        ttft_mean = short_stats.get("ttft_s", {}).get("mean")
        print(f"  {'T/s (short, mean)':<35} {f'{tps_mean:.1f}' if tps_mean else 'N/A'}")
        print(f"  {'TTFT (short, mean)':<35} {f'{ttft_mean:.3f}s' if ttft_mean else 'N/A'}")

    if result.get("ability"):
        print(f"  {'Ability Score':<35} {result['ability']['score']} ({result['ability']['pass_rate']}%)")

    if result.get("agent_fitness"):
        print(f"  {'Agent Fitness Score':<35} {result['agent_fitness']['score']} ({result['agent_fitness']['pass_rate']}%)")

    if result.get("plausibility"):
        p = result["plausibility"]
        conf_icon = {"HIGH": "🟢", "MEDIUM": "🟡", "LOW": "🔴"}
        print(f"  {'Confidence':<35} {conf_icon.get(p.get('confidence',''),'')} {p.get('confidence','N/A')}")
        for w in p.get("warnings", []):
            print(f"  {'':<35} ⚠ {w[:60]}")

    print(f"  {'─'*70}")

    # Speed
    if result.get("speed"):
        print(f"\n  ⚡ SPEED")
        print(f"  {'─'*50}")
        header = f"  {'Test':<15} {'TPS Mean':>10} {'TPS Med':>10} {'TTFT Mean':>10} {'TTFT Med':>10}"
        print(header)
        print(f"  {'─'*50}")
        for key in ["short_50", "medium_200", "long_800"]:
            data = result["speed"].get(key)
            if not data:
                continue
            if data.get("error"):
                print(f"  {key:<15} {'ERROR':>10}")
                continue
            stats = data["stats"]
            tps = stats.get("tokens_per_sec", {})
            ttft = stats.get("ttft_s", {})
            print(f"  {key:<15} {tps.get('mean','-'):>8.1f}  {tps.get('median','-'):>8.1f}  {ttft.get('mean','-'):>8.2f}s {ttft.get('median','-'):>8.2f}s")

    # Ability
    if result.get("ability"):
        print(f"\n  🧠 ABILITY: {result['ability']['score']} ({result['ability']['pass_rate']}%)")
        print(f"  {'─'*50}")
        for test_id, test in result["ability"]["tests"].items():
            status = "✅" if test.get("passed") else "❌"
            t = test.get("timing", {})
            extra = ""
            if not test.get("passed") and test.get("details"):
                # Show first FAIL detail for failed tests
                fails = [d for d in test["details"] if "FAIL" in str(d)]
                if fails:
                    extra = f" — {fails[0].strip()}"
            elif test.get("passed") and test.get("details"):
                # Show score for passed tests with partial failures
                fails = [d for d in test["details"] if "FAIL" in str(d)]
                passes = [d for d in test["details"] if "PASS" in str(d)]
                if fails:
                    extra = f" ({len(passes)}/{len(passes)+len(fails)} checks)"
            err = f" — {test['error']}" if test.get("error") else ""
            print(f"  {status} {test['test']}{extra}{err}")

    # Agent Fitness
    if result.get("agent_fitness"):
        print(f"\n  🤖 AGENT FITNESS: {result['agent_fitness']['score']} ({result['agent_fitness']['pass_rate']}%)")
        print(f"  {'─'*50}")
        for test_id, test in result["agent_fitness"]["tests"].items():
            status = "✅" if test.get("passed") else "❌"
            extra = ""
            if test.get("details"):
                fails = [d for d in test["details"] if "FAIL" in str(d)]
                if fails:
                    extra = f" — {fails[0].strip()}"
            # Show conciseness rating
            if test_id == "conciseness" and test.get("metrics", {}).get("rating"):
                extra = f" — Rating: {test['metrics']['rating']} ({test['metrics']['ratio']}x)"
            err = f" — {test['error']}" if test.get("error") else ""
            print(f"  {status} {test['test']}{extra}{err}")

    print(f"\n{'='*80}\n")


# ──────────────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────────────

def main():
    global API_BASE

    parser = argparse.ArgumentParser(
        description="bench-llm — Local LLM Benchmarking Tool for LM Studio",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\
            Examples:
              %(prog)s qwen3.6-27b-fable-fusion      # Full benchmark
              %(prog)s gemma-4-31b-it --speed-only    # Speed only
              %(prog)s nemotron-nano-4b --quick       # Fast smoke test
              %(prog)s --list                         # List models
              %(prog)s --check <model-id>             # Pre-flight check
              %(prog)s <model-id> --json              # Machine-readable JSON output
        """),
    )
    parser.add_argument("model", nargs="?", help="Model ID to benchmark/check")
    parser.add_argument("--list", action="store_true", help="List available models and exit")
    parser.add_argument("--check", action="store_true", help="Pre-flight check: resolve model file and report status")
    parser.add_argument("--speed-only", action="store_true", help="Run only speed tests")
    parser.add_argument("--ability-only", action="store_true", help="Run only ability tests")
    parser.add_argument("--agent-only", action="store_true", help="Run only agent fitness tests")
    parser.add_argument("--quick", action="store_true", help="Quick smoke test (1 run per speed test)")
    parser.add_argument("--json", action="store_true", help="Output only machine-readable JSON")
    parser.add_argument("--api-base", default=API_BASE, help=f"API base URL (default: {API_BASE})")

    args = parser.parse_args()

    # Global API_BASE override
    if args.api_base != API_BASE:
        API_BASE = args.api_base

    # List models
    if args.list:
        models = list_models()
        print(f"\nModels available at {API_BASE}:")
        for m in sorted(models):
            print(f"  • {m}")
        print(f"\n  Total: {len(models)} model(s)")
        return

    # Pre-flight check
    if args.check:
        if not args.model:
            print("ERROR: --check requires a model ID")
            sys.exit(1)
        preflight_check(args.model, output_json=args.json)
        return

    # Model required for benchmarking
    if not args.model:
        parser.print_help()
        print("\nERROR: model ID required for benchmarking (or use --list)")
        sys.exit(1)

    model_id = args.model

    # Run benchmark
    try:
        run_benchmark(
            model_id,
            speed_only=args.speed_only,
            quick=args.quick,
            ability_only=args.ability_only,
            agent_only=args.agent_only,
            json_only=args.json,
        )
    except KeyboardInterrupt:
        print("\n\n  ⚠  Benchmark interrupted by user.")
        sys.exit(130)
    except Exception as e:
        print(f"\n  💥 FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
