"""Base embedding adapter interface for AI/ML embedding operations."""

import contextlib
import typing as t
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum

import numpy as np
from pydantic import BaseModel, Field, SecretStr
from acb.cleanup import CleanupMixin
from acb.config import AdapterBase, Settings
from acb.ssl_config import SSLConfigMixin


class EmbeddingModel(str, Enum):
    """Standard embedding model identifiers."""

    # OpenAI Models
    TEXT_EMBEDDING_3_SMALL = "text-embedding-3-small"
    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"

    # HuggingFace/Sentence Transformers Models
    ALL_MINILM_L6_V2 = "all-MiniLM-L6-v2"
    ALL_MPNET_BASE_V2 = "all-mpnet-base-v2"
    MULTI_QA_MPNET_BASE_DOT_V1 = "multi-qa-mpnet-base-dot-v1"
    ALL_DISTILROBERTA_V1 = "all-distilroberta-v1"
    PARAPHRASE_MULTILINGUAL_MPNET_BASE_V2 = "paraphrase-multilingual-mpnet-base-v2"

    # ONNX Models (optimized versions)
    ONNX_ALL_MINILM_L6_V2 = "onnx-all-MiniLM-L6-v2"
    ONNX_ALL_MPNET_BASE_V2 = "onnx-all-mpnet-base-v2"

    # Liquid AI LFM Models
    LIQUID_EFFICIENT_BASE = "liquid-efficient-base"
    LIQUID_EDGE_OPTIMIZED = "liquid-edge-optimized"


class PoolingStrategy(str, Enum):
    """Embedding pooling strategies."""

    MEAN = "mean"
    MAX = "max"
    CLS = "cls"
    WEIGHTED_MEAN = "weighted_mean"


class VectorNormalization(str, Enum):
    """Vector normalization methods."""

    L2 = "l2"
    L1 = "l1"
    NONE = "none"


@dataclass
class EmbeddingRequest:
    """Request for embedding generation."""

    texts: list[str]
    model: str | EmbeddingModel | None = None
    normalize: bool = True
    pooling: PoolingStrategy = PoolingStrategy.MEAN
    chunk_size: int | None = None
    batch_size: int = 32


@dataclass
class EmbeddingResponse:
    """Response from embedding generation."""

    embeddings: list[list[float]]
    model: str
    dimensions: int
    tokens_used: int | None = None
    processing_time: float | None = None


class EmbeddingResult(BaseModel):
    """Single embedding result."""

    text: str
    embedding: list[float]
    model: str
    dimensions: int
    tokens: int | None = None
    metadata: dict[str, t.Any] = Field(default_factory=dict)


class EmbeddingBatch(BaseModel):
    """Batch embedding results."""

    results: list[EmbeddingResult]
    total_tokens: int | None = None
    processing_time: float | None = None
    model: str
    batch_size: int


class EmbeddingBaseSettings(Settings, SSLConfigMixin):
    """Base settings for embedding adapters."""

    model: str = Field(default="text-embedding-3-small")
    api_key: SecretStr | None = Field(default=None)
    base_url: str | None = Field(default=None)
    max_retries: int = Field(default=3)
    timeout: float = Field(default=30.0)
    batch_size: int = Field(default=32)
    max_tokens_per_batch: int = Field(default=8192)
    normalize_embeddings: bool = Field(default=True)
    cache_embeddings: bool = Field(default=True)
    cache_ttl: int = Field(default=3600)  # 1 hour

    # Edge optimization settings
    memory_limit_mb: int = Field(default=512)
    enable_model_caching: bool = Field(default=True)
    prefetch_models: bool = Field(default=False)

    # Advanced processing settings
    chunk_size: int = Field(default=512)
    chunk_overlap: int = Field(default=50)
    pooling_strategy: PoolingStrategy = Field(default=PoolingStrategy.MEAN)
    normalization: VectorNormalization = Field(default=VectorNormalization.L2)


