from typing import Any, Optional
from uuid import UUID, uuid4
from mloda.core.abstract_plugins.components.base_artifact import BaseArtifact
from mloda.core.abstract_plugins.components.input_data.api.base_api_data import BaseApiData
from mloda.core.abstract_plugins.compute_framework import ComputeFramework
from mloda.core.core.cfw_manager import CfwManager
from mloda.core.core.step.abstract_step import Step
from mloda.core.abstract_plugins.feature_group import FeatureGroup, format_feature_group_class
from mloda.core.abstract_plugins.components.feature_set import FeatureSet
from mloda.core.abstract_plugins.components.parallelization_modes import ParallelizationMode


class FeatureGroupStep(Step):
    def __init__(
        self,
        feature_group: type[FeatureGroup],
        features: FeatureSet,
        required_uuids: set[UUID],
        compute_framework: type[ComputeFramework],
        children_if_root: Optional[set[UUID]] = None,
        api_input_data: BaseApiData | bool = False,
    ) -> None:
        if children_if_root is None:
            children_if_root = set()
        self.feature_group = feature_group
        self.features = features
        self.required_uuids = required_uuids
        self.compute_framework = compute_framework
        self.children_if_root = frozenset(children_if_root.union({f.uuid for f in features.features}))
        self.api_input_data = api_input_data

        self.uuid = uuid4()

        self.step_is_done = False

        self.need_to_upload = False

        # Currently, also used for joinsteps without tfs. This might be a bug.
        self.tfs_ids: set[UUID] = set()

    def get_uuids(self) -> set[UUID]:
        return {feature.uuid for feature in self.features.features}

    def execute(
        self,
        cfw_register: CfwManager,
        cfw: ComputeFramework,
        from_cfw: Optional[ComputeFramework | UUID] = None,  # Not used in this implementation
        data: Optional[Any] = None,
    ) -> Optional[Any]:
        self.location = cfw_register.get_location()

        runtime_artifacts = cfw_register.get_runtime_artifacts()
        if runtime_artifacts is not None and self.feature_group.artifact() is not None:
            self.features.resolve_artifact_for_runtime(runtime_artifacts)

        if self.api_input_data:
            data = self.get_api_input_data(data, cfw_register)

        data = self.run_calculate_feature(cfw, data)
        self.save_artifact(self.features, cfw_register)

        # return_data_type_rule

        if self.location:
            if self.need_to_upload:
                cfw.upload_finished_data(self.location)
                cfw_register.add_uuid_flyway_datasets(cfw.uuid, set(self.children_if_root))
            return data
        return None

    def run_calculate_feature(self, cfw: ComputeFramework, data: Optional[Any] = None) -> Any:
        if self.feature_group.calculate_feature is None:
            raise ValueError("FeatureGroup calculate_feature is not implemented")

        data = cfw.run_calculation(self.feature_group, self.features, self.location, data)
        cfw.validate_expected_framework(self.location)
        return data

    def save_artifact(self, features: FeatureSet, cfw_register: CfwManager) -> None:
        artifact = self.validate_set_artifact_to_save(features)
        if artifact is None:
            return

        artifact_or_meta_information = artifact.save(features, features.save_artifact)

        # Handle multiple artifacts case (when artifact.save returns a dict of artifacts)
        if isinstance(artifact_or_meta_information, dict) and isinstance(features.save_artifact, dict):
            # This is a multiple artifacts case - save each artifact individually
            for artifact_key, artifact_path in artifact_or_meta_information.items():
                cfw_register.set_artifact_to_save(artifact_key, artifact_path)
        else:
            # Single artifact case (legacy behavior)
            assert features.artifact_to_save is not None  # validated in validate_set_artifact_to_save
            cfw_register.set_artifact_to_save(features.artifact_to_save, artifact_or_meta_information)

    def validate_set_artifact_to_save(self, features: FeatureSet) -> None | type[BaseArtifact]:
        if features.artifact_to_save is None:
            return None

        if features.save_artifact is None:
            raise ValueError(
                f"No artifact to save although it was requested. {self.feature_group} {features.artifact_to_save}."
            )

        artifact = self.feature_group.artifact()
        if artifact is None:
            raise ValueError(
                f"Artifact is not implemented for {self.feature_group}, but requested {features.artifact_to_save}"
            )
        return artifact

    def get_api_input_data(self, data: Any, cfw_register: CfwManager) -> Any:
        if data is not None:
            raise ValueError(
                f"Data is not None, but api_input_data is not False. {format_feature_group_class(self.feature_group)}."
            )

        if not isinstance(self.api_input_data, BaseApiData):
            raise ValueError(f"Api input data is not a tuple. {format_feature_group_class(self.feature_group)}.")

        api_data = cfw_register.get_api_data_by_name(self.api_input_data.get_api_input_name())

        data = self.api_input_data.get_data_by_using_api_data(api_data)

        if data is None:
            raise ValueError(
                f"Data is None: {self.feature_group} although we have an api_input_data {self.api_input_data.get_api_input_name()}."
            )

        return data

    def get_parallelization_mode(self) -> set[ParallelizationMode]:
        return self.compute_framework.supported_parallelization_modes()

    def add_value_to_children_if_root(self, value: UUID) -> None:
        self.children_if_root = self.children_if_root | frozenset([value])
