"""Embedding providers for vector memory.

Supports multiple embedding backends:
- Kimi API (cloud-based)
- Ollama (local)
"""

from __future__ import annotations

import os
from abc import ABC, abstractmethod
from typing import Any

import httpx


class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers."""

    @abstractmethod
    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for a list of texts.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors (list of floats).
        """
        pass

    @abstractmethod
    def embed_sync(self, texts: list[str]) -> list[list[float]]:
        """Synchronous version of embed.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors (list of floats).
        """
        pass

    @abstractmethod
    def get_dimension(self) -> int:
        """Get the dimension of the embedding vectors.

        Returns:
            Dimension of the embedding vectors.
        """
        pass


class KimiEmbeddingProvider(EmbeddingProvider):
    """Kimi API embedding provider.

    Uses Moonshot AI's embedding API for high-quality embeddings.
    """

    DEFAULT_BASE_URL = "https://api.moonshot.cn/v1"
    DEFAULT_MODEL = "moonshot-embedding-v1"
    DIMENSION = 2048

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        model: str | None = None,
    ) -> None:
        """Initialize the Kimi embedding provider.

        Args:
            api_key: Kimi API key. If not provided, uses KIMI_API_KEY env var.
            base_url: Base URL for Kimi API.
            model: Embedding model to use.
        """
        self.api_key = api_key or os.environ.get("KIMI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "Kimi API key is required. Set KIMI_API_KEY environment variable "
                "or pass api_key parameter."
            )

        self.base_url = base_url or self.DEFAULT_BASE_URL
        self.model = model or self.DEFAULT_MODEL
        self._client: httpx.AsyncClient | None = None
        self._sync_client: httpx.Client | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=httpx.Timeout(60.0, connect=30.0),
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.api_key}",
                },
            )
        return self._client

    def _get_sync_client(self) -> httpx.Client:
        """Get or create sync HTTP client."""
        if self._sync_client is None or self._sync_client.is_closed:
            self._sync_client = httpx.Client(
                base_url=self.base_url,
                timeout=httpx.Timeout(60.0, connect=30.0),
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.api_key}",
                },
            )
        return self._sync_client

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings using Kimi API.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors.
        """
        client = await self._get_client()

        # Kimi API supports batching, but let's limit batch size
        batch_size = 50
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]

            response = await client.post(
                "/embeddings",
                json={
                    "model": self.model,
                    "input": batch,
                },
            )
            response.raise_for_status()
            data = response.json()

            # Sort by index to maintain order
            embeddings = sorted(data.get("data", []), key=lambda x: x.get("index", 0))
            all_embeddings.extend([e["embedding"] for e in embeddings])

        return all_embeddings

    def embed_sync(self, texts: list[str]) -> list[list[float]]:
        """Synchronous version of embed.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors.
        """
        client = self._get_sync_client()

        batch_size = 50
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]

            response = client.post(
                "/embeddings",
                json={
                    "model": self.model,
                    "input": batch,
                },
            )
            response.raise_for_status()
            data = response.json()

            embeddings = sorted(data.get("data", []), key=lambda x: x.get("index", 0))
            all_embeddings.extend([e["embedding"] for e in embeddings])

        return all_embeddings

    def get_dimension(self) -> int:
        """Get embedding dimension."""
        return self.DIMENSION

    async def close(self) -> None:
        """Close async HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()

    def close_sync(self) -> None:
        """Close sync HTTP client."""
        if self._sync_client and not self._sync_client.is_closed:
            self._sync_client.close()


class OllamaEmbeddingProvider(EmbeddingProvider):
    """Ollama local embedding provider.

    Uses local Ollama instance for embeddings.
    """

    DEFAULT_BASE_URL = "http://localhost:11434"
    DEFAULT_MODEL = "nomic-embed-text"

    # Model dimensions
    MODEL_DIMENSIONS = {
        "nomic-embed-text": 768,
        "mxbai-embed-large": 1024,
        "snowflake-arctic-embed": 1024,
    }

    def __init__(
        self,
        base_url: str | None = None,
        model: str | None = None,
    ) -> None:
        """Initialize the Ollama embedding provider.

        Args:
            base_url: Base URL for Ollama API.
            model: Embedding model to use.
        """
        self.base_url = base_url or self.DEFAULT_BASE_URL
        self.model = model or self.DEFAULT_MODEL
        self._client: httpx.AsyncClient | None = None
        self._sync_client: httpx.Client | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=httpx.Timeout(60.0, connect=10.0),
            )
        return self._client

    def _get_sync_client(self) -> httpx.Client:
        """Get or create sync HTTP client."""
        if self._sync_client is None or self._sync_client.is_closed:
            self._sync_client = httpx.Client(
                base_url=self.base_url,
                timeout=httpx.Timeout(60.0, connect=10.0),
            )
        return self._sync_client

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings using Ollama API.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors.
        """
        client = await self._get_client()
        all_embeddings = []

        # Ollama embeddings are typically done one at a time
        for text in texts:
            response = await client.post(
                "/api/embeddings",
                json={
                    "model": self.model,
                    "prompt": text,
                },
            )
            response.raise_for_status()
            data = response.json()
            all_embeddings.append(data["embedding"])

        return all_embeddings

    def embed_sync(self, texts: list[str]) -> list[list[float]]:
        """Synchronous version of embed.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors.
        """
        client = self._get_sync_client()
        all_embeddings = []

        for text in texts:
            response = client.post(
                "/api/embeddings",
                json={
                    "model": self.model,
                    "prompt": text,
                },
            )
            response.raise_for_status()
            data = response.json()
            all_embeddings.append(data["embedding"])

        return all_embeddings

    def get_dimension(self) -> int:
        """Get embedding dimension for the current model."""
        return self.MODEL_DIMENSIONS.get(self.model, 768)

    async def close(self) -> None:
        """Close async HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()

    def close_sync(self) -> None:
        """Close sync HTTP client."""
        if self._sync_client and not self._sync_client.is_closed:
            self._sync_client.close()

    async def health_check(self) -> dict[str, Any]:
        """Check if Ollama is available.

        Returns:
            Health status dictionary.
        """
        try:
            client = await self._get_client()
            response = await client.get("/api/tags")
            response.raise_for_status()
            return {
                "status": "healthy",
                "base_url": self.base_url,
                "model": self.model,
            }
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e),
                "base_url": self.base_url,
            }


def create_embedding_provider(
    provider: str = "auto",
    **kwargs: Any,
) -> EmbeddingProvider:
    """Create an embedding provider.

    Args:
        provider: Provider type ("kimi", "ollama", or "auto").
        **kwargs: Additional arguments for the provider.

    Returns:
        Configured embedding provider.

    Raises:
        ValueError: If provider type is unknown.
    """
    if provider == "kimi":
        return KimiEmbeddingProvider(**kwargs)
    elif provider == "ollama":
        return OllamaEmbeddingProvider(**kwargs)
    elif provider == "auto":
        # Try Kimi first if API key is available
        if os.environ.get("KIMI_API_KEY"):
            return KimiEmbeddingProvider(**kwargs)
        else:
            return OllamaEmbeddingProvider(**kwargs)
    else:
        raise ValueError(f"Unknown embedding provider: {provider}")
