from __future__ import annotations

from typing import Any, List, Optional, Set, Type, Union, final
from abc import ABC
from mloda_core.abstract_plugins.components.base_artifact import BaseArtifact
from mloda_core.abstract_plugins.components.data_access_collection import DataAccessCollection
from mloda_core.abstract_plugins.components.data_types import DataType

from mloda_core.abstract_plugins.components.domain import Domain
from mloda_core.abstract_plugins.components.feature_name import FeatureName
from mloda_core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
from mloda_core.abstract_plugins.components.input_data.base_input_data import BaseInputData
from mloda_core.abstract_plugins.components.input_data.creator.data_creator import DataCreator
from mloda_core.abstract_plugins.compute_frame_work import ComputeFrameWork
from mloda_core.abstract_plugins.components.feature import Feature
from mloda_core.abstract_plugins.components.feature_set import FeatureSet
from mloda_core.abstract_plugins.components.options import Options
from mloda_core.abstract_plugins.components.index.index import Index
from mloda_core.abstract_plugins.components.utils import get_all_subclasses


class AbstractFeatureGroup(ABC):
    """
    Mostly implement:
    input_features, except it is a primary source

    Implement if necessary:
    - match_feature_group_criteria - default is the class name
    - domain - default is the default domain
    - compute_framework_rule - default is true, which sets the compute framework to all available compute frameworks
    - index_columns - default is None
    - return_data_type_rule - default is None
    """

    def __init__(self) -> None:
        pass

    @classmethod
    def input_data(cls) -> Optional[BaseInputData]:
        """This function should return the input data class used for this feature group."""
        return None

    @classmethod
    def validate_input_features(cls, data: Any, features: FeatureSet) -> Optional[bool]:
        """This function should be used to validate the input data."""
        return None

    @classmethod
    def validate_output_features(cls, data: Any, features: FeatureSet) -> Optional[bool]:
        """This function should be used to validate the output data."""
        return None

    @classmethod
    def calculate_feature(cls, data: Any, features: FeatureSet) -> Any:
        """This function should be used to calculate the feature."""
        return None

    @staticmethod
    def artifact() -> Type[BaseArtifact] | None:
        """Returns the artifact associated with this feature group.

        Artifacts are data generated by a feature group and can be used by other feature groups.
        This is necessary for scenarios such as embeddings, where the output of one feature group
        serves as an input for another, enabling complex data transformations and feature engineering
        workflows.

        This method should be overridden by subclasses to provide the specific artifact
        that the feature group generates or uses. If no artifact is associated with the
        feature group, this method should return None.
        """
        return None

    @final
    @classmethod
    def load_artifact(cls, features: FeatureSet) -> Any:
        """
        Convenience function to load an artifact associated with the given FeatureSet.

        This method utilizes the `artifact` method to retrieve the specific artifact class
        associated with the feature group. It then calls the `load` method of the artifact class
        to load the artifact data.
        """
        artifact = cls.artifact()
        if artifact is None:
            raise ValueError(f"Artifact load is called, but not implemented: {cls.get_class_name()}.")
        return artifact.load(features)

    def set_feature_name(self, config: Options, feature_name: FeatureName) -> FeatureName:
        return feature_name

    @classmethod
    def return_data_type_rule(cls, feature: Feature) -> Optional[DataType]:
        """This function should return the data type for the feature group
        if this feature group has a fixed return data type.
        """
        return None

    def input_features(self, options: Options, feature_name: FeatureName) -> Optional[Set[Feature]]:
        """This function should return the input features for the feature group
        if this feature is dependent on other features.
        Else, return None"""
        raise NotImplementedError

    @classmethod
    def index_columns(cls) -> Optional[List[Index]]:
        """This function should return the index columns to merge on, by names, by the user set index or by the features.

        However, here we can validate given indexe by the aforementioned sources, if implemented
        """
        return None

    @classmethod
    def supports_index(cls, index: Index) -> Optional[bool]:
        """This function should return true if the incoming index is supported by the feature group."""

        supported_index_columns = cls.index_columns()

        if supported_index_columns is None:
            return None

        for supported_index_column in supported_index_columns:
            if index.is_a_part_of_(supported_index_column):
                return True

        return False

    @classmethod
    def match_feature_group_criteria(
        cls,
        feature_name: Union[FeatureName, str],
        options: Options,
        data_access_collection: Optional[DataAccessCollection] = None,
    ) -> bool:
        """This function returns true if the feature_name is matching with this feature group.
        The if statement contains the rules.
        Each case has different use cases. You can disallow them be removing them. However, often you can just use the default.
        If you want to implement a concrete implementation, e.g.just accept specific names, then you can overwrite this function.
        """

        if isinstance(feature_name, FeatureName):
            feature_name = feature_name.name

        if cls().is_root(options, feature_name):
            input_data_class = cls.input_data()

            if input_data_class is not None:
                if isinstance(input_data_class, DataCreator):
                    return input_data_class.matches(feature_name, options, None)

                if isinstance(input_data_class, ApiInputData):
                    return input_data_class.matches(feature_name, options, None)

                return input_data_class.matches(feature_name, options, data_access_collection)

        if cls.feature_name_equal_to_class_name(feature_name):
            return True

        if cls.feature_name_contains_class_name_as_prefix(feature_name):
            return True

        if feature_name in cls.feature_names_supported():
            return True

        return False

    @classmethod
    def feature_names_supported(cls) -> Set[str]:
        """
        The reponse could contain all or just a single feature name that are supported by this feature group.
        This way one can add custom feature names to the feature group in a very simple manner.

        This function is a convenience functionality. It is not necessary to implement this function.
        """
        return set()

    @classmethod
    def feature_name_equal_to_class_name(cls, feature_name: str) -> bool:
        """
        This functionality is useful if a feature name is equal to the class name.
        Use cases are very complex use cases like scores or very specific implementations for embeddings.
        """
        return feature_name == cls.get_class_name()

    @classmethod
    def feature_name_contains_class_name_as_prefix(cls, feature_name: str) -> bool:
        return feature_name.startswith(cls.prefix())

    @classmethod
    def get_domain(cls) -> Domain:
        """This function should return the domain for the feature group"""
        return Domain.get_default_domain()

    @classmethod
    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFrameWork]]]:
        """This function should return the compute framework rule for the feature group.
        Overwrite if necessary.

        True indicates that the feature group creator does not care about the compute framework.
        If the feature group creator wants to define the compute framework, return a set of compute frameworks.
        """
        return True

    @final
    @classmethod
    def compute_framework_definition(cls) -> Set[Type[ComputeFrameWork]]:
        """Use compute_framework_rule to set the compute framework definition."""

        rule = cls.compute_framework_rule()

        """If FG creator does not care, we allow every framework."""
        if rule is True:
            return get_all_subclasses(ComputeFrameWork)
        if isinstance(rule, bool):
            raise Exception("Compute framework rule for is not a set of compute frameworks.")
        return rule

    @classmethod
    def get_class_name(cls) -> str:
        return cls.__name__

    def __eq__(self, another: Any) -> bool:
        if isinstance(another, AbstractFeatureGroup):
            return self.get_class_name() == another.get_class_name()
        raise Exception(f"Cannot compare AbstractFeatureGroup with another type. {another} ")

    def __hash__(self) -> int:
        return hash(self.get_class_name())

    @final
    def is_root(self, options: Options, feature_name: str | FeatureName) -> bool:
        try:
            if not isinstance(feature_name, FeatureName):
                feature_name = FeatureName(feature_name)

            if self.input_features(options, feature_name) is None:
                """
                A feature could be a root feature, if it does not have any input features. 
                But a feature could be flexible depending on the options.                
                """
                return True
        except NotImplementedError:
            """
            If this is not implemented, then this is a root feature.
            """
            return True
        except Exception:
            """
            If this is not None, then this might create errors due to input feature definition.
            Thus, an error just means that this is not a root feature.
            """
            pass
        return False

    @classmethod
    def prefix(cls) -> str:
        """
        This is a convention, which means we can refer to a class via this.
        """
        return f"{cls.get_class_name()}_"
