"""
Abstract base class for embedding model clients.

Defines a provider-agnostic interface for creating text embeddings from either
single text input or batched text input.
"""

from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import Any, Dict, List, Optional, TypeVar

from pydantic import BaseModel

from agentbyte.component import ComponentLoader
from agentbyte.middleware.base import BaseMiddleware, MiddlewareChain, OperationType

from .types import EmbeddingResult

ConfigT = TypeVar("ConfigT", bound="BaseEmbeddingClientConfig")


class BaseEmbeddingClientConfig(BaseModel):
    """Base configuration class for embedding client serialization."""

    model: str = "text-embedding-3-small"
    config: Dict[str, Any] = {}


class BaseEmbeddingClient(ABC, ComponentLoader):
    """
    Abstract base class for embedding clients.

    Implementations should use dependency-injected provider SDK clients and map
    provider-specific exceptions to the shared LLM error hierarchy.
    """

    def __init__(
        self,
        model: str,
        client: Any,
        config: Optional[Dict[str, Any]] = None,
        middlewares: Optional[List[BaseMiddleware]] = None,
        **kwargs: Any,
    ):
        self.model = model
        self.client = client
        self.config = config or {}
        self.config.update(kwargs)
        self._middleware_chain = MiddlewareChain(middlewares or [])

    @abstractmethod
    async def create(self, input_text: str, **kwargs: Any) -> EmbeddingResult:
        """Create one embedding result from one text input.

        Consumers should read `result.embedding` for the single returned vector.
        The batch-oriented `result.embeddings` field still contains one item.
        """

    @abstractmethod
    async def create_batch(
        self,
        input_texts: List[str],
        **kwargs: Any,
    ) -> EmbeddingResult:
        """Create one embedding result from a batch of text inputs.

        Consumers should read `result.embeddings` for the full batch in request
        order. The `result.embedding` field mirrors the first vector.
        """

    async def _execute_embedding_operation(
        self,
        input_texts: List[str],
        operation_kwargs: Dict[str, Any],
        executor: Callable[[List[str], Dict[str, Any]], Awaitable[EmbeddingResult]],
    ) -> EmbeddingResult:
        """Run embedding execution through middleware chain."""
        if not input_texts:
            raise ValueError("input_texts cannot be empty")

        payload: Dict[str, Any] = {
            "input": list(input_texts),
            "kwargs": dict(operation_kwargs),
        }

        async def operation(data: Dict[str, Any]) -> EmbeddingResult:
            texts = data.get("input")
            if not isinstance(texts, list) or not texts:
                raise ValueError("input_texts cannot be empty")

            if not all(isinstance(item, str) for item in texts):
                raise ValueError("input_texts must contain only strings")

            kwargs = data.get("kwargs", {})
            if not isinstance(kwargs, dict):
                raise ValueError("embedding kwargs must be a dictionary")

            return await executor(texts, kwargs)

        return await self._middleware_chain.execute(
            operation=OperationType.EMBEDDING_CALL.value,
            agent_name=self.__class__.__name__,
            agent_context=None,
            data=payload,
            func=operation,
            metadata={"model": self.model},
        )

    def add_middleware(self, middleware: BaseMiddleware) -> None:
        """Register middleware for embedding operations."""
        self._middleware_chain.add(middleware)

    def remove_middleware(self, middleware: BaseMiddleware) -> None:
        """Unregister middleware for embedding operations."""
        self._middleware_chain.remove(middleware)

    @abstractmethod
    def _to_config(self) -> BaseEmbeddingClientConfig:
        """Serialize client state to a config model."""

    @classmethod
    @abstractmethod
    def _from_config(cls, config: BaseEmbeddingClientConfig) -> "BaseEmbeddingClient":
        """Reconstruct a client from serialized config."""


__all__ = ["BaseEmbeddingClient", "BaseEmbeddingClientConfig"]
