from collections import defaultdict
from copy import deepcopy
from typing import Any, Optional
from uuid import UUID
import uuid

from mloda.core.abstract_plugins.components.index.add_index_feature import create_index_feature
from mloda.core.abstract_plugins.components.index.index import Index
from mloda.core.abstract_plugins.components.input_data.api.api_input_data_collection import (
    ApiInputDataCollection,
)
from mloda.core.abstract_plugins.components.plugin_option.plugin_collector import PluginCollector
from mloda.core.filter.global_filter import GlobalFilter
from mloda.core.prepare.accessible_plugins import PreFilterPlugins
from mloda.core.abstract_plugins.components.feature_name import FeatureName
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.compute_framework import ComputeFramework
from mloda.core.prepare.execution_plan import ExecutionPlan
from mloda.core.prepare.graph.build_graph import BuildGraph
from mloda.core.prepare.resolve_graph import ResolveGraph
from mloda.core.runtime.run import ExecutionOrchestrator
from mloda.core.prepare.identify_feature_group import IdentifyFeatureGroupClass
from mloda.core.runtime.flight.runner_flight_server import ParallelRunnerFlightServer
from mloda.core.abstract_plugins.feature_group import FeatureGroup, format_feature_group_class
from mloda.core.abstract_plugins.components.feature import Feature
from mloda.core.abstract_plugins.components.feature_collection import Features
from mloda.core.abstract_plugins.components.options import Options
from mloda.core.abstract_plugins.components.link import JoinType, Link
from mloda.core.abstract_plugins.components.validators.link_validator import LinkValidator


