from __future__ import annotations

from typing import Any, Callable, Dict, 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_group_version import FeatureGroupVersion
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.components.match_data.match_data import MatchData
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 description(cls) -> str:
        """
        Returns a description for this feature group.

        The method returns the class's own docstring if it has been overridden from
        the base class's docstring. Otherwise, it falls back to the class name.
        This behavior allows subclasses to easily customize their description.
        """
        base_doc = (AbstractFeatureGroup.__doc__ or "").strip()
        current_doc = (cls.__doc__ or "").strip()

        if current_doc and current_doc != base_doc:
            return current_doc
        return cls.get_class_name()

    @classmethod
    def version(cls) -> str:
        """
        Returns a composite version identifier for this feature group.

        The version identifier is generated by combining:
          - the version of the 'mloda' package (retrieved from package metadata),
          - the module name where the feature group is defined, and
          - a SHA-256 hash of the feature group class's source code.

        This composite identifier uniquely represents the implementation state of the feature group,
        making it easier to detect changes, manage compatibility, and debug issues.

        If you need to change the version of the feature group, you can do so by subclassing
        FeatureGroupVersion and overriding the version method. This allows you to create a new version system.
        """
        return FeatureGroupVersion.version(cls)

    @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:
        """
        Allows modification of the feature name based on configuration.

        This method provides a hook for subclasses to modify the feature name based on the
        provided options and the initial feature name. The default implementation simply
        returns the original feature name, but subclasses can override this method to implement
        custom logic for feature name modification.
        """
        return feature_name

    @staticmethod
    def apply_naming_convention(
        result: Any, feature_name: str, suffix_generator: Optional[Callable[[int], str]] = None
    ) -> Dict[str, Any]:
        """
        Applies naming convention to multi-column results.

        For 2D arrays with multiple columns, creates a dictionary mapping
        column names to column data using the pattern: feature_name~0, feature_name~1, etc.
        """
        if hasattr(result, "shape") and len(result.shape) > 1 and result.shape[1] > 1:
            output = {}
            for col_idx in range(result.shape[1]):
                suffix = suffix_generator(col_idx) if suffix_generator else str(col_idx)
                col_name = f"{feature_name}~{suffix}"
                output[col_name] = result[:, col_idx]
            return output
        return {}

    @staticmethod
    def get_column_base_feature(column_name: str) -> str:
        """
        Extracts the base feature name from a column name by stripping the ~N suffix.

        Returns the original name if no ~N suffix exists.
        """
        return column_name.split("~")[0]

    @staticmethod
    def expand_feature_columns(feature_name: str, num_columns: int) -> List[str]:
        """
        Generates a list of column names with ~N suffixes.

        Returns a list of column names following the pattern: ["feature~0", "feature~1", ...]
        """
        return [f"{feature_name}~{i}" for i in range(num_columns)]

    @staticmethod
    def resolve_multi_column_feature(feature_name: str, available_columns: Set[str]) -> List[str]:
        """
        Resolves a feature name to its corresponding column(s) in the available columns.

        If the exact feature_name exists in available_columns, returns [feature_name].
        Otherwise, finds all columns matching the pattern {feature_name}~* and returns them sorted.
        If no matches found, returns [feature_name] (caller will handle the error).
        """
        if feature_name in available_columns:
            return [feature_name]

        matching_columns = [col for col in available_columns if col.startswith(f"{feature_name}~")]

        if matching_columns:
            return sorted(matching_columns)

        return [feature_name]

    @classmethod
    def return_data_type_rule(cls, feature: Feature) -> Optional[DataType]:
        """
        Specifies a fixed return data type for this feature group, if applicable.

        If this feature group always returns a specific data type, this method should
        return that data type. Otherwise, it should return None, indicating that the
        data type is not fixed and may vary depending on the input or computation.
        """
        return None

    def input_features(self, options: Options, feature_name: FeatureName) -> Optional[Set[Feature]]:
        """
        Defines the input features required by this feature group.

        If this feature depends on other features as input, this method should return a set
        containing those features. If it does not depend on any other features (i.e., it is
        a root feature), it should return None.

        The specific input features may depend on the provided options and the feature name.
        """
        raise NotImplementedError

    @classmethod
    def index_columns(cls) -> Optional[List[Index]]:
        """
        Specifies the index columns used for merging or joining data.

        This method should return a list of Index objects representing the columns to be
        used as indices for merging or joining dataframes. The indices can be defined by
        name, by user-set index, or by the features themselves.

        This method also provides an opportunity to validate given indices against the
        aforementioned sources, if implemented.
        """
        return None

    @classmethod
    def supports_index(cls, index: Index) -> Optional[bool]:
        """
        Indicates whether this feature group supports the given index.

        This method should return True if the feature group supports the given index,
        False otherwise. If the feature group does not have any specific index
        requirements, it should return None.
        """
        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 _matches_input_data(
        cls, feature_name: str, options: Options, data_access_collection: Optional[DataAccessCollection]
    ) -> bool:
        """
        Helper function to check if the input data matches.
        """
        input_data_class = cls.input_data()

        if input_data_class is None:
            return False

        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)

    @classmethod
    def match_feature_group_criteria(
        cls,
        feature_name: Union[FeatureName, str],
        options: Options,
        data_access_collection: Optional[DataAccessCollection] = None,
    ) -> bool:
        """
        Determines whether this feature group matches the given criteria.

        This method returns True if the feature_name matches the criteria defined by this
        feature group, and False otherwise. The criteria may include the feature name,
        options, and data access collection.

        The if statement contains the rules. Each case has different use cases.
        You can disallow them by 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_and_matches_input_data(feature_name, options, data_access_collection):
            return True

        if cls._matches_data(feature_name, options, data_access_collection):
            return True

        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]:
        """
        Returns a set of feature names that are explicitly supported by this feature group.

        This method provides a way to specify a set of feature names that this feature group
        is designed to handle. It can be used to add custom feature names to the feature
        group in a 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:
        """
        Checks if the given feature name is equal to the class name of this feature group.

        This functionality is useful in cases where the feature name directly corresponds
        to the class name, such as with 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:
        """
        Checks if the given feature name starts with the class name of this feature group
        as a prefix.
        """
        return feature_name.startswith(cls.prefix())

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

    @classmethod
    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFrameWork]]]:
        """
        Defines the rule for determining the compute framework to use for this feature group.

        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]]:
        """
        Determines the set of compute frameworks supported by this feature group based on the
        `compute_framework_rule`.
        """

        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

    @final
    @classmethod
    def get_class_name(cls) -> str:
        """
        Returns the name of the class.
        """
        return cls.__name__

    def __eq__(self, another: Any) -> bool:
        """
        Checks if this feature group is equal to another object.
        """
        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:
        """
        Returns the hash code for this feature group.
        """
        return hash(self.get_class_name())

    @final
    def is_root(self, options: Options, feature_name: str | FeatureName) -> bool:
        """
        Determines whether this feature is a root feature (i.e., does not depend on any
        other features).
        """
        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:
        """
        Returns the prefix used for feature names associated with this feature group.

        This is a convention, which means we can refer to a class via this.
        """
        return f"{cls.get_class_name()}_"

    @classmethod
    def _is_root_and_matches_input_data(
        cls, feature_name: str, options: Options, data_access_collection: Optional[DataAccessCollection]
    ) -> bool:
        """
        Checks if the feature group is a root and matches input data.
        """
        return cls().is_root(options, feature_name) and cls._matches_input_data(
            feature_name, options, data_access_collection
        )

    @final
    @classmethod
    def _matches_data(
        cls, feature_name: str, options: Options, data_access_collection: Optional[DataAccessCollection]
    ) -> bool:
        """
        This functionality is for matching data, when a data access is necessary.
        This is relevant for compute frameworks which need a connection object.

        To be used, create a class like this:

        class MyMatchData(AbstractFeatureGroup, MatchData):
            ...

        and then create the function match_data_access.
        """

        if not issubclass(cls, MatchData):
            return False

        return cls.matches(feature_name, options, data_access_collection)
