"""Embedding helpers for OpenAI-compatible and local providers."""

from __future__ import annotations

import importlib
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast

from agno.knowledge.embedder.openai import OpenAIEmbedder
from agno.utils.log import log_info, log_warning

from mindroom.tool_system.dependencies import ensure_optional_deps

if TYPE_CHECKING:
    from agno.knowledge.embedder.base import Embedder
    from openai.types.create_embedding_response import CreateEmbeddingResponse

_OPENAI_EMBEDDING_DIMENSIONS = {
    "text-embedding-3-large": 3072,
    "text-embedding-3-small": 1536,
}
_MEM0_OPENAI_DEFAULT_DIMENSIONS = 1536
DEFAULT_SENTENCE_TRANSFORMERS_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
_SENTENCE_TRANSFORMERS_DEPENDENCIES = ["sentence-transformers"]
_SENTENCE_TRANSFORMERS_EXTRA = "sentence_transformers"


def _default_dimensions(model: str) -> int | None:
    """Return the default dimensions for models that support the parameter."""
    return _OPENAI_EMBEDDING_DIMENSIONS.get(model)


def effective_knowledge_embedder_signature(
    provider: str,
    model: str,
    *,
    host: str | None = None,
    dimensions: int | None = None,
) -> tuple[str, str, str, str]:
    """Return the knowledge embedder settings that affect indexing behavior."""
    effective_host = host if provider in {"openai", "ollama"} else ""
    effective_dimensions = dimensions
    if provider == "openai" and effective_dimensions is None:
        effective_dimensions = _default_dimensions(model)
    elif provider in {"ollama", "sentence_transformers"}:
        effective_dimensions = None
    return (
        provider,
        model,
        effective_host or "",
        str(effective_dimensions) if effective_dimensions is not None else "",
    )


def effective_mem0_embedder_signature(
    provider: str,
    model: str,
    *,
    host: str | None = None,
    dimensions: int | None = None,
) -> tuple[str, str, str, str]:
    """Return the Mem0 embedder settings that affect memory collection compatibility."""
    effective_host = host if provider in {"openai", "ollama"} else ""
    effective_dimensions = dimensions
    if provider == "openai" and effective_dimensions is None:
        effective_dimensions = _MEM0_OPENAI_DEFAULT_DIMENSIONS
    elif provider in {"ollama", "sentence_transformers"}:
        effective_dimensions = None
    return (
        provider,
        model,
        effective_host or "",
        str(effective_dimensions) if effective_dimensions is not None else "",
    )


def ensure_sentence_transformers_dependencies() -> None:
    """Install the optional local sentence-transformers runtime when needed."""
    ensure_optional_deps(_SENTENCE_TRANSFORMERS_DEPENDENCIES, _SENTENCE_TRANSFORMERS_EXTRA)


def create_sentence_transformers_embedder(
    model: str = DEFAULT_SENTENCE_TRANSFORMERS_MODEL,
    *,
    dimensions: int | None = None,
) -> Embedder:
    """Create a local sentence-transformers embedder after ensuring its optional extra exists."""
    ensure_sentence_transformers_dependencies()
    module = importlib.import_module("agno.knowledge.embedder.sentence_transformer")
    embedder_class = cast("Any", module.SentenceTransformerEmbedder)
    if dimensions is None:
        return cast("Embedder", embedder_class(id=model))
    return cast("Embedder", embedder_class(id=model, dimensions=dimensions))


@dataclass
class MindRoomOpenAIEmbedder(OpenAIEmbedder):
    """Avoid forcing OpenAI defaults onto arbitrary OpenAI-compatible hosts."""

    _dimensions_explicit: bool = field(init=False, default=False, repr=False)

    def __post_init__(self) -> None:
        """Track whether dimensions came from explicit config."""
        self._dimensions_explicit = self.dimensions is not None
        if self.dimensions is None:
            self.dimensions = _default_dimensions(self.id)

    def _should_send_dimensions(self) -> bool:
        return self.dimensions is not None and (self._dimensions_explicit or self.id in _OPENAI_EMBEDDING_DIMENSIONS)

    def _request_params(self, input_value: str | list[str]) -> dict[str, Any]:
        request: dict[str, Any] = {
            "input": input_value,
            "model": self.id,
            "encoding_format": self.encoding_format,
        }
        if self.user is not None:
            request["user"] = self.user
        if self._should_send_dimensions():
            request["dimensions"] = self.dimensions
        if self.request_params:
            request.update(self.request_params)
        return request

    # NOTE: These overrides intentionally mirror agno's async/embedder methods
    # because upstream inlines request construction instead of calling a shared helper.
    # Keep them aligned with agno when upgrading that dependency.
    def response(self, text: str) -> CreateEmbeddingResponse:
        """Request a single embedding synchronously."""
        return self.client.embeddings.create(**self._request_params(text))

    async def async_get_embedding(self, text: str) -> list[float]:
        """Request a single embedding asynchronously."""
        try:
            response: CreateEmbeddingResponse = await self.aclient.embeddings.create(**self._request_params(text))
            return response.data[0].embedding
        except Exception as e:
            log_warning(e)
            return []

    async def async_get_embedding_and_usage(self, text: str) -> tuple[list[float], dict[str, Any] | None]:
        """Request one embedding and its usage payload asynchronously."""
        try:
            response = await self.aclient.embeddings.create(**self._request_params(text))
            embedding = response.data[0].embedding
            usage = response.usage
            return embedding, usage.model_dump() if usage else None
        except Exception as e:
            log_warning(f"Error getting embedding: {e}")
            return [], None

    async def async_get_embeddings_batch_and_usage(
        self,
        texts: list[str],
    ) -> tuple[list[list[float]], list[dict[str, Any] | None]]:
        """Request embeddings for a batch of texts and return per-item usage."""
        all_embeddings: list[list[float]] = []
        all_usage: list[dict[str, Any] | None] = []
        log_info(f"Getting embeddings and usage for {len(texts)} texts in batches of {self.batch_size} (async)")

        for i in range(0, len(texts), self.batch_size):
            batch_texts = texts[i : i + self.batch_size]
            try:
                response: CreateEmbeddingResponse = await self.aclient.embeddings.create(
                    **self._request_params(batch_texts),
                )
                batch_embeddings = [data.embedding for data in response.data]
                all_embeddings.extend(batch_embeddings)

                usage_dict = response.usage.model_dump() if response.usage else None
                all_usage.extend([usage_dict] * len(batch_embeddings))
            except Exception as e:
                log_warning(f"Error in async batch embedding: {e}")
                for text in batch_texts:
                    try:
                        embedding, usage = await self.async_get_embedding_and_usage(text)
                        all_embeddings.append(embedding)
                        all_usage.append(usage)
                    except Exception as inner:
                        log_warning(f"Error in individual async embedding fallback: {inner}")
                        all_embeddings.append([])
                        all_usage.append(None)

        return all_embeddings, all_usage
