"""Modal GPU EZ 임베딩 모델을 위한 LangChain Embeddings 래퍼."""

from __future__ import annotations

from typing import Any

from langchain_core.embeddings import Embeddings
from pydantic import BaseModel, ConfigDict, Field


class ModalGpuEzEmbeddings(BaseModel, Embeddings):
    """modal_gpu_ez를 사용하는 LangChain Embeddings 구현.

    Example:
        .. code-block:: python

            from langchain_modal_gpu_ez import ModalGpuEzEmbeddings

            embeddings = ModalGpuEzEmbeddings(
                model_id="BAAI/bge-small-en-v1.5",
                gpu="Local",
            )
            vectors = embeddings.embed_documents(["hello", "world"])
    """

    model_id: str = "BAAI/bge-small-en-v1.5"
    gpu: str = "Local"
    model_kwargs: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(
        extra="forbid",
        protected_namespaces=(),
    )

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        """텍스트 리스트를 임베딩 벡터 리스트로 변환한다."""
        import modal_gpu_ez as mg

        result = mg.use(
            self.model_id,
            self.gpu,
            input=texts,
            **self.model_kwargs,
        )
        return result  # type: ignore[return-value]

    def embed_query(self, text: str) -> list[float]:
        """단일 텍스트를 임베딩 벡터로 변환한다."""
        return self.embed_documents([text])[0]

    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
        """비동기로 텍스트 리스트를 임베딩 벡터 리스트로 변환한다."""
        import modal_gpu_ez as mg

        result = await mg.async_use(
            self.model_id,
            self.gpu,
            input=texts,
            **self.model_kwargs,
        )
        return result  # type: ignore[return-value]

    async def aembed_query(self, text: str) -> list[float]:
        """비동기로 단일 텍스트를 임베딩 벡터로 변환한다."""
        result = await self.aembed_documents([text])
        return result[0]
