"""vLLM provider — local GPU inference via OpenAI-compatible API."""

from __future__ import annotations

import json
import time
import urllib.error
import urllib.request

from orxaq.orchestration.providers.base import LLMProvider, LLMResponse


class VLLMProvider(LLMProvider):
    """Provider for local vLLM server (OpenAI-compatible API)."""

    name = "vllm"

    def __init__(
        self,
        endpoint: str = "http://localhost:8000/v1",
        models: list[str] | None = None,
    ) -> None:
        self.endpoint = endpoint.rstrip("/")
        self.models = models or ["default"]

    def generate(
        self,
        prompt: str,
        model: str | None = None,
        max_tokens: int = 1024,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Generate via vLLM's OpenAI-compatible /chat/completions endpoint."""
        use_model = model or self.models[0]
        url = f"{self.endpoint}/chat/completions"

        payload = {
            "model": use_model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature,
        }

        start = time.monotonic()
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=data,
            headers={"Content-Type": "application/json"},
        )

        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                result = json.loads(resp.read())
        except (urllib.error.URLError, OSError) as e:
            return LLMResponse(
                text="",
                model=use_model,
                provider=self.name,
                metadata={"error": str(e)},
            )

        elapsed = (time.monotonic() - start) * 1000

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

        return LLMResponse(
            text=text,
            model=use_model,
            provider=self.name,
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
            latency_ms=elapsed,
        )

    def health_check(self) -> dict:
        """Check vLLM server availability via /models endpoint."""
        url = f"{self.endpoint}/models"
        try:
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=5) as resp:
                result = json.loads(resp.read())
            model_ids = [m.get("id", "") for m in result.get("data", [])]
            return {
                "healthy": True,
                "provider": self.name,
                "models": model_ids or self.models,
                "endpoint": self.endpoint,
            }
        except (urllib.error.URLError, OSError) as e:
            return {
                "healthy": False,
                "provider": self.name,
                "models": self.models,
                "endpoint": self.endpoint,
                "error": str(e),
            }
