from typing import Any, Optional
from uuid import UUID, uuid4

from mloda.core.abstract_plugins.components.error_utils import internal_invariant_error
from mloda.core.abstract_plugins.components.framework_transformer.cfw_transformer import (
    ComputeFrameworkTransformer,
)
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
from mloda.core.runtime.flight.flight_server import FlightServer


class TransformFrameworkStep(Step):
    def __init__(
        self,
        from_framework: type[ComputeFramework],
        to_framework: type[ComputeFramework],
        required_uuids: set[UUID],
        from_feature_group: type[FeatureGroup],
        to_feature_group: type[FeatureGroup],
        link_id: Optional[UUID] = None,
        right_framework_uuids: Optional[set[UUID]] = None,
    ) -> None:
        if right_framework_uuids is None:
            right_framework_uuids = set()
        self.from_framework = from_framework
        self.to_framework = to_framework
        self.required_uuids = required_uuids
        self.uuid = uuid4()
        self.from_feature_group = from_feature_group
        self.to_feature_group = to_feature_group
        self.link_id = link_id
        self.transformer = ComputeFrameworkTransformer()

        # This variable is only set, if the TFS was requested by a joinstep.
        self.right_framework_uuid: Optional[UUID] = None
        if len(right_framework_uuids) > 0:
            self.right_framework_uuid = next(iter(right_framework_uuids))

        self.step_is_done = False

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, TransformFrameworkStep):
            return False
        return (
            self.from_framework == other.from_framework
            and self.to_framework == other.to_framework
            and self.from_feature_group == other.from_feature_group
            and self.to_feature_group == other.to_feature_group
        )

    def __hash__(self) -> int:
        return hash((self.from_framework, self.to_framework, self.from_feature_group, self.to_feature_group))

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

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

        if from_cfw is None:
            raise ValueError(
                internal_invariant_error(
                    "from_cfw is None when executing TransformFrameworkStep.",
                    f"to_framework={self.to_framework.get_class_name()}, required_uuids={self.required_uuids}",
                    "The source compute framework must be provided to transform data between frameworks.",
                )
            )

        data = self.get_data(from_cfw)
        column_names = self.get_column_names(cfw_register, from_cfw)

        data = self.transform(cfw, data, column_names)

        cfw.set_data(data)
        cfw.set_column_names()

        if self.location:
            cfw.upload_finished_data(self.location)
            return data
        return None

    def get_column_names(self, cfw_register: CfwManager, from_cfw: ComputeFramework | UUID) -> set[str]:
        if self.location and isinstance(from_cfw, UUID):
            return cfw_register.get_column_names(from_cfw)

        if isinstance(from_cfw, UUID):
            raise ValueError("From_cfw is a UUID, but we are not using flightserver.")

        return from_cfw.get_column_names()

    def get_data(self, cfw: ComputeFramework | UUID) -> Any:
        """
        This method is used to get the data from the compute framework.
        If we are using multiprocessing, we use flightserver to transport the data.

        If we are not using multiprocessing, we just get the data from the compute framework.
        """
        if isinstance(cfw, UUID) and self.location:
            data = FlightServer.download_table(self.location, str(cfw))
            return data

        if isinstance(cfw, UUID):
            raise ValueError("From_cfw is a UUID, but we are not using flightserver.")

        return cfw.get_data()

    def set_data(self, cfw: ComputeFramework, data: Any) -> None:
        cfw.set_data(data)

    def transform(self, cfw: ComputeFramework, data: Any, feature_names: set[str]) -> Any:
        if self.equal_frameworks():
            return data

        _from_fw = self.from_framework.expected_data_framework()
        _to_fw = self.to_framework.expected_data_framework()

        # Try to find a transformation chain (direct or through PyArrow)
        transformation_chain = self.transformer.get_transformation_chain(_from_fw, _to_fw)

        if transformation_chain is None:
            raise KeyError(
                f"No transformation path found from {_from_fw} to {_to_fw}. "
                f"Available transformers: {list(self.transformer.transformer_map.keys())}"
            )

        # Apply transformations in sequence
        current_fw = _from_fw
        for i, transformer_cls in enumerate(transformation_chain):
            # Determine target framework for this transformation step
            if i == len(transformation_chain) - 1:
                # Last step: transform to final target
                target_fw = _to_fw
            else:
                # Intermediate step: find what this transformer outputs
                # Look for the transformer in the map to determine its output type
                for (src, dst), trans in self.transformer.transformer_map.items():
                    if trans == transformer_cls and src == current_fw:
                        target_fw = dst
                        break

            data = transformer_cls.transform(current_fw, target_fw, data, cfw.framework_connection_object)
            current_fw = target_fw  # Update current framework for next iteration

        return data

    def equal_frameworks(self) -> bool:
        if self.from_framework.expected_data_framework() == self.to_framework.expected_data_framework():
            return True
        return False
