from __future__ import annotations

import random
import time

from langchain_core.embeddings import Embeddings
from mistralai import Mistral

from memory_agent.errors import is_rate_limit_error


class MistralEmbedEmbeddings(Embeddings):
    """Embeddings client backed by Mistral embeddings API."""

    def __init__(
        self,
        model_name: str,
        api_token: str,
        max_retries: int = 6,
        base_delay_seconds: float = 1.0,
        max_delay_seconds: float = 16.0,
    ) -> None:
        self._model_name = model_name
        self._client = Mistral(api_key=api_token)
        self._max_retries = max_retries
        self._base_delay_seconds = base_delay_seconds
        self._max_delay_seconds = max_delay_seconds

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        response = self._create_embeddings_with_retry(texts=texts)
        vectors = self._extract_vectors(response=response)
        if len(vectors) != len(texts):
            raise RuntimeError(
                f"Unexpected embeddings count from Mistral API: got {len(vectors)}, expected {len(texts)}."
            )
        return vectors

    def embed_query(self, text: str) -> list[float]:
        embeddings = self.embed_documents([text])
        return embeddings[0]

    def _create_embeddings_with_retry(self, texts: list[str]) -> object:
        for attempt in range(self._max_retries + 1):
            try:
                return self._client.embeddings.create(
                    model=self._model_name,
                    inputs=texts,
                )
            except Exception as error:
                if not is_rate_limit_error(error) or attempt >= self._max_retries:
                    raise
                delay = min(
                    self._max_delay_seconds,
                    self._base_delay_seconds * (2**attempt),
                ) + random.uniform(0.0, 0.25)
                time.sleep(delay)
        raise RuntimeError("Failed to request embeddings after retries.")

    def _extract_vectors(self, response: object) -> list[list[float]]:
        payload_data = getattr(response, "data", None)
        if payload_data is None and hasattr(response, "model_dump"):
            payload_data = response.model_dump().get("data", [])
        if payload_data is None and isinstance(response, dict):
            payload_data = response.get("data", [])
        if payload_data is None:
            raise RuntimeError("Unexpected embeddings response format from Mistral API.")

        indexed_vectors: list[tuple[int | None, list[float]]] = []
        for item in payload_data:
            if isinstance(item, dict):
                index = item.get("index")
                raw_vector = item.get("embedding")
            else:
                index = getattr(item, "index", None)
                raw_vector = getattr(item, "embedding", None)
            if raw_vector is None:
                raise RuntimeError("Missing embedding vector in Mistral API response.")
            indexed_vectors.append((index, [float(value) for value in raw_vector]))

        if indexed_vectors and all(index is not None for index, _ in indexed_vectors):
            indexed_vectors.sort(key=lambda pair: int(pair[0]))
        return [vector for _, vector in indexed_vectors]