"""All builtin functions."""

from typing import Annotated, Any, TYPE_CHECKING

from .typing import Float32, Vector, TypeAttr
from . import op, llm

# Libraries that are heavy to import. Lazily import them later.
if TYPE_CHECKING:
    import sentence_transformers


class ParseJson(op.FunctionSpec):
    """Parse a text into a JSON object."""


class SplitRecursively(op.FunctionSpec):
    """Split a document (in string) recursively."""


class ExtractByLlm(op.FunctionSpec):
    """Extract information from a text using a LLM."""

    llm_spec: llm.LlmSpec
    output_type: type
    instruction: str | None = None


class SentenceTransformerEmbed(op.FunctionSpec):
    """
    `SentenceTransformerEmbed` embeds a text into a vector space using the [SentenceTransformer](https://huggingface.co/sentence-transformers) library.

    Args:

        model: The name of the SentenceTransformer model to use.
        args: Additional arguments to pass to the SentenceTransformer constructor. e.g. {"trust_remote_code": True}
    """

    model: str
    args: dict[str, Any] | None = None


@op.executor_class(gpu=True, cache=True, behavior_version=1)
class SentenceTransformerEmbedExecutor:
    """Executor for SentenceTransformerEmbed."""

    spec: SentenceTransformerEmbed
    _model: "sentence_transformers.SentenceTransformer"

    def analyze(self, text: Any) -> type:
        import sentence_transformers  # pylint: disable=import-outside-toplevel

        args = self.spec.args or {}
        self._model = sentence_transformers.SentenceTransformer(self.spec.model, **args)
        dim = self._model.get_sentence_embedding_dimension()
        result: type = Annotated[
            Vector[Float32, dim],  # type: ignore
            TypeAttr("cocoindex.io/vector_origin_text", text.analyzed_value),
        ]
        return result

    def __call__(self, text: str) -> list[Float32]:
        result: list[Float32] = self._model.encode(text).tolist()
        return result
