#!/usr/bin/env python3
"""LLMHosts E2E hardware verification script.

Runs automated health, latency, and throughput checks against a running
LLMHosts instance. Designed to be run on actual GPU machines.

Usage:
    python verify.py [--host HOST] [--port PORT] [--model MODEL]
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from urllib.request import Request, urlopen


def check_health(base_url: str) -> dict:
    """Check /health and /ready endpoints."""
    results = {"health": False, "ready": False, "version": "unknown"}

    try:
        resp = urlopen(f"{base_url}/health", timeout=10)
        data = json.loads(resp.read())
        results["health"] = resp.status == 200
        results["version"] = data.get("version", "unknown")
    except Exception as e:
        results["health_error"] = str(e)

    try:
        resp = urlopen(f"{base_url}/ready", timeout=10)
        results["ready"] = resp.status == 200
    except Exception as e:
        results["ready_error"] = str(e)

    return results


def check_models(base_url: str) -> list[str]:
    """List available models via /v1/models."""
    try:
        resp = urlopen(f"{base_url}/v1/models", timeout=10)
        data = json.loads(resp.read())
        return [m["id"] for m in data.get("data", [])]
    except Exception:
        return []


def check_inference(base_url: str, model: str, prompt: str = "Say hello in exactly 5 words.") -> dict:
    """Run a single inference request and measure latency."""
    payload = json.dumps(
        {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 50,
        }
    ).encode()

    req = Request(
        f"{base_url}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json"},
    )

    start = time.monotonic()
    try:
        resp = urlopen(req, timeout=120)
        data = json.loads(resp.read())
        elapsed = time.monotonic() - start

        content = data["choices"][0]["message"]["content"]
        tokens = data.get("usage", {})

        return {
            "success": True,
            "latency_ms": round(elapsed * 1000),
            "content": content[:100],
            "prompt_tokens": tokens.get("prompt_tokens", 0),
            "completion_tokens": tokens.get("completion_tokens", 0),
            "model": data.get("model", model),
        }
    except Exception as e:
        return {"success": False, "error": str(e), "latency_ms": round((time.monotonic() - start) * 1000)}


def check_throughput(base_url: str, model: str, num_requests: int = 5) -> dict:
    """Run multiple requests and measure throughput statistics."""
    latencies = []
    successes = 0

    for i in range(num_requests):
        result = check_inference(base_url, model, f"Count from 1 to {i + 3}.")
        if result["success"]:
            latencies.append(result["latency_ms"])
            successes += 1

    if not latencies:
        return {"success": False, "error": "All requests failed"}

    return {
        "success": True,
        "total_requests": num_requests,
        "successful": successes,
        "failed": num_requests - successes,
        "latency_p50_ms": round(statistics.median(latencies)),
        "latency_p95_ms": round(sorted(latencies)[int(len(latencies) * 0.95)] if len(latencies) > 1 else latencies[0]),
        "latency_mean_ms": round(statistics.mean(latencies)),
        "latency_min_ms": min(latencies),
        "latency_max_ms": max(latencies),
    }


def check_cache(base_url: str, model: str) -> dict:
    """Verify cache by sending identical requests and checking for speed improvement."""
    prompt = "What is 2+2? Answer with just the number."

    # First request (cold)
    r1 = check_inference(base_url, model, prompt)
    if not r1["success"]:
        return {"success": False, "error": "First request failed"}

    # Second identical request (should hit cache)
    r2 = check_inference(base_url, model, prompt)
    if not r2["success"]:
        return {"success": False, "error": "Second request failed"}

    speedup = r1["latency_ms"] / max(r2["latency_ms"], 1)

    return {
        "success": True,
        "cold_latency_ms": r1["latency_ms"],
        "warm_latency_ms": r2["latency_ms"],
        "speedup": round(speedup, 2),
        "cache_likely_hit": speedup > 1.5,
    }


def run_all_checks(host: str, port: int, model: str | None) -> dict:
    """Run the complete verification battery."""
    base_url = f"http://{host}:{port}"
    results: dict = {"base_url": base_url, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ")}

    print(f"\n{'=' * 60}")
    print("LLMHosts E2E Hardware Verification")
    print(f"Target: {base_url}")
    print(f"{'=' * 60}\n")

    # 1. Health
    print("[1/5] Health check...", end=" ", flush=True)
    health = check_health(base_url)
    results["health"] = health
    print("OK" if health["health"] else "FAIL")

    # 2. Models
    print("[2/5] Model discovery...", end=" ", flush=True)
    models = check_models(base_url)
    results["models"] = models
    print(f"{len(models)} models" if models else "FAIL")

    # Pick model
    if model:
        test_model = model
    elif models:
        test_model = models[0]
    else:
        print("\nERROR: No models available and no --model specified")
        results["error"] = "No models available"
        return results

    print(f"    Using model: {test_model}")

    # 3. Single inference
    print("[3/5] Inference test...", end=" ", flush=True)
    inference = check_inference(base_url, test_model)
    results["inference"] = inference
    print(f"{inference['latency_ms']}ms" if inference["success"] else "FAIL")

    # 4. Throughput
    print("[4/5] Throughput test (5 requests)...", end=" ", flush=True)
    throughput = check_throughput(base_url, test_model, 5)
    results["throughput"] = throughput
    if throughput["success"]:
        print(f"p50={throughput['latency_p50_ms']}ms, p95={throughput['latency_p95_ms']}ms")
    else:
        print("FAIL")

    # 5. Cache
    print("[5/5] Cache test...", end=" ", flush=True)
    cache = check_cache(base_url, test_model)
    results["cache"] = cache
    if cache["success"]:
        hit = "HIT" if cache["cache_likely_hit"] else "MISS"
        print(f"{cache['warm_latency_ms']}ms (cache {hit}, {cache['speedup']}x speedup)")
    else:
        print("FAIL")

    # Summary
    print(f"\n{'=' * 60}")
    passed = sum(1 for k in ["health", "inference", "throughput", "cache"] if results.get(k, {}).get("success"))
    total = 4
    print(f"Results: {passed}/{total} checks passed")
    if health.get("version"):
        print(f"Version: {health['version']}")
    print(f"{'=' * 60}\n")

    return results


def main() -> None:
    parser = argparse.ArgumentParser(description="LLMHosts E2E Hardware Verification")
    parser.add_argument("--host", default="127.0.0.1", help="LLMHosts host (default: 127.0.0.1)")
    parser.add_argument("--port", type=int, default=4000, help="LLMHosts port (default: 4000)")
    parser.add_argument("--model", default=None, help="Model to test (default: first available)")
    parser.add_argument("--json", action="store_true", help="Output results as JSON")
    args = parser.parse_args()

    results = run_all_checks(args.host, args.port, args.model)

    if args.json:
        print(json.dumps(results, indent=2))

    # Exit code: 0 if health + inference passed
    health_ok = results.get("health", {}).get("health", False)
    inference_ok = results.get("inference", {}).get("success", False)
    sys.exit(0 if health_ok and inference_ok else 1)


if __name__ == "__main__":
    main()
