"""Multimodal embedder backed by a vLLM OpenAI-compatible HTTP server.

vLLM's ``/v1/embeddings`` endpoint is a superset of OpenAI's:

- Text inputs use the standard ``input: list[str]`` field — one HTTP call
  returns N embeddings.
- Image inputs use a ``messages`` array carrying an ``image_url`` content
  part with a base64 data URI. One image per HTTP call.

Models like ``Qwen/Qwen3-VL-Embedding-8B`` and ``jinaai/jina-embeddings-v4``
ship with chat templates that map both shapes into a shared vector space.
"""

import base64
import io
from typing import TYPE_CHECKING, Any

import httpx

from haiku.rag.embeddings import EmbedderWrapper

if TYPE_CHECKING:
    from PIL import Image as PILImage


class VLLMMultimodalEmbedder(EmbedderWrapper):
    supports_images = True

    def __init__(
        self,
        model_name: str,
        vector_dim: int,
        base_url: str,
        api_key: str | None = None,
        timeout: float = 60.0,
    ):
        super().__init__(embedder=None, vector_dim=vector_dim)
        self._model_name = model_name
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._timeout = timeout

    def _headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        return headers

    async def _post(self, body: dict[str, Any]) -> list[list[float]]:
        try:
            async with httpx.AsyncClient(timeout=self._timeout) as client:
                response = await client.post(
                    f"{self._base_url}/embeddings",
                    json=body,
                    headers=self._headers(),
                )
                response.raise_for_status()
                payload = response.json()
        except httpx.ConnectError as e:
            raise ValueError(
                f"Could not connect to vLLM at {self._base_url}. "
                f"Ensure the service is running. Error: {e}"
            ) from e
        except httpx.TimeoutException as e:
            raise ValueError(
                f"Request to vLLM timed out after {self._timeout}s. Error: {e}"
            ) from e
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                raise ValueError(
                    "Authentication failed against vLLM. Check the API key."
                ) from e
            raise ValueError(f"HTTP error from vLLM: {e}") from e

        data = payload.get("data") or []
        if not data:
            raise ValueError(f"vLLM returned no embeddings: {payload}")
        return [list(d["embedding"]) for d in data]

    async def embed_query(self, text: str) -> list[float]:
        rows = await self._post(
            {
                "model": self._model_name,
                "input": [text],
                "encoding_format": "float",
            }
        )
        return rows[0]

    async def embed_documents(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        return await self._post(
            {
                "model": self._model_name,
                "input": texts,
                "encoding_format": "float",
            }
        )

    async def embed_image(self, image: "bytes | PILImage.Image") -> list[float]:
        rows = await self._post(
            {
                "model": self._model_name,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": _to_data_uri(image)},
                            }
                        ],
                    }
                ],
                "encoding_format": "float",
            }
        )
        return rows[0]


def _to_data_uri(image: "bytes | PILImage.Image") -> str:
    """Render an image as a ``data:image/png;base64,...`` URI."""
    if isinstance(image, bytes):
        return f"data:image/png;base64,{base64.b64encode(image).decode('ascii')}"

    from PIL import Image as PILImageModule

    if isinstance(image, PILImageModule.Image):
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        return (
            f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}"
        )

    raise TypeError(f"Unsupported image type: {type(image)!r}")