"""Production control-plane integration for an OpenAI-compatible vLLM server.

vLLM supports per-request LoRA selection, but its public server API does not
expose arbitrary residual-stream hooks. This backend is explicit about that
boundary instead of pretending activation cards execute.
"""

from __future__ import annotations

import asyncio
import time

import httpx

from ..contracts import BackendCapabilities
from ..errors import UnsupportedCapability
from .base import Backend, BackendInput, BackendResult


class VLLMBackend(Backend):
    def __init__(
        self,
        base_url: str,
        model: str,
        *,
        api_key: str | None = None,
        timeout_seconds: float = 60.0,
    ):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key
        self.timeout_seconds = timeout_seconds
        self._capabilities = BackendCapabilities(
            backend="vllm",
            model=model,
            batching=True,
            kv_cache=True,
            activation_cards=False,
            per_request_adapters=True,
            metrics=("latency_ms", "input_tokens", "output_tokens"),
            limitations=(
                "Residual cards require a future vLLM worker plugin.",
                "Current integration governs prompts, outputs, tenant state, and LoRA routing.",
            ),
        )

    @property
    def capabilities(self) -> BackendCapabilities:
        return self._capabilities

    async def generate_batch(self, requests: list[BackendInput]) -> list[BackendResult]:
        if any(
            spec.get("kind") == "steering"
            for request in requests
            for spec in request.activation_specs
        ):
            raise UnsupportedCapability(
                "vLLM HTTP mode cannot execute residual-stream cards"
            )
        if any(
            spec.get("kind") not in {"adapter"}
            for request in requests
            for spec in request.activation_specs
        ):
            raise UnsupportedCapability("vLLM received an unknown activation card")
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        async with httpx.AsyncClient(
            base_url=self.base_url,
            headers=headers,
            timeout=self.timeout_seconds,
        ) as client:
            return list(
                await asyncio.gather(
                    *[self._one(client, request) for request in requests]
                )
            )

    async def _one(
        self, client: httpx.AsyncClient, request: BackendInput
    ) -> BackendResult:
        adapters = [
            spec for spec in request.activation_specs if spec.get("kind") == "adapter"
        ]
        if len(adapters) > 1:
            raise UnsupportedCapability(
                "vLLM HTTP mode allows at most one adapter card per request"
            )
        if adapters:
            selected_model = str(adapters[0].get("served_model", "")).strip()
            if not selected_model:
                raise UnsupportedCapability(
                    "vLLM adapter cards require a signed served_model"
                )
        else:
            selected_model = self.model
        started = time.perf_counter()
        response = await client.post(
            "/v1/chat/completions",
            json={
                "model": selected_model,
                "messages": [{"role": "user", "content": request.prompt}],
                "max_tokens": request.max_new_tokens,
                "temperature": 0,
            },
        )
        response.raise_for_status()
        body = response.json()
        usage = body.get("usage") or {}
        return BackendResult(
            text=body["choices"][0]["message"]["content"],
            input_tokens=int(usage.get("prompt_tokens", 0)),
            output_tokens=int(usage.get("completion_tokens", 0)),
            latency_ms=round((time.perf_counter() - started) * 1000, 3),
            metadata={"served_model": selected_model},
        )

