from abc import ABC
from typing import Any, Optional, final
from uuid import UUID, uuid4
from mloda.core.abstract_plugins.components.data_types import DataType
from mloda.core.abstract_plugins.components.framework_transformer.cfw_transformer import (
    ComputeFrameworkTransformer,
)
from mloda.core.abstract_plugins.components.merge.base_merge_engine import BaseMergeEngine
import pyarrow as pa

from mloda.core.abstract_plugins.function_extender import (
    Extender,
    ExtenderHook,
    _CompositeExtender,
)
from mloda.core.abstract_plugins.components.feature_name import FeatureName
from mloda.core.abstract_plugins.components.parallelization_modes import ParallelizationMode
from mloda.core.filter.filter_engine import BaseFilterEngine
from mloda.core.abstract_plugins.components.mask.base_mask_engine import BaseMaskEngine
from mloda.core.runtime.flight.flight_server import FlightServer


class ComputeFramework(ABC):
    """
    Documentation ComputeFramework:

    This class is used to define the compute framework.

    A compute framework must be in the namespace of the python project and must inherit from this class.
    This way, we can run feature computation on multiple frameworks.
    Usecases:
    - Online and Offline computation
    - Testing
    - Migrations from one framework to another

    As you can have multiple compute frameworks, you can define a feature group to be computed on multiple compute frameworks.
    However, in the space of a run, it is necessary that there is only exactly one way to compute a feature of feature group,
    we can limit the compute framework from three sides:
    1) feature definition - only one compute framework can be set for one feature
    2) feature group definition - compute frameworks that can be used for the feature group
    3) from api request side - limits compute frameworks which can be used

    Of course, for migrations, we have allow multiple definitions of the same feature group on different compute frameworks.
    This usecase however is currently not supported, as one could just run the module twice and compare the result datasets for now.
    """

    def __init__(
        self,
        mode: ParallelizationMode,
        children_if_root: frozenset[UUID],
        uuid: UUID = uuid4(),
        function_extender: Optional[set[Extender]] = None,
    ) -> None:
        """This class is initialized step execution."""
        self.mode = mode
        self.data: Any = None
        self.children_if_root = children_if_root
        self.already_calculated_children_tracker: set[UUID] = set()
        self.column_names: set[str] = set()
        self.function_extender = function_extender if function_extender is not None else set()

        self.uuid = uuid

        self.transformer = ComputeFrameworkTransformer()

        # collects all datasets which were created based on this feature group except if it is the !final! feature
        self.object_ids: list[str] = []

        # connection object for frameworks that need persistent connections (e.g., DuckDB, Spark)
        self.framework_connection_object: Optional[Any] = None

    @classmethod
    def expected_data_framework(cls) -> Any:
        """
        This function should return the expected data framework for the compute framework.
        However, we only need to set it if we really want to be sure that the datatype is correct.
        """
        return None

    @classmethod
    def filter_engine(cls) -> type[BaseFilterEngine]:
        """
        This function should return the filtered data.
        The BaseFilterEngine should be overwritten by the appropriate ComputeFramework if needed
        """
        raise NotImplementedError

    @classmethod
    def mask_engine(cls) -> type[BaseMaskEngine] | None:
        """Return the mask engine for building boolean masks inside calculate_feature().

        Override in framework plugins that support masking.
        Returns None by default (mask support is optional).
        """
        return None

    def transform(
        self,
        data: Any,
        feature_names: set[str],
    ) -> Any:
        """This function should be used to transform the data.
        The idea here is that we can transform the data to a common format.
        At the end, this format is represented by the expected_data_framework
        and thus, defined by the compute framework abstraction.

        If you wish not to transform the data, you can leave this.
        This is not relevant if you stay in one compute framework as you don t need to switch it.
        """

        transformed_data = self.apply_compute_framework_transformer(data)
        if transformed_data is not None:
            return transformed_data
        return data

    @final
    def apply_compute_framework_transformer(self, data: Any) -> Any:
        """
        This part is also refactored to be more readable.

        The part, where we add single columns etc is not done yet.
        """
        _from_fw = type(data)
        _to_fw = self.expected_data_framework()
        transformer_cls = self.transformer.transformer_map.get((_from_fw, _to_fw), None)
        if transformer_cls is not None:
            return transformer_cls.transform(_from_fw, _to_fw, data, self.framework_connection_object)

        return None

    def select_data_by_column_names(
        self, data: Any, selected_feature_names: set[FeatureName], column_ordering: Optional[str] = None
    ) -> Any:
        """
        If you only want to store the requested features, implement this functionality depending on your framework.

        e.g. if you are using pyarrow, you can use the following code:
        return data.select([f.name for f in selected_feature_names])
        """
        return data

    @classmethod
    def merge_engine(cls) -> type[BaseMergeEngine]:
        """
        This function should return a subclass of the BaseMergeEngine.
        With this, we can merge data from the same compute framework.

        This implementation is optional.
        """
        raise NotImplementedError(f"Merge functionality is for this compute framework not implemented {cls.__name__}.")

    def set_framework_connection_object(self, framework_connection_object: Optional[Any] = None) -> None:
        """
        Some compute frameworks (e.g., DuckDB, Spark) require sharing their connection
        with merge engines to ensure data consistency. Override this method in
        subclasses that need to provide a connection object.
        """
        self.framework_connection_object = None

    @classmethod
    def _connection_matches(cls, conn: Any) -> bool:
        """Return True if `conn` is a connection this framework can adopt.

        Default returns False. SQL frameworks override with a single isinstance check;
        the surrounding fail-fast policy lives in `pick_connection_from_dac`.
        """
        return False

    @classmethod
    def pick_connection_from_dac(cls, data_access_collection: Any) -> Optional[Any]:
        """Pick the matching connection for this framework from the DataAccessCollection.

        Called at Engine setup (once per session), not on the per-request run path.
        Returns None if no connection matches, the single match if exactly one does,
        and raises ValueError if multiple connections match - the caller has to
        disambiguate explicitly rather than let a non-deterministic set iteration decide.
        """
        if data_access_collection is None:
            return None
        matches = [
            conn for conn in data_access_collection.initialized_connection_objects if cls._connection_matches(conn)
        ]
        if not matches:
            return None
        if len(matches) > 1:
            raise ValueError(
                f"{cls.__name__}.pick_connection_from_dac: DataAccessCollection contains "
                f"{len(matches)} connections matching this framework; expected at most one."
            )
        return matches[0]

    @final
    def get_framework_connection_object(self) -> Any:
        """This method retrieves the connection object set by `set_framework_connection_object`."""
        return self.framework_connection_object

    def set_column_names(self) -> None:
        self.column_names = self._extract_column_names(self.data)

    @staticmethod
    def is_available() -> bool:
        """
        Check if the compute framework's dependencies are available.
        Subclasses should override this to check for their specific dependencies if needed.

        Returns:
            bool: True if all required dependencies are installed, False otherwise
        """
        return True  # Default implementation assumes no external dependencies

    @classmethod
    def supported_parallelization_modes(cls) -> set[ParallelizationMode]:
        """Returns parallelization modes this framework supports. Override to restrict."""
        return {ParallelizationMode.SYNC, ParallelizationMode.THREADING, ParallelizationMode.MULTIPROCESSING}

    @final
    def run_calculation(
        self, feature_group: Any, features: Any, location: str | None, data: Optional[Any] = None
    ) -> Optional[Any]:
        # case multiprocessing or case base api input feature
        if data is not None:
            self.data = data

        self.run_validate_input_features(feature_group, features)

        # every one does this
        features = self.set_filter_engine(features)
        features = self.set_mask_engine(features)
        data = self.run_calculate_feature(feature_group, features)
        data = self.run_final_filter(data, features, feature_group)

        names = features.get_all_names()

        if not isinstance(data, self.expected_data_framework()):
            # if data is not in the expected data framework, we need to transform it and for this, we may need the framework connection object
            self.set_framework_connection_object(features.get_options_key(feature_group.get_class_name()))
            self.data = self.transform(data, names)
        else:
            self.data = data

        self.set_column_names()

        self.run_validate_output_features(feature_group, features)

        # case threading/sync
        if not location:
            return None

        # case multiprocessing
        # return data to be used in next step of this framework in this process
        if len(self.children_if_root) > len(self.already_calculated_children_tracker) + len(features.features):
            return self.data

        # upload finished dataset to flight server
        self.data = self.upload_finished_data(location)
        return self.data

    @final
    def run_final_filter(self, data: Any, features: Any, feature_group: Any = None) -> Any:
        if features.filter_engine is None:
            return data

        if feature_group is not None:
            fg_preference = feature_group.final_filters()
            if fg_preference is False:
                return data
            if fg_preference is True:
                self._validate_filter_columns(data, features, feature_group)
                filter_engine = features.filter_engine()
                return filter_engine.apply_filters(data, features)

        try:
            if features.filter_engine().final_filters() is False:
                return data
        except NotImplementedError:
            return data

        self._validate_filter_columns(data, features, feature_group)
        filter_engine = features.filter_engine()
        result = filter_engine.apply_filters(data, features)

        return result

    def _validate_filter_columns(self, data: Any, features: Any, feature_group: Any) -> None:
        """Validate that filter columns exist and have compatible dtypes before row elimination.

        Called on every path where the framework will apply filters, regardless of
        whether the FeatureGroup or the FilterEngine requested elimination.
        Uses filter_engine.applicable_filters() to determine which columns will
        be processed, so this check and apply_filters() share the same gate.
        """
        filter_engine_cls = features.filter_engine()
        applicable = filter_engine_cls.applicable_filters(features)
        if not applicable:
            return

        data_columns = self._extract_column_names(data)
        fg_name = feature_group.get_class_name() if feature_group is not None else "Unknown"
        for single_filter in applicable:
            col = single_filter.filter_feature.name
            if col not in data_columns:
                raise ValueError(
                    f"{fg_name} output is missing filter column '{col}' "
                    f"required for row elimination. "
                    f"This is a bug in the FeatureGroup, not in your filter. "
                    f"Available columns: {sorted(data_columns)}"
                )
            self._validate_filter_dtype_compatibility(data, single_filter, fg_name)

    def _validate_filter_dtype_compatibility(self, data: Any, single_filter: Any, fg_name: str) -> None:
        """Check that the filter column dtype is compatible with the filter value type."""
        col = single_filter.filter_feature.name
        col_dtype = self._extract_column_dtype(data, col)
        if col_dtype is None:
            return

        filter_values = self._collect_filter_values(single_filter)
        if not filter_values:
            return

        col_dtype_str = str(col_dtype).lower()
        col_is_string = self._is_string_dtype(col_dtype_str)
        col_is_numeric = self._is_numeric_dtype(col_dtype_str)

        for val in filter_values:
            if isinstance(val, str) and col_is_numeric:
                raise ValueError(
                    f"Filter on '{col}' expects string comparison but column has type {col_dtype}. "
                    f"This usually means the FeatureGroup '{fg_name}' modified the filter "
                    f"column's type during calculation."
                )
            if isinstance(val, (int, float)) and col_is_string:
                raise ValueError(
                    f"Filter on '{col}' expects numeric comparison but column has type {col_dtype}. "
                    f"This usually means the FeatureGroup '{fg_name}' modified the filter "
                    f"column's type during calculation."
                )

    @staticmethod
    def _collect_filter_values(single_filter: Any) -> list[Any]:
        """Extract all comparable values from a filter's parameters."""
        values: list[Any] = []
        param = single_filter.parameter
        if param.value is not None:
            values.append(param.value)
        if param.values is not None:
            values.extend(param.values)
        if param.min_value is not None:
            values.append(param.min_value)
        if param.max_value is not None:
            values.append(param.max_value)
        return values

    _STRING_TYPES: frozenset[str] = frozenset(
        {
            "str",
            "string",
            "utf8",
            "large_utf8",
            "large_string",
            "object",
            "varchar",
            "nvarchar",
            "text",
            "ntext",
            # Spark
            "stringtype()",
        }
    )

    _NUMERIC_TYPES: frozenset[str] = frozenset(
        {
            # Python / generic
            "int",
            "float",
            "double",
            "long",
            "short",
            "byte",
            # PyArrow / Pandas / Polars sized types
            "int8",
            "int16",
            "int32",
            "int64",
            "uint8",
            "uint16",
            "uint32",
            "uint64",
            "float16",
            "float32",
            "float64",
            # SQL types
            "integer",
            "bigint",
            "smallint",
            "tinyint",
            "hugeint",
            "ubigint",
            "uinteger",
            "usmallint",
            "utinyint",
            "uhugeint",
            "real",
            "numeric",
            "number",
            # Spark types
            "longtype()",
            "integertype()",
            "shorttype()",
            "bytetype()",
            "floattype()",
            "doubletype()",
        }
    )

    _NUMERIC_PREFIXES: tuple[str, ...] = ("decimal",)

    @staticmethod
    def _is_string_dtype(dtype_str: str) -> bool:
        """Classify whether a dtype string represents a string/text type."""
        return dtype_str in ComputeFramework._STRING_TYPES

    @staticmethod
    def _is_numeric_dtype(dtype_str: str) -> bool:
        """Classify whether a dtype string represents a numeric type."""
        if dtype_str in ComputeFramework._NUMERIC_TYPES:
            return True
        return any(dtype_str.startswith(p) for p in ComputeFramework._NUMERIC_PREFIXES)

    def _extract_column_names(self, data: Any) -> set[str]:
        """Extract column names from the data.

        Called only with the non-None output of calculate_feature.
        """
        raise NotImplementedError

    def _extract_column_dtype(self, data: Any, column_name: str) -> Optional[str]:
        """Extract the dtype of a specific column as a human-readable string.

        Returns None when the framework does not support dtype extraction.
        Subclasses should override to return dtype strings like "int64", "string", etc.
        """
        return None

    def _extract_column_data_type(self, data: Any, column_name: str) -> Optional[DataType]:
        """Resolve a column's native dtype to the unified DataType enum.

        Default: return None. Frameworks must override to convert their native
        type system (pa.DataType, pd.api.types, polars.DataType, pyiceberg types,
        pyspark types, duckdb types) directly to DataType, no string round-trip.
        Returning None means the validator skips that column (graceful no-op).
        """
        return None

    @final
    def set_filter_engine(self, features: Any) -> Any:
        """We set the filter engine of the feature set here.
        With this, we can run filters on the data during the calculate process.

        This is needed as we have different times when we want to apply filters.

        -   We can apply filters on final data sets e.g. reading csv like pandas
        -   We can apply filters on the data during read e.g. reading from a database
        -   We can apply filters on the data during the calculation process e.g. pyarrow
        """
        try:
            features.filter_engine = self.filter_engine

            if not issubclass(features.filter_engine(), BaseFilterEngine):
                raise ValueError(f"Filter engine {self.filter_engine} not supported by {self.__class__.__name__}")

        except NotImplementedError:
            pass

        return features

    @final
    def set_mask_engine(self, features: Any) -> Any:
        """Wire the mask engine onto the FeatureSet.

        The mask engine provides boolean mask primitives for use during
        calculate_feature(). It is set here so that FeatureGroups can access
        it via features.mask_engine.
        """
        engine = self.mask_engine()
        if engine is not None:
            features.mask_engine = engine
        return features

    @final
    def run_validate_input_features(self, feature_group: Any, features: Any) -> None:
        if self.data is None:
            return

        extender = self.get_function_extender(ExtenderHook.VALIDATE_INPUT_FEATURE)
        if extender is None:
            feature_group.validate_input_features(self.data, features)
        else:
            extender(feature_group.validate_input_features, self.data, features)

    @final
    def run_validate_output_features(self, feature_group: Any, features: Any) -> None:
        if self.data is None:
            return

        from mloda.core.abstract_plugins.components.validators.datatype_validator import DataTypeValidator

        DataTypeValidator.validate(
            features,
            lambda col_name: self._extract_column_data_type(self.data, col_name),
        )

        extender = self.get_function_extender(ExtenderHook.VALIDATE_OUTPUT_FEATURE)
        if extender is None:
            feature_group.validate_output_features(self.data, features)
        else:
            extender(feature_group.validate_output_features, self.data, features)

    @final
    def get_column_names(self) -> set[str]:
        return self.column_names

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

    @final
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, ComputeFramework):
            return False
        return type(self) is type(other) and self.children_if_root == other.children_if_root

    @final
    def __hash__(self) -> int:
        return hash((id(type(self)), self.children_if_root))

    def validate_expected_framework(self, location: Optional[str] = None) -> None:
        """
        Validates that the data is in the expected framework.
        Only touch is if your framework supports multiple data frameworks.
        """
        if self.expected_data_framework() is None:
            return

        if self.data is None:
            return

        # If location is a string, it means it is a uuid of the object in arrow flight.
        if isinstance(location, str) and self.data is not None:
            return

        if not isinstance(self.data, self.expected_data_framework()):
            expected = self.expected_data_framework()
            raise ValueError(
                f"Data type {type(self.data).__name__} is not supported by {self.__class__.__name__}.\n"
                f"Expected type: {expected.__name__ if expected else 'not specified'}.\n"
                "Ensure your FeatureGroup's calculate_feature returns data in the "
                "framework's expected format, or register a ComputeFrameworkTransformer "
                "to convert between formats."
            )

    @final
    def add_already_calculated_children_and_drop_if_possible(
        self, children: set[UUID], location: Optional[str] = None
    ) -> bool | frozenset[UUID]:
        # if len(self.object_ids) > 1:
        #    if location is None:
        #        raise ValueError("Location is not set")
        #    self.drop_data(set(self.object_ids[:-1]), location)

        self.already_calculated_children_tracker.update(children)

        if self.children_if_root.issubset(self.already_calculated_children_tracker):
            self.drop_last_data(location)
            return True

        if len(self.object_ids) > 0:
            return self.children_if_root

        return False

    @final
    def get_function_extender(self, wrapper_function_enum: ExtenderHook) -> Optional[Extender]:
        matching_extenders = []
        for extender in self.function_extender:
            if wrapper_function_enum in extender.wraps():
                matching_extenders.append(extender)

        if len(matching_extenders) == 0:
            return None
        if len(matching_extenders) == 1:
            return matching_extenders[0]

        sorted_extenders = sorted(matching_extenders, key=lambda e: e.priority)
        return _CompositeExtender(sorted_extenders, wrapper_function_enum)

    @final
    def run_calculate_feature(self, feature_group: Any, features: Any) -> Any:
        extender = self.get_function_extender(ExtenderHook.FEATURE_GROUP_CALCULATE_FEATURE)

        try:
            if extender is None:
                return feature_group.calculate_feature(self.data, features)
            return extender(feature_group.calculate_feature, self.data, features)
        except KeyError as e:
            # Provide helpful error message for missing columns
            self._raise_helpful_missing_column_error(feature_group, e)

    def _raise_helpful_missing_column_error(self, feature_group: Any, error: KeyError) -> None:
        """
        Raises a helpful ValueError suggesting the KeyError might be due to missing Links.
        """
        feature_name = feature_group.get_class_name()
        error_str = str(error)

        error_message = f"""
Feature '{feature_name}' failed with a KeyError: {error_str}

This might be caused by missing Links when your feature has multiple dependencies.

When a feature depends on multiple input features, you must provide explicit Links to specify
how to merge them. Without Links, the framework cannot determine how to combine the data.

Option 1: Explicit JoinSpec (works with any feature group):
    from mloda.user import Link, JoinSpec

    links = {{
        Link.inner(
            JoinSpec(RootFeatureA, "shared_column"),
            JoinSpec(RootFeatureB, "shared_column"),
        )
    }}

Option 2: Shorthand via index_columns() (requires feature groups to define index_columns()):
    from mloda.user import Link

    links = {{
        Link.inner_on(RootFeatureA, RootFeatureB)
    }}

    mloda.run_all(
        features=[Feature.int32_of("{feature_name}")],
        links=links,
        ...
    )

Available join types:
- Link.inner(left, right)    - Keep only matching rows from both sides
- Link.left(left, right)     - Keep all rows from left, matching from right
- Link.right(left, right)    - Keep all rows from right, matching from left
- Link.outer(left, right)    - Keep all rows from both sides
- Link.inner_on(left, right) - Shorthand using index_columns() definitions
"""
        raise ValueError(error_message.strip())

    @final
    def set_data(self, data: Any) -> None:
        self.data = data

    @final
    def get_data(self) -> Any:
        return self.data

    @final
    def get_object_ids(self) -> list[str]:
        return self.object_ids

    @final
    def upload_finished_data(self, location: str) -> str:
        """Uploads table by using its own cfw uuid."""
        return self.upload_table(location, self.uuid)

    @final
    def upload_table(self, location: str, object_id: Optional[UUID] = None) -> str:
        if object_id is None:
            object_id = uuid4()
        _object_id = str(object_id)

        # transform to pa.Table for better parquet support
        if not isinstance(self.data, pa.Table):
            _from_fw = type(self.data)
            _to_fw = pa.Table

            transformer_cls = self.transformer.transformer_map.get((_from_fw, _to_fw), None)
            if transformer_cls is not None:
                self.data = transformer_cls.transform(_from_fw, _to_fw, self.data, self.framework_connection_object)

        FlightServer.upload_table(location, self.data, _object_id)
        self.object_ids.append(_object_id)
        return _object_id

    @final
    @classmethod
    def convert_flight_server_data_back(cls, data: Any, transformer: ComputeFrameworkTransformer) -> Any:
        if not isinstance(data, pa.Table):
            return data
        if isinstance(data, cls.expected_data_framework()):
            return data

        _from_fw = type(data)
        _to_fw = cls.expected_data_framework()

        transformer_cls = transformer.transformer_map.get((_from_fw, _to_fw), None)
        if transformer_cls is not None:
            return transformer_cls.transform(_from_fw, _to_fw, data, None)

        raise ValueError(
            f"Conversion from {type(data)} to {cls.expected_data_framework()} is not supported. This can happen when a FlightServer was used."
        )

    @final
    def drop_data(self, table_keys: set[str], location: str) -> None:
        FlightServer.drop_tables(location, table_keys)

    @final
    def drop_last_data(self, location: Optional[str] = None) -> None:
        if isinstance(self.data, str) and location:
            self.drop_data({self.data}, location)

        self.data = None

    @final
    def get_uuid(self) -> UUID:
        if self.uuid is None:
            raise ValueError("UUID is not set")
        return self.uuid

    @final
    def identify_naming_convention(
        self,
        selected_feature_names: set[FeatureName],
        column_names: set[str],
        ordering: Optional[str] = None,
    ) -> set[str] | list[str]:
        """
        Identifies columns that match feature names or follow the naming convention pattern.

        This method supports the multiple result columns pattern, where a feature group can
        return multiple related columns using the naming convention 'feature_name~column_suffix'.

        Args:
            selected_feature_names: A set of FeatureName objects representing the requested features
            column_names: A set of strings representing the available column names in the data
            ordering: Optional ordering mode - None for Set, "alphabetical" or "request_order" for List

        Returns:
            A set of column names (default) or list if ordering is specified

        Raises:
            ValueError: If no matching columns are found or invalid ordering value

        Example:
            If selected_feature_names contains 'Temperature' and column_names contains
            'Temperature~mean', 'Temperature~max', 'Temperature~min', this method will return
            all three columns as they follow the naming convention.
        """
        if ordering is not None and ordering not in ("alphabetical", "request_order"):
            raise ValueError(f"Invalid ordering value: '{ordering}'. Must be None, 'alphabetical', or 'request_order'.")

        feature_name_strings = {str(f) for f in selected_feature_names}
        _selected_feature_names: set[str] = set()

        for col in column_names:
            for feature_name in feature_name_strings:
                if col == feature_name:
                    _selected_feature_names.add(col)
                    continue

                if col.startswith(f"{feature_name}~"):
                    _selected_feature_names.add(col)

        if not _selected_feature_names:
            raise ValueError(
                f"No columns found that match feature names {feature_name_strings} or follow the naming convention 'feature_name~column_name'"
            )

        if ordering is None:
            return _selected_feature_names

        if ordering == "alphabetical":
            return sorted(_selected_feature_names)

        # ordering == "request_order"
        result: list[str] = []
        for feature in selected_feature_names:
            feature_name_str = str(feature)
            matching_cols = []
            for col in _selected_feature_names:
                if col == feature_name_str or col.startswith(f"{feature_name_str}~"):
                    matching_cols.append(col)
            matching_cols.sort()
            result.extend(matching_cols)
        return result
