"""
Multi-backend embedding client.

Dispatches embedding requests to OpenAI-compatible APIs (OpenAI, vLLM,
llama-cpp, Ollama) or Amazon Bedrock via direct HTTP with Bearer token auth.
"""

import base64
import json
import logging
import time
from typing import Optional

import requests

from library.metrics import (
    EMBEDDING_API_ERRORS_TOTAL,
    EMBEDDING_BATCH_DURATION,
    EMBEDDING_TOKENS_TOTAL,
    EMBEDDINGS_GENERATED_TOTAL,
)

logger = logging.getLogger(__name__)


class EmbeddingClient:
    """
    Client for generating text and image embeddings via multiple backends.

    Dispatches based on ``LLMApi.api_type``:

    * ``openai``, ``vllm``, ``llama-cpp``, ``ollama`` — OpenAI-compatible
      ``POST /embeddings``
    * ``bedrock`` — Amazon Bedrock Runtime ``POST /model/{id}/invoke``
      with Bearer token auth
    """

    def __init__(self, embedding_model, user=None):
        """
        :param embedding_model: ``LLMModel`` instance for embeddings.
        :param user: Optional Django user for usage tracking.
        """
        self.model = embedding_model
        self.api = embedding_model.api
        self.user = user
        self.base_url = self.api.base_url.rstrip("/")
        self.model_name = self.model.name
        self.api_type = self.api.api_type
        self.timeout = self.api.timeout_seconds or 120

        logger.info(
            "EmbeddingClient initialized api=%s model=%s api_type=%s base_url=%s",
            self.api.name,
            self.model_name,
            self.api_type,
            self.base_url,
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def embed_text(self, text: str) -> list[float]:
        """
        Generate an embedding vector for a single text string.

        :param text: Text to embed.
        :returns: Embedding vector as list of floats.
        """
        if self.api_type == "bedrock":
            return self._embed_bedrock_single(text)
        return self._embed_openai_single(text)

    def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """
        Generate embeddings for multiple texts.

        :param texts: List of text strings.
        :returns: List of embedding vectors.
        """
        if self.api_type == "bedrock":
            return self._embed_bedrock_batch(texts)
        return self._embed_openai_batch(texts)

    def embed_image(self, image_data: bytes, image_ext: str = "png") -> Optional[list[float]]:
        """
        Generate a multimodal embedding for an image.

        Requires a model with ``supports_multimodal=True``.

        :param image_data: Raw image bytes.
        :param image_ext: Image format extension.
        :returns: Embedding vector, or None if not supported.
        """
        if not self.model.supports_multimodal:
            logger.debug(
                "Model %s does not support multimodal, skipping image embedding",
                self.model_name,
            )
            return None

        b64 = base64.b64encode(image_data).decode("utf-8")
        mime_type = f"image/{image_ext}" if image_ext != "jpg" else "image/jpeg"

        if self.api_type == "bedrock":
            return self._embed_bedrock_image(b64, mime_type)
        return self._embed_openai_image(b64, mime_type)

    # ------------------------------------------------------------------
    # OpenAI-compatible backend
    # ------------------------------------------------------------------

    def _embed_openai_single(self, text: str) -> list[float]:
        """Embed a single text via OpenAI-compatible /embeddings endpoint."""
        result = self._embed_openai_batch([text])
        return result[0]

    def _embed_openai_batch(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts via OpenAI-compatible /embeddings endpoint."""
        url = f"{self.base_url}/embeddings"
        payload = {"input": texts, "model": self.model_name}
        headers = self._openai_headers()

        logger.debug(
            "OpenAI embedding request texts=%d model=%s",
            len(texts),
            self.model_name,
        )

        with EMBEDDING_BATCH_DURATION.labels(
            model_name=self.model_name, api_type=self.api_type
        ).time():
            try:
                resp = requests.post(
                    url,
                    json=payload,
                    headers=headers,
                    timeout=self.timeout * max(1, len(texts) // 10),
                )
                if resp.status_code != 200:
                    logger.error(
                        "OpenAI embedding failed status=%d body=%s",
                        resp.status_code,
                        resp.text[:500],
                    )
                resp.raise_for_status()
                data = resp.json()
            except requests.RequestException as exc:
                EMBEDDING_API_ERRORS_TOTAL.labels(
                    model_name=self.model_name,
                    api_type=self.api_type,
                    error_type=type(exc).__name__,
                ).inc()
                logger.error("OpenAI embedding request failed: %s", exc)
                raise

        embeddings = self._parse_openai_response(data)

        # Metrics
        EMBEDDINGS_GENERATED_TOTAL.labels(
            model_name=self.model_name,
            api_type=self.api_type,
            content_type="text",
        ).inc(len(embeddings))
        EMBEDDING_TOKENS_TOTAL.labels(model_name=self.model_name).inc(
            sum(len(t.split()) for t in texts)
        )

        self._log_usage(len(texts), sum(len(t.split()) for t in texts))

        logger.debug(
            "OpenAI embedding response texts=%d dimensions=%d",
            len(embeddings),
            len(embeddings[0]) if embeddings else 0,
        )
        return embeddings

    def _embed_openai_image(self, b64_image: str, mime_type: str) -> Optional[list[float]]:
        """Embed an image via OpenAI-compatible multimodal endpoint."""
        url = f"{self.base_url}/embeddings"
        payload = {
            "input": [{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}}],
            "model": self.model_name,
        }
        headers = self._openai_headers()

        try:
            resp = requests.post(url, json=payload, headers=headers, timeout=self.timeout)
            resp.raise_for_status()
            data = resp.json()
            embeddings = self._parse_openai_response(data)
            if embeddings:
                EMBEDDINGS_GENERATED_TOTAL.labels(
                    model_name=self.model_name,
                    api_type=self.api_type,
                    content_type="image",
                ).inc()
                return embeddings[0]
        except Exception as exc:
            EMBEDDING_API_ERRORS_TOTAL.labels(
                model_name=self.model_name,
                api_type=self.api_type,
                error_type=type(exc).__name__,
            ).inc()
            logger.warning("Image embedding failed: %s", exc)

        return None

    def _openai_headers(self) -> dict[str, str]:
        """Build headers for OpenAI-compatible requests."""
        headers = {"Content-Type": "application/json"}
        if self.api.api_key:
            headers["Authorization"] = f"Bearer {self.api.api_key}"
        return headers

    def _parse_openai_response(self, data) -> list[list[float]]:
        """
        Parse embedding response from various OpenAI-compatible formats.

        Handles:
        - OpenAI standard: ``{"data": [{"embedding": [...], "index": 0}]}``
        - Direct list of dicts: ``[{"embedding": [...]}]``
        - Direct list of vectors: ``[[0.1, 0.2, ...]]``
        - Dict with embeddings key: ``{"embeddings": [[...]]}``

        :param data: Parsed JSON response.
        :returns: List of embedding vectors.
        """
        if isinstance(data, list):
            if data and isinstance(data[0], dict) and "embedding" in data[0]:
                return [item["embedding"] for item in data]
            return data

        if isinstance(data, dict):
            if "data" in data:
                return [
                    item["embedding"]
                    for item in sorted(data["data"], key=lambda x: x.get("index", 0))
                ]
            if "embedding" in data:
                return [data["embedding"]]
            if "embeddings" in data:
                return data["embeddings"]

        raise ValueError(f"Unexpected embedding response format: {type(data)}")

    # ------------------------------------------------------------------
    # Amazon Bedrock backend
    # ------------------------------------------------------------------

    def _embed_bedrock_single(self, text: str) -> list[float]:
        """Embed a single text via Bedrock Runtime InvokeModel."""
        url = f"{self.base_url}/model/{self.model_name}/invoke"
        headers = {
            "Authorization": f"Bearer {self.api.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        body = {"inputText": text, "normalize": True}

        # Include dimensions if the model supports configurable output
        if self.model.vector_dimensions:
            body["dimensions"] = self.model.vector_dimensions

        logger.debug(
            "Bedrock embedding request model=%s text_len=%d",
            self.model_name,
            len(text),
        )

        with EMBEDDING_BATCH_DURATION.labels(
            model_name=self.model_name, api_type=self.api_type
        ).time():
            try:
                resp = requests.post(
                    url, json=body, headers=headers, timeout=self.timeout
                )
                if resp.status_code != 200:
                    logger.error(
                        "Bedrock embedding failed status=%d body=%s",
                        resp.status_code,
                        resp.text[:500],
                    )
                resp.raise_for_status()
                data = resp.json()
            except requests.RequestException as exc:
                EMBEDDING_API_ERRORS_TOTAL.labels(
                    model_name=self.model_name,
                    api_type=self.api_type,
                    error_type=type(exc).__name__,
                ).inc()
                logger.error("Bedrock embedding request failed: %s", exc)
                raise

        embedding = data.get("embedding")
        if not embedding:
            raise ValueError(f"Bedrock response missing 'embedding' key: {list(data.keys())}")

        token_count = data.get("inputTextTokenCount", len(text.split()))

        EMBEDDINGS_GENERATED_TOTAL.labels(
            model_name=self.model_name,
            api_type=self.api_type,
            content_type="text",
        ).inc()
        EMBEDDING_TOKENS_TOTAL.labels(model_name=self.model_name).inc(token_count)
        self._log_usage(1, token_count)

        logger.debug(
            "Bedrock embedding response dimensions=%d tokens=%d",
            len(embedding),
            token_count,
        )
        return embedding

    def _embed_bedrock_batch(self, texts: list[str]) -> list[list[float]]:
        """
        Embed multiple texts via Bedrock (client-side loop).

        Bedrock InvokeModel accepts one input at a time, so we loop.

        :param texts: List of text strings.
        :returns: List of embedding vectors.
        """
        embeddings = []
        for i, text in enumerate(texts):
            embedding = self._embed_bedrock_single(text)
            embeddings.append(embedding)
            if (i + 1) % 10 == 0:
                logger.debug(
                    "Bedrock batch progress %d/%d", i + 1, len(texts)
                )
        return embeddings

    def _embed_bedrock_image(self, b64_image: str, mime_type: str) -> Optional[list[float]]:
        """Embed an image via Bedrock multimodal endpoint."""
        url = f"{self.base_url}/model/{self.model_name}/invoke"
        headers = {
            "Authorization": f"Bearer {self.api.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        body = {
            "inputImage": b64_image,
            "normalize": True,
        }
        if self.model.vector_dimensions:
            body["dimensions"] = self.model.vector_dimensions

        try:
            resp = requests.post(url, json=body, headers=headers, timeout=self.timeout)
            resp.raise_for_status()
            data = resp.json()
            embedding = data.get("embedding")
            if embedding:
                EMBEDDINGS_GENERATED_TOTAL.labels(
                    model_name=self.model_name,
                    api_type=self.api_type,
                    content_type="image",
                ).inc()
                return embedding
        except Exception as exc:
            EMBEDDING_API_ERRORS_TOTAL.labels(
                model_name=self.model_name,
                api_type=self.api_type,
                error_type=type(exc).__name__,
            ).inc()
            logger.warning("Bedrock image embedding failed: %s", exc)

        return None

    # ------------------------------------------------------------------
    # Usage tracking
    # ------------------------------------------------------------------

    def _log_usage(self, text_count: int, token_count: int):
        """
        Log embedding usage to LLMUsage model.

        :param text_count: Number of texts embedded.
        :param token_count: Approximate token count.
        """
        try:
            from llm_manager.models import LLMUsage

            LLMUsage.objects.create(
                model=self.model,
                user=self.user,
                input_tokens=token_count,
                output_tokens=0,
                cached_tokens=0,
                total_cost=(token_count / 1000) * float(self.model.input_cost_per_1k),
                purpose="embeddings",
            )
        except Exception as exc:
            logger.warning("Failed to log embedding usage: %s", exc)