class EmbeddingAdapter(AdapterBase, CleanupMixin, ABC):
    """Base embedding adapter with standard interface."""

    def __init__(self, settings: EmbeddingBaseSettings | None = None) -> None:
        super().__init__()
        self._settings = settings or EmbeddingBaseSettings()
        self._client = None
        self._model_cache: dict[str, t.Any] = {}
        self._ready_event = None

    @property
    def settings(self) -> EmbeddingBaseSettings:
        """Get adapter settings."""
        return self._settings

    @property
    def client(self) -> t.Any:
        """Get the embedding client."""
        return self._client

    async def __aenter__(self) -> "EmbeddingAdapter":
        """Async context manager entry."""
        await self._ensure_client()
        return self

    async def __aexit__(
        self: t.Any, exc_type: t.Any, exc_val: t.Any, exc_tb: t.Any
    ) -> None:
        """Async context manager exit."""
        await self.cleanup()

    # Public API Methods

    async def embed_text(
        self,
        text: str,
        model: str | EmbeddingModel | None = None,
        normalize: bool | None = None,
        **kwargs: t.Any,
    ) -> list[float]:
        """Generate embedding for a single text."""
        result = await self.embed_texts(
            [text],
            model=model,
            normalize=normalize,
            **kwargs,
        )
        return result.results[0].embedding

    async def embed_texts(
        self,
        texts: list[str],
        model: str | EmbeddingModel | None = None,
        normalize: bool | None = None,
        batch_size: int | None = None,
        **kwargs: t.Any,
    ) -> EmbeddingBatch:
        """Generate embeddings for multiple texts."""
        return await self._embed_texts(
            texts,
            model=model or self._settings.model,
            normalize=normalize
            if normalize is not None
            else self._settings.normalize_embeddings,
            batch_size=batch_size or self._settings.batch_size,
            **kwargs,
        )

    async def embed_documents(
        self,
        documents: list[str],
        chunk_size: int | None = None,
        chunk_overlap: int | None = None,
        model: str | EmbeddingModel | None = None,
        **kwargs: t.Any,
    ) -> list[EmbeddingBatch]:
        """Embed large documents with chunking."""
        return await self._embed_documents(
            documents,
            chunk_size=chunk_size or self._settings.chunk_size,
            chunk_overlap=chunk_overlap or self._settings.chunk_overlap,
            model=model or self._settings.model,
            **kwargs,
        )

    async def compute_similarity(
        self,
        embedding1: list[float],
        embedding2: list[float],
        method: str = "cosine",
    ) -> float:
        """Compute similarity between two embeddings."""
        return await self._compute_similarity(embedding1, embedding2, method)

    async def get_model_info(
        self,
        model: str | EmbeddingModel | None = None,
    ) -> dict[str, t.Any]:
        """Get information about an embedding model."""
        return await self._get_model_info(model or self._settings.model)

    async def list_models(self) -> list[dict[str, t.Any]]:
        """List available embedding models."""
        return await self._list_models()

    # Abstract methods that implementations must provide

    @abstractmethod
    async def _ensure_client(self) -> t.Any:
        """Ensure client is initialized."""

    @abstractmethod
    async def _embed_texts(
        self,
        texts: list[str],
        model: str,
        normalize: bool,
        batch_size: int,
        **kwargs: t.Any,
    ) -> EmbeddingBatch:
        """Implementation-specific text embedding."""

    @abstractmethod
    async def _embed_documents(
        self,
        documents: list[str],
        chunk_size: int,
        chunk_overlap: int,
        model: str,
        **kwargs: t.Any,
    ) -> list[EmbeddingBatch]:
        """Implementation-specific document embedding."""

    @abstractmethod
    async def _compute_similarity(
        self,
        embedding1: list[float],
        embedding2: list[float],
        method: str,
    ) -> float:
        """Implementation-specific similarity computation."""

    @abstractmethod
    async def _get_model_info(self, model: str) -> dict[str, t.Any]:
        """Implementation-specific model information."""

    @abstractmethod
    async def _list_models(self) -> list[dict[str, t.Any]]:
        """Implementation-specific model listing."""

    # Utility methods

    def _normalize_vector(
        self,
        vector: list[float],
        method: VectorNormalization = VectorNormalization.L2,
    ) -> list[float]:
        """Normalize a vector using specified method."""
        if method == VectorNormalization.NONE:
            return vector

        np_vector = np.array(vector)

        if method == VectorNormalization.L2:
            norm = np.linalg.norm(np_vector)
            if norm == 0:
                return vector
            normalized: list[float] = (np_vector / norm).tolist()
            return normalized
        if method == VectorNormalization.L1:
            norm = np.sum(np.abs(np_vector))
            if norm == 0:
                return vector
            normalized_l1: list[float] = (np_vector / norm).tolist()
            return normalized_l1

        return vector

    def _chunk_text(self, text: str, chunk_size: int, overlap: int = 0) -> list[str]:
        """Split text into overlapping chunks."""
        if len(text) <= chunk_size:
            return [text]

        chunks = []
        start = 0

        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]
            chunks.append(chunk)

            if end >= len(text):
                break

            start = end - overlap

        return chunks

    def _batch_texts(self, texts: list[str], batch_size: int) -> list[list[str]]:
        """Split texts into batches."""
        return [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]

    async def cleanup(self) -> None:
        """Clean up resources."""
        if hasattr(self, "_client") and self._client:
            if hasattr(self._client, "close"):
                await self._client.close()
            elif hasattr(self._client, "__aexit__"):
                with contextlib.suppress(Exception):
                    await self._client.__aexit__(None, None, None)

        self._client = None
        self._model_cache.clear()

        await super().cleanup()


# Embedding utilities
class EmbeddingUtils:
    """Utility functions for embedding operations."""

    @staticmethod
    def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
        """Compute cosine similarity between two vectors."""
        np_vec1 = np.array(vec1)
        np_vec2 = np.array(vec2)

        dot_product = np.dot(np_vec1, np_vec2)
        norm1 = np.linalg.norm(np_vec1)
        norm2 = np.linalg.norm(np_vec2)

        if 0 in (norm1, norm2):
            return 0.0

        return float(dot_product / (norm1 * norm2))

    @staticmethod
    def euclidean_distance(vec1: list[float], vec2: list[float]) -> float:
        """Compute Euclidean distance between two vectors."""
        np_vec1 = np.array(vec1)
        np_vec2 = np.array(vec2)
        return float(np.linalg.norm(np_vec1 - np_vec2))

    @staticmethod
    def dot_product(vec1: list[float], vec2: list[float]) -> float:
        """Compute dot product between two vectors."""
        np_vec1 = np.array(vec1)
        np_vec2 = np.array(vec2)
        return float(np.dot(np_vec1, np_vec2))

    @staticmethod
    def manhattan_distance(vec1: list[float], vec2: list[float]) -> float:
        """Compute Manhattan distance between two vectors."""
        np_vec1 = np.array(vec1)
        np_vec2 = np.array(vec2)
        return float(np.sum(np.abs(np_vec1 - np_vec2)))


# Type aliases for convenience
EmbeddingVector = list[float]
EmbeddingMatrix = list[list[float]]
