# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors


from functools import cached_property
from typing import List, Any, Optional, Tuple

import numpy as np

from pydantic import ConfigDict, PrivateAttr

from ..util import attempt_import_or_raise
from .base import EmbeddingFunction
from .registry import register
from .utils import TEXT


@register("huggingface")
class TransformersEmbeddingFunction(EmbeddingFunction):
    """
    An embedding function that can use any model from the transformers library.

    Parameters:
    ----------
    name : str
        The name of the model to use. This should be a model name that can be loaded
        by transformers.AutoModel.from_pretrained. For example, "bert-base-uncased".
        default: "colbert-ir/colbertv2.0""
    device : str
        The device to use for the model. Default is "cpu".
    show_progress_bar : bool
        Whether to show a progress bar when loading the model. Default is True.
    trust_remote_code : bool
        Whether or not to allow for custom models defined on the HuggingFace
        Hub in their own modeling files. This option should only be set to True
        for repositories you trust and in which you have read the code, as it
        will execute code present on the Hub on your local machine. When
        restored from table metadata, this is only enabled for models allowed with
        [allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
    revision : Optional[str]
        The model revision to load, such as a commit hash. Default is None.

    to download package, run :
        `pip install transformers`
    you may need to install pytorch as well - `https://pytorch.org/get-started/locally/`

    """

    name: str = "colbert-ir/colbertv2.0"
    device: str = "cpu"
    trust_remote_code: bool = False
    revision: Optional[str] = None
    _tokenizer: Any = PrivateAttr()
    _model: Any = PrivateAttr()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._ndims = None
        transformers = attempt_import_or_raise("transformers")
        self._tokenizer = transformers.AutoTokenizer.from_pretrained(
            self.name, revision=self.revision
        )
        self._model = transformers.AutoModel.from_pretrained(
            self.name, trust_remote_code=self.trust_remote_code, revision=self.revision
        )
        self._model.to(self.device)

    def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
        return (self.name, self.revision)

    model_config = ConfigDict(ignored_types=(cached_property,))

    def ndims(self):
        self._ndims = self._model.config.hidden_size
        return self._ndims

    def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]:
        return self.compute_source_embeddings(query)

    def compute_source_embeddings(self, texts: TEXT, *args, **kwargs) -> List[np.array]:
        texts = self.sanitize_input(texts)
        embedding = []
        for text in texts:
            encoding = self._tokenizer(
                text, return_tensors="pt", padding=True, truncation=True
            ).to(self.device)
            emb = self._model(**encoding).last_hidden_state.mean(dim=1).squeeze()
            embedding.append(emb.tolist())

        return embedding


@register("colbert")
class ColbertEmbeddings(TransformersEmbeddingFunction):
    """
    An embedding function that uses the colbert model from the huggingface library.

    Parameters:
    ----------
    name : str
        The name of the model to use. This should be a model name that can be loaded
        by transformers.AutoModel.from_pretrained. For example, "bert-base-uncased".
        default: "colbert-ir/colbertv2.0""

    to download package, run :
        `pip install transformers`
    you may need to install pytorch as well - `https://pytorch.org/get-started/locally/`

    """

    name: str = "colbert-ir/colbertv2.0"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)