#
#   Copyright 2023 Logical Clocks AB
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
from __future__ import annotations

import json
from dataclasses import dataclass

import humps
from hopsworks_common import client
from hopsworks_common.client.exceptions import FeatureStoreException
from hsfs import util
from hsfs.core.vector_db_client import VectorDbClient


class SimilarityFunctionType:
    """Enumeration class representing different types of similarity functions.

    Parameters:
        L2: Represents L2 norm similarity function.
        COSINE: Represents cosine similarity function.
        DOT_PRODUCT: Represents dot product similarity function.
    """

    L2: str = "l2_norm"
    COSINE: str = "cosine"
    DOT_PRODUCT: str = "dot_product"

    def __init__(self) -> None:
        # Fix for the doc
        raise NotImplementedError("This class should not be instantiated.")


@dataclass
class HsmlModel:
    """Data class storing the metadata of a hsml model."""

    model_registry_id: int
    model_name: str
    model_version: int

    @classmethod
    def from_response_json(cls, json_dict):
        json_decamelized = humps.decamelize(json_dict)

        return cls(
            model_registry_id=json_decamelized.get("model_registry_id"),
            model_name=json_decamelized.get("model_name"),
            model_version=json_decamelized.get("model_version"),
        )

    @classmethod
    def from_model(cls, model):
        return cls(
            model_registry_id=model.model_registry_id,
            model_name=model.name,
            model_version=model.version,
        )

    # should get from backend because of authorisation check (unshared project etc)
    def get_model(self):
        try:
            from hsml.model import Model
        except ModuleNotFoundError as err:
            raise FeatureStoreException(
                "Model is attached to embedding feature but hsml library is not installed."
                "Install hsml library before getting the feature group."
            ) from err

        path_params = [
            "project",
            client.get_instance()._project_id,
            "modelregistries",
            self.model_registry_id,
            "models",
            self.model_name + "_" + str(self.model_version),
        ]
        query_params = {"expand": "trainingdatasets"}

        model_json = client.get_instance()._send_request(
            "GET", path_params, query_params
        )
        return Model.from_response_json(model_json)

    def json(self):
        return json.dumps(self, cls=util.Encoder)

    def to_dict(self):
        return {
            "modelRegistryId": self.model_registry_id,
            "modelName": self.model_name,
            "modelVersion": self.model_version,
        }


class EmbeddingFeature:
    """Represents an embedding feature.

    Parameters:
        name: The name of the embedding feature.
        dimension: The dimensionality of the embedding feature.
        similarity_function_type:
            The type of similarity function used for the embedding feature.
            Available functions are `L2`, `COSINE`, and `DOT_PRODUCT`.
            Defaults to `SimilarityFunctionType.L2`.
        model: `hsml.model.Model` A Model in hsml.
        feature_group: The feature group object that contains the embedding feature.
        embedding_index: `EmbeddingIndex` The index for managing embedding features.
    """

    def __init__(
        self,
        name: str = None,
        dimension: int = None,
        similarity_function_type: SimilarityFunctionType = SimilarityFunctionType.L2,
        model=None,
        feature_group=None,
        embedding_index=None,
    ):
        self._name = name
        self._dimension = dimension
        self._similarity_function_type = similarity_function_type
        self._model = model
        self._feature_group = feature_group
        self._embedding_index = embedding_index

    @classmethod
    def from_response_json(cls, json_dict):
        json_decamelized = humps.decamelize(json_dict)
        hsml_model_json = json_decamelized.get("model")
        hsml_model = (
            HsmlModel.from_response_json(hsml_model_json).get_model()
            if hsml_model_json
            else None
        )

        return cls(
            name=json_decamelized.get("name"),
            dimension=json_decamelized.get("dimension"),
            similarity_function_type=json_decamelized.get("similarity_function_type"),
            model=hsml_model,
        )

    @property
    def name(self):
        """str: The name of the embedding feature."""
        return self._name

    @property
    def dimenstion(self):
        """The dimensionality of the embedding feature.

        This one is excluded from the docs as the name is misspelled but kept to avoid breaking the API.
        """
        return self._dimension

    @property
    def dimension(self):
        """int: The dimensionality of the embedding feature."""
        return self._dimension

    @property
    def similarity_function_type(self):
        """SimilarityFunctionType: The type of similarity function used for the embedding feature."""
        return self._similarity_function_type

    @property
    def model(self):
        """hsml.model.Model: The Model in hsml."""
        return self._model

    @property
    def feature_group(self):
        """FeatureGroup: The feature group object that contains the embedding feature."""
        return self._feature_group

    @feature_group.setter
    def feature_group(self, feature_group):
        """Set the feature group object that contains the embedding feature.

        Args:
            feature_group (FeatureGroup): The feature group object.
        """
        self._feature_group = feature_group

    @property
    def embedding_index(self):
        """EmbeddingIndex: The index for managing embedding features."""
        return self._embedding_index

    @embedding_index.setter
    def embedding_index(self, embedding_index):
        """Set the index for managing embedding features.

        Args:
            embedding_index (EmbeddingIndex): The embedding index object.
        """
        self._embedding_index = embedding_index

    def json(self):
        """Serialize the EmbeddingFeature object to a JSON string."""
        return json.dumps(self, cls=util.Encoder)

    def to_dict(self):
        """Convert the EmbeddingFeature object to a dictionary.

        Returns:
            `dict`: A dictionary representation of the EmbeddingFeature object.
        """
        d = {
            "name": self._name,
            "dimension": self._dimension,
            "similarityFunctionType": self._similarity_function_type,
        }
        if self._model:
            d["model"] = HsmlModel.from_model(self._model)
        return d

    def __repr__(self):
        return self.json()


