from typing import Optional, cast
from mloda.core.abstract_plugins.compute_framework import ComputeFramework
from mloda.core.abstract_plugins.components.feature_collection import Features
from mloda.core.abstract_plugins.components.utils import get_all_subclasses
from mloda.core.abstract_plugins.components.parallelization_modes import ParallelizationMode


class SetupComputeFramework:
    """A class to create the compute framework and do basic validation."""

    def __init__(
        self,
        user_compute_frameworks: set[type[ComputeFramework]] | Optional[list[str]],
        features: Features,
        parallelization_modes: Optional[set[ParallelizationMode]] = None,
    ) -> None:
        available_compute_frameworks = get_all_subclasses(ComputeFramework)

        if user_compute_frameworks:
            if isinstance(user_compute_frameworks, list):
                user_set_compute_frameworks: set[str | type[ComputeFramework]] = set(user_compute_frameworks)
            else:
                user_set_compute_frameworks = cast(set[str | type[ComputeFramework]], user_compute_frameworks)

            available_compute_frameworks = self.filter_user_set_in_available_sub_classes(
                user_set_compute_frameworks, available_compute_frameworks
            )

        if parallelization_modes is not None:
            available_compute_frameworks = self._filter_by_parallelization_modes(
                available_compute_frameworks, parallelization_modes
            )
            if not available_compute_frameworks:
                raise ValueError(
                    f"No compute frameworks support the requested parallelization modes: {parallelization_modes}."
                )

        self.validate_if_at_least_one_feature_compute_framework_is_in_available_compute_framework(
            features, available_compute_frameworks
        )

        self.compute_frameworks = available_compute_frameworks

    def validate_if_at_least_one_feature_compute_framework_is_in_available_compute_framework(
        self, features: Features, available_compute_frameworks: set[type[ComputeFramework]]
    ) -> None:
        for feature in features.collection:
            if feature.compute_frameworks and not any(
                cf in available_compute_frameworks for cf in feature.compute_frameworks
            ):
                raise ValueError(
                    f"Feature {feature.name} has compute frameworks {feature.compute_frameworks} not in {available_compute_frameworks}."
                )

    def _filter_by_parallelization_modes(
        self,
        compute_frameworks: set[type[ComputeFramework]],
        parallelization_modes: set[ParallelizationMode],
    ) -> set[type[ComputeFramework]]:
        return {
            cfw
            for cfw in compute_frameworks
            if cfw.supported_parallelization_modes().intersection(parallelization_modes)
        }

    def filter_user_set_in_available_sub_classes(
        self,
        api_request_compute_frameworks: set[str | type[ComputeFramework]],
        sub_classes: set[type[ComputeFramework]],
    ) -> set[type[ComputeFramework]]:
        compute_frameworks = set()
        compute_frameworks = {
            sub
            for sub in sub_classes
            if sub.get_class_name() in api_request_compute_frameworks or sub in api_request_compute_frameworks
        }

        if not compute_frameworks:
            available_names = sorted(cls.get_class_name() for cls in sub_classes)
            raise ValueError(
                f"No given compute frameworks {api_request_compute_frameworks} found in "
                f"available compute frameworks: {available_names}."
            )
        return compute_frameworks
