"""vLLM instance discovery and management.

Discovers running vLLM instances via their OpenAI-compatible API.
Default: http://localhost:8000
"""

from __future__ import annotations

import json
import logging
import re
import time

import httpx
from pydantic import BaseModel

from llmhosts.constants import VLLM_DEFAULT_HOST

logger = logging.getLogger(__name__)


class VLLMModel(BaseModel):
    """A model served by vLLM (OpenAI /v1/models response)."""

    id: str
    object: str = "model"
    owned_by: str = "vllm"
    max_model_len: int | None = None


class VLLMMetrics(BaseModel):
    """vLLM server metrics (throughput, queue depth, etc.)."""

    throughput_tps: float = 0.0  # tokens per second
    avg_queue_depth: float = 0.0
    active_requests: int = 0
    gpu_memory_used_gb: float = 0.0
    gpu_memory_total_gb: float = 0.0


class VLLMHealth(BaseModel):
    """vLLM health status."""

    healthy: bool
    version: str = ""
    uptime_seconds: float = 0.0
    models_loaded: int = 0


class VLLMDiscovery:
    """Discover and interact with vLLM instances.

    vLLM exposes an OpenAI-compatible API at /v1/models, /v1/chat/completions, etc.
    """

    def __init__(self, host: str = VLLM_DEFAULT_HOST, timeout: float = 5.0) -> None:
        self._host = host.rstrip("/")
        self._timeout = timeout
        self._client = httpx.AsyncClient(base_url=self._host, timeout=httpx.Timeout(timeout))

    async def __aenter__(self) -> VLLMDiscovery:
        return self

    async def __aexit__(self, exc_type: type | None, exc_val: BaseException | None, exc_tb: object) -> None:
        await self.close()

    async def is_available(self) -> bool:
        """Check if a vLLM instance is running."""
        try:
            resp = await self._client.get("/v1/models")
            return resp.status_code == 200
        except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPError) as exc:
            logger.debug("vLLM not reachable at %s: %s", self._host, exc)
            return False

    async def list_models(self) -> list[VLLMModel]:
        """List models served by vLLM."""
        try:
            resp = await self._client.get("/v1/models")
            resp.raise_for_status()
            data = resp.json()
            models_data = data.get("data", [])
            result: list[VLLMModel] = []
            for m in models_data:
                if isinstance(m, dict):
                    model_id = m.get("id", "")
                    max_len = m.get("root", {}).get("max_model_len") if isinstance(m.get("root"), dict) else None
                    result.append(
                        VLLMModel(
                            id=model_id,
                            object=m.get("object", "model"),
                            owned_by=m.get("owned_by", "vllm"),
                            max_model_len=max_len,
                        )
                    )
                else:
                    result.append(VLLMModel(id=str(m)))
            return result
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            logger.error("Failed to list vLLM models: %s", exc)
            return []

    async def get_metrics(self) -> VLLMMetrics:
        """Get vLLM server metrics (throughput, queue depth, etc.)."""
        try:
            resp = await self._client.get("/metrics")
            if resp.status_code != 200:
                return VLLMMetrics()

            text = resp.text
            metrics = VLLMMetrics()

            # Parse Prometheus-format metrics
            # vllm:num_requests_running{...} 0.0
            running_match = re.search(r"vllm:num_requests_running[^{}]*(?:\{[^}]*\})?\s+([\d.]+)", text)
            if running_match:
                metrics.active_requests = int(float(running_match.group(1)))

            waiting_match = re.search(r"vllm:num_requests_waiting[^{}]*(?:\{[^}]*\})?\s+([\d.]+)", text)
            if waiting_match:
                metrics.avg_queue_depth = float(waiting_match.group(1))

            # throughput_tps: vLLM doesn't expose directly; would need time-delta on generation_tokens
            # Default 0.0; compare_throughput does actual benchmarking
            metrics.throughput_tps = 0.0

            return metrics
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            logger.debug("Failed to fetch vLLM metrics: %s", exc)
            return VLLMMetrics()

    async def health_check(self) -> VLLMHealth:
        """Check vLLM health status."""
        available = await self.is_available()
        models_count = 0
        version = ""
        uptime = 0.0

        if available:
            models = await self.list_models()
            models_count = len(models)

            # Try /version endpoint (some vLLM versions expose it)
            try:
                ver_resp = await self._client.get("/version")
                if ver_resp.status_code == 200:
                    version = ver_resp.text.strip()
            except (httpx.HTTPError, httpx.TimeoutException):
                pass

        return VLLMHealth(
            healthy=available,
            version=version,
            uptime_seconds=uptime,
            models_loaded=models_count,
        )

    async def chat_for_benchmark(self, model: str, messages: list[dict[str, str]]) -> tuple[int, float]:
        """Run a minimal chat completion for throughput benchmarking. Returns (tokens, elapsed_sec)."""
        payload = {
            "model": model,
            "messages": messages,
            "stream": True,
            "max_tokens": 50,
        }
        start = time.perf_counter()
        token_count = 0
        try:
            async with self._client.stream("POST", "/v1/chat/completions", json=payload) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if line.startswith("data: ") and line != "data: [DONE]":
                        chunk = line[6:]
                        if chunk.strip():
                            try:
                                obj = json.loads(chunk)
                                for choice in obj.get("choices", []):
                                    delta = choice.get("delta", {})
                                    if delta.get("content"):
                                        token_count += 1
                            except json.JSONDecodeError:
                                pass
        except (httpx.HTTPError, httpx.TimeoutException):
            pass
        elapsed = time.perf_counter() - start
        return token_count, max(elapsed, 0.001)

    async def close(self) -> None:
        """Close the underlying httpx client."""
        await self._client.aclose()

    def __repr__(self) -> str:
        return f"VLLMDiscovery(host={self._host!r})"