class EmbeddingIndex:
    """Represents an index for managing embedding features.

    Parameters:
        index_name:
            The name of the embedding index.
            The name of the project index is used if not provided.
        features: A list of `EmbeddingFeature` objects for the features that contain embeddings that should be indexed for similarity search.
        col_prefix:
            The prefix to be added to column names when using project index.
            It is managed by Hopsworks and should not be provided.

    Example:
        ```python
        embedding_index = EmbeddingIndex()
        embedding_index.add_embedding(name="user_vector", dimension=256)
        embeddings = embedding_index.get_embeddings()
        ```
    """

    def __init__(
        self,
        index_name: str | None = None,
        features: list[EmbeddingFeature] | None = None,
        col_prefix: str | None = None,
    ):
        self._index_name = index_name
        if features is None:
            self._features = {}
        else:
            self._features = {feat.name: feat for feat in features}
        self._feature_group = None
        self._col_prefix = col_prefix
        self._vector_db_client = None

    def add_embedding(
        self,
        name: str,
        dimension: int,
        similarity_function_type: SimilarityFunctionType
        | None = SimilarityFunctionType.L2,
        model=None,
    ):
        """Adds a new embedding feature to the index.

        Example:
            ```python
            embedding_index = EmbeddingIndex()
            embedding_index.add_embedding(name="user_vector", dimension=256)

            # Attach a hsml model to the embedding feature
            embedding_index = EmbeddingIndex()
            embedding_index.add_embedding(name="user_vector", dimension=256, model=hsml_model)
            ```

        Parameters:
            name: The name of the embedding feature.
            dimension: The dimensionality of the embedding feature.
            similarity_function_type: The type of similarity function to be used.
            model: `hsml.model.Model | None` The hsml model used to generate the embedding.
        """
        self._features[name] = EmbeddingFeature(
            name, dimension, similarity_function_type, model=model
        )

    def get_embedding(self, name: str) -> EmbeddingFeature:
        """Get `EmbeddingFeature` associated with the feature name.

        Parameters:
            name: The name of the embedding feature.

        Returns:
            The `EmbeddingFeature` associated with the name.
        """
        feat = self._features.get(name)
        if feat:
            feat.feature_group = self._feature_group
            feat.embedding_index = self
        return feat

    def get_embeddings(self) -> list[EmbeddingFeature]:
        """Returns the list of `EmbeddingFeature` objects associated with the index.

        Returns:
            All embedding features in the index.
        """
        for feat in self._features.values():
            feat.feature_group = self._feature_group
            feat.embedding_index = self
        return self._features.values()

    def count(self, options: map = None) -> int:
        """Count the number of records in the feature group.

        Parameters:
            options:
                The options used for the request to the vector database.
                The keys are attribute values of [`OpensearchRequestOption`][hopsworks_common.core.opensearch.OpensearchRequestOption].

        Returns:
            The number of records in the feature group.

        Raises:
            ValueError: If the feature group is not initialized.
            hopsworks.client.exceptions.FeatureStoreException: If an error occurs during the count operation.
        """
        if self._vector_db_client is None:
            self._vector_db_client = VectorDbClient(self._feature_group.select_all())
        return self._vector_db_client.count(self.feature_group, options=options)

    @classmethod
    def from_response_json(cls, json_dict):
        json_decamelized = humps.decamelize(json_dict)
        return cls(
            index_name=json_decamelized.get("index_name"),
            features=[
                EmbeddingFeature.from_response_json(f)
                for f in json_decamelized.get("features")
            ],
            col_prefix=json_decamelized.get("col_prefix"),
        )

    @property
    def feature_group(self):
        """FeatureGroup: The feature group object that contains the embedding feature."""
        return self._feature_group

    @feature_group.setter
    def feature_group(self, feature_group):
        """Setter for the feature group object."""
        self._feature_group = feature_group

    @property
    def index_name(self):
        """str: The name of the embedding index."""
        return self._index_name

    @property
    def col_prefix(self):
        """str: The prefix to be added to column names."""
        return self._col_prefix

    def json(self):
        """Serialize the EmbeddingIndex object to a JSON string."""
        return json.dumps(self, cls=util.Encoder)

    def to_dict(self) -> dict:
        """Convert the EmbeddingIndex object to a dictionary.

        Returns:
            A dictionary representation of the EmbeddingIndex object.
        """
        return {
            "indexName": self._index_name,
            "features": [feature.to_dict() for feature in self._features.values()],
            "colPrefix": self._col_prefix,
        }

    def __repr__(self):
        return self.json()