class Engine:
    def __init__(
        self,
        features: Features,
        compute_frameworks: set[type[ComputeFramework]],
        links: Optional[set[Link]],
        data_access_collection: Optional[DataAccessCollection] = None,
        global_filter: Optional[GlobalFilter] = None,
        api_input_data_collection: Optional[ApiInputDataCollection] = None,
        plugin_collector: Optional[PluginCollector] = None,
        column_ordering: Optional[str] = None,
    ) -> None:
        # setup variables which track the primary sources and the compute platforms
        self.feature_group_collection: dict[type[FeatureGroup], set[Feature]] = defaultdict(set)

        # use global filters
        self.global_filter = global_filter

        # Tracks feature relation to its parents
        self.feature_link_parents: dict[UUID, set[UUID]] = defaultdict(set)

        # get accessible feature groups and their compute platforms
        self.accessible_plugins = PreFilterPlugins(compute_frameworks, plugin_collector).get_accessible_plugins()
        # get links
        LinkValidator.validate_links(links)
        self.links = links

        # set api input collection if relevant
        self.api_input_data_collection = api_input_data_collection

        # TODO
        self.plugin_collector = plugin_collector
        self.copy_compute_frameworks = deepcopy(compute_frameworks)

        self.data_access_collection = data_access_collection
        self.column_ordering = column_ordering
        self.execution_planner = self.create_setup_execution_plan(features)
        self.tfs_connection_map = self._resolve_tfs_connection_map()

    def _resolve_tfs_connection_map(self) -> dict[type[ComputeFramework], Any]:
        """Resolve a connection per TFS destination framework at setup time.

        Walks the planned TFS steps once and asks each destination framework class to pick
        a matching connection from the DataAccessCollection. The resulting dict is consumed
        on the run path by ComputeFrameworkExecutor without any further DAC scanning.
        """
        connection_map: dict[type[ComputeFramework], Any] = {}
        if self.data_access_collection is None:
            return connection_map
        for tfs in self.execution_planner.tfs_collecion:
            cfw_class = tfs.to_framework
            if cfw_class in connection_map:
                continue
            conn = cfw_class.pick_connection_from_dac(self.data_access_collection)
            if conn is not None:
                connection_map[cfw_class] = conn
        return connection_map

    def compute(self, flight_server: Optional[ParallelRunnerFlightServer] = None) -> ExecutionOrchestrator:
        execution_plan_copy = deepcopy(self.execution_planner)
        orchestrator = ExecutionOrchestrator(
            execution_plan_copy,
            flight_server,
            column_ordering=self.column_ordering,
            tfs_connection_map=self.tfs_connection_map,
        )
        if isinstance(orchestrator, ExecutionOrchestrator):
            return orchestrator
        raise ValueError("ExecutionOrchestrator setup failed.")

    def create_setup_execution_plan(self, features: Features) -> ExecutionPlan:
        self.setup_features_recursion(features)

        graph_builder = BuildGraph(self.feature_link_parents, self.feature_group_collection)
        graph_builder.build_graph_from_feature_links()
        graph = graph_builder.graph

        # resolve graph into a queue
        resolver = ResolveGraph(graph, self.links)
        resolver.create_initial_queue()

        resolver.set_nodes_per_feature_group()

        planned_queue = resolver.resolve_links()

        planned_queue = resolver.resolver_compute_framework.links(
            planned_queue, resolver.resolver_links.get_link_trekker()
        )

        execution_planner = ExecutionPlan(self.global_filter, self.api_input_data_collection)
        execution_planner.create_execution_plan(planned_queue, graph, resolver.resolver_links.get_link_trekker())
        return execution_planner

    def setup_features_recursion(self, features: Features) -> None:
        for feature in features:
            self._process_feature(feature, features)

    def _process_feature(self, feature: Feature, features: Features) -> None:
        """Processes a single feature by delegating tasks to helper methods."""

        # Drop feature chainer parser properties
        feature_group_class, compute_frameworks = self._identify_feature_group_and_frameworks(feature)
        feature_group = feature_group_class()

        self._set_feature_name(feature, feature_group)
        self._set_compute_framework_and_data_type(feature, compute_frameworks, feature_group_class)

        added = self.add_feature_to_collection(feature_group_class, feature, features.child_uuid)

        if added:
            parent_domain = feature.domain.name if feature.domain else None
            self._handle_input_features_recursion(
                feature_group_class, feature.uuid, feature.options, feature.name, parent_domain=parent_domain
            )

        if self.global_filter:
            self._add_filter_feature(feature_group_class, feature_group, feature, features)

        if feature_group.index_columns():
            self._add_index_feature(feature_group_class, feature_group, feature, features)
        elif self.links:
            self._add_index_feature_from_links(feature_group_class, feature_group, feature, features)

    def _set_feature_name(self, feature: Feature, feature_group: FeatureGroup) -> None:
        """Sets the feature name using the feature group's logic."""
        feature.name = feature_group.set_feature_name(feature.options, feature.name)

    def _set_compute_framework_and_data_type(
        self,
        feature: Feature,
        compute_frameworks: set[type[ComputeFramework]],
        feature_group_class: type[FeatureGroup],
    ) -> None:
        """Sets the compute framework and data type for the feature."""
        feature = self.set_compute_framework(feature, compute_frameworks)
        feature.data_type = self.set_data_type(feature, feature_group_class)

    def _identify_feature_group_and_frameworks(
        self, feature: Feature
    ) -> tuple[type[FeatureGroup], set[type[ComputeFramework]]]:
        """Identifies the feature group class and compute frameworks for a given feature."""
        identifier = IdentifyFeatureGroupClass(
            feature, self.accessible_plugins, self.links, self.data_access_collection
        )
        return identifier.get()

    def _add_index_feature(
        self,
        feature_group_class: type[FeatureGroup],
        feature_group: FeatureGroup,
        feature: Feature,
        features: Features,
    ) -> None:
        indexes = feature_group_class.index_columns()
        if indexes is None:
            raise ValueError(f"Feature group {format_feature_group_class(feature_group_class)} has no indexes defined.")

        if self.links is None:
            return

        for index in indexes:
            self._process_index_feature(feature_group_class, feature_group, feature, features, index)

    def _add_index_feature_from_links(
        self,
        feature_group_class: type[FeatureGroup],
        feature_group: FeatureGroup,
        feature: Feature,
        features: Features,
    ) -> None:
        """Adds index features derived from JoinSpec links when the feature group has no index_columns defined."""
        if self.links is None:
            return

        feature_name_str = feature.name

        for link in self.links:
            if link.jointype in (JoinType.APPEND, JoinType.UNION):
                continue
            if link.left_feature_group == feature_group_class and feature_name_str in link.left_index.index:
                return
            if link.right_feature_group == feature_group_class and feature_name_str in link.right_index.index:
                return

        for link in self.links:
            if link.jointype in (JoinType.APPEND, JoinType.UNION):
                continue
            if link.left_feature_group == feature_group_class:
                self._create_and_add_index_feature(
                    feature_group_class, feature_group, feature, features, link.left_index
                )
            if link.right_feature_group == feature_group_class:
                self._create_and_add_index_feature(
                    feature_group_class, feature_group, feature, features, link.right_index
                )

    def _process_index_feature(
        self,
        feature_group_class: type[FeatureGroup],
        feature_group: FeatureGroup,
        feature: Feature,
        features: Features,
        index: Index,
    ) -> None:
        """Processes the index feature for both left and right links."""
        if self.links is None:
            return

        for link in self.links:
            if link.left_feature_group == feature_group_class and link.left_index == index:
                self._create_and_add_index_feature(feature_group_class, feature_group, feature, features, index)

            if link.right_feature_group == feature_group_class and link.right_index == index:
                self._create_and_add_index_feature(feature_group_class, feature_group, feature, features, index)

    def _create_and_add_index_feature(
        self,
        feature_group_class: type[FeatureGroup],
        feature_group: FeatureGroup,
        feature: Feature,
        features: Features,
        index: Index,
    ) -> None:
        """Creates and adds a new index feature to the collection."""
        new_index_feature = create_index_feature(index, feature_group, feature)
        self.add_feature_to_collection(feature_group_class, new_index_feature, features.child_uuid, True)

    def _add_filter_feature(
        self,
        feature_group_class: type[FeatureGroup],
        feature_group: FeatureGroup,
        feature: Feature,
        features: Features,
    ) -> None:
        if self.global_filter:
            matched_filters = self.global_filter.identify_matched_filters(
                feature_group_class, feature, self.data_access_collection
            )

            for match in matched_filters:
                match.filter_feature.name = feature_group.set_feature_name(
                    match.filter_feature.options, match.filter_feature.name
                )
                # We assign a new UUID to the filter feature to ensure
                # it is treated as a separate instance from the original filter feature
                match.filter_feature.uuid = uuid.uuid4()
                self.global_filter.add_filter_to_collection(feature_group_class, feature.name, match)
                self.add_feature_to_collection(feature_group_class, match.filter_feature, features.child_uuid)

    def add_feature_link_to_links(self, feature: Feature) -> None:
        """With this functionality, we can add links with a feature instead via mloda API."""

        if feature.link is None:
            return

        if self.links is None:
            self.links = {feature.link}
        else:
            self.links.add(feature.link)

    def add_feature_to_collection(
        self,
        feature_group_class: type[FeatureGroup],
        feature: Feature,
        child_uuid: Optional[UUID],
        if_index_feature: bool = False,
    ) -> bool:
        feature_collection = self.feature_group_collection[feature_group_class]

        if feature not in feature_collection:
            self.add_feature_link_to_links(feature)

            self.feature_link_parents[feature.uuid] = set()
            feature_collection.add(feature)
            return True

        if child_uuid:
            # Find the wanted_uuid in feature_collection
            wanted_uuid = next((f.uuid for f in feature_collection if feature == f), None)

            if wanted_uuid is not None:
                self._update_feature_link_parents(child_uuid, feature.uuid, wanted_uuid, if_index_feature)

        return False

    def _update_feature_link_parents(
        self, child_uuid: UUID, original_uuid: UUID, wanted_uuid: UUID, if_index_feature: bool
    ) -> None:
        """Updates the feature link parents based on whether it's an index feature or not."""
        if not if_index_feature:
            if original_uuid in self.feature_link_parents[child_uuid]:
                self.feature_link_parents[child_uuid].remove(original_uuid)
            self.feature_link_parents[child_uuid].add(wanted_uuid)
        else:
            self.feature_link_parents[child_uuid].add(wanted_uuid)

    def _handle_input_features_recursion(
        self,
        feature_group_class: type[FeatureGroup],
        uuid: UUID,
        options: Options,
        feature_name: FeatureName,
        parent_domain: Optional[str] = None,
    ) -> None:
        """Handles recursion for input features of a feature group."""
        feature_group = feature_group_class()

        # options = deepcopy(options)

        try:
            input_features = feature_group.input_features(options, feature_name)
        except NotImplementedError:  # This means, it is a root feature.
            input_features = None

        if input_features:
            features = Features(
                list(input_features), child_options=options, child_uuid=uuid, parent_domain=parent_domain
            )
            if features.child_uuid is None:
                raise ValueError(f"Features {features} has no parent uuid although it should have one.")
            self.feature_link_parents[features.child_uuid] = features.parent_uuids
            self.setup_features_recursion(features)

    def set_compute_framework(self, feature: Feature, compute_frameworks: set[type[ComputeFramework]]) -> Feature:
        """
        This function leads to that the feature has always a compute framework set!
        """
        if feature.compute_frameworks:
            if feature.get_compute_framework() not in compute_frameworks:
                raise ValueError(
                    f"Feature {feature.name} does not support compute framework {feature.compute_frameworks}."
                )
        else:
            feature.compute_frameworks = compute_frameworks
        return feature

    def set_data_type(self, feature: Feature, feature_group_class: type[FeatureGroup]) -> Optional[DataType]:
        fg_data_type = feature_group_class.return_data_type_rule(feature)
        if feature.data_type and fg_data_type:
            if feature.data_type != fg_data_type:
                raise ValueError(
                    f"Feature {feature.name} has a data type mismatch with feature group {feature_group_class}."
                )
            return fg_data_type
        return fg_data_type or feature.data_type
