from functools import partial
from typing import Any, Callable, Generator, Optional, TYPE_CHECKING, Tuple, Union

from metaflow.debug import debug
from metaflow.exception import MetaflowException
from metaflow.user_configs.config_parameters import (
    ConfigValue,
    resolve_delayed_evaluator,
)

if TYPE_CHECKING:
    import metaflow.decorators
    import metaflow.flowspec
    import metaflow.parameters


class MutableStep:
    """
    A MutableStep is a wrapper passed to the `CustomStepDecorator`'s `evaluate` method
    to allow the decorator to interact with the step and providing easy methods to
    modify the behavior of the step.
    """

    def __init__(
        self,
        flow_spec: "metaflow.flowspec.FlowSpec",
        step: Union[
            Callable[["metaflow.decorators.FlowSpecDerived"], None],
            Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
        ],
    ):
        self._mutable_container = MutableFlow(flow_spec)
        self._my_step = step

    @property
    def flow(self) -> "MutableFlow":
        """
        The flow that contains this step

        Returns
        -------
        MutableFlow
            The flow that contains this step
        """
        return self._mutable_container

    @property
    def decorators(self) -> Generator["metaflow.decorators.StepDecorator", None, None]:
        """
        Iterate over all the decorators of this step. Note that the same type of decorator
        may be present multiple times and no order is guaranteed.

        Yields
        ------
        metaflow.decorators.StepDecorator
            A decorator of the step
        """
        for deco in self._my_step.decorators:
            yield deco

    def add_decorator(self, deco_type: partial, **kwargs) -> None:
        """
        Add a Metaflow decorator to a step.

        Parameters
        ----------
        deco_type : partial
            The decorator class to add to this step
        """
        # Prevent circular import
        from metaflow.decorators import DuplicateStepDecoratorException, StepDecorator

        # Validate deco_type
        if (
            not isinstance(deco_type, partial)
            or len(deco_type.args) != 1
            or not issubclass(deco_type.args[0], StepDecorator)
        ):
            raise TypeError("add_decorator takes a StepDecorator")

        deco_type = deco_type.args[0]
        if (
            deco_type.name in [deco.name for deco in self._my_step.decorators]
            and not deco_type.allow_multiple
        ):
            raise DuplicateStepDecoratorException(deco_type.name, self._my_step)

        debug.userconf_exec(
            "Mutable decorator adding step decorator %s to step %s"
            % (deco_type.name, self._my_step.name)
        )
        self._my_step.decorators.append(
            deco_type(attributes=kwargs, statically_defined=True)
        )

    def remove_decorator(self, deco_name: str, all: bool = True, **kwargs) -> bool:
        """
        Remove one or more Metaflow decorators from a step.

        Some decorators can be applied multiple times to a step. This method allows you
        to choose which decorator to remove or just remove all of them or one of them.

        Parameters
        ----------
        deco_name : str
            Name of the decorator to remove
        all : bool, default True
            If True, remove all instances of the decorator that match the filters
            passed using kwargs (or all the instances of the decorator if no filters are
            passed). If False, removes only the first found instance of the decorator.

        Returns
        -------
        bool
            Returns True if at least one decorator was removed.
        """
        new_deco_list = []
        did_remove = False
        for deco in self._my_step.decorators:
            if deco.name == deco_name:
                # Evaluate all the configuration values if any
                deco.init()

                # Check filters
                match_ok = True
                if kwargs:
                    for k, v in kwargs.items():
                        match_ok = k in deco.attributes and deco.attributes[k] == v
                        if match_ok is False:
                            break
                if match_ok:
                    did_remove = True
                    debug.userconf_exec(
                        "Mutable decorator removing step decorator %s from step %s"
                        % (deco.name, self._my_step.name)
                    )
                else:
                    new_deco_list.append(deco)
            else:
                new_deco_list.append(deco)
            if did_remove and not all:
                break

        self._my_step.decorators = new_deco_list
        return did_remove


class MutableFlow:
    def __init__(self, flow_spec: "metaflow.flowspec.FlowSpec"):
        self._flow_cls = flow_spec

    @property
    def decorators(self) -> Generator["metaflow.decorators.FlowDecorator", None, None]:
        """
        Iterate over all the decorators of this flow. Note that the same type of decorator
        may be present multiple times and no order is guaranteed.

        Yields
        ------
        metaflow.decorators.FlowDecorator
            A decorator of the flow
        """
        for decos in self._flow_cls._flow_decorators.values():
            for deco in decos:
                yield deco

    @property
    def configs(self) -> Generator[Tuple[str, ConfigValue], None, None]:
        """
        Iterate over all user configurations in this flow

        Use this to parameterize your flow based on configuration. As an example, the
        `evaluate` method of your `CustomFlowDecorator` can use this to add an
        environment decorator.
        ```
        class MyDecorator(CustomFlowDecorator):
            def evaluate(flow: MutableFlow):
                val = next(flow.configs)[1].steps.start.cpu
                flow.start.add_decorator(environment, vars={'mycpu': val})
                return flow

        @MyDecorator()
        class TestFlow(FlowSpec):
            config = Config('myconfig.json')

            @step
            def start(self):
                pass
        ```
        can be used to add an environment decorator to the `start` step.

        Yields
        ------
        Tuple[str, ConfigValue]
            Iterates over the configurations of the flow
        """
        from metaflow.flowspec import _FlowState

        # When configs are parsed, they are loaded in _flow_state[_FlowState.CONFIGS]
        for name, value in self._flow_cls._flow_state.get(
            _FlowState.CONFIGS, {}
        ).items():
            yield name, ConfigValue(value) if value is not None else None

    @property
    def parameters(self) -> Generator[Tuple[str, Any], None, None]:
        for var, param in self._flow_cls._get_parameters():
            if param.IS_CONFIG_PARAMETER:
                continue
            yield var, param

    @property
    def steps(self) -> Generator[Tuple[str, MutableStep], None, None]:
        """
        Iterate over all the steps in this flow. The order of the steps
        returned is not guaranteed.

        Yields
        ------
        Tuple[str, MutableStep]
            A tuple with the step name and the step proxy
        """
        for var in dir(self._flow_cls):
            potential_step = getattr(self._flow_cls, var)
            if callable(potential_step) and hasattr(potential_step, "is_step"):
                yield var, MutableStep(self._flow_cls, potential_step)

    def add_parameter(
        self, name: str, value: "metaflow.parameters.Parameter", overwrite: bool = False
    ) -> None:
        from metaflow.parameters import Parameter

        if hasattr(self._flow_cls, name) and not overwrite:
            raise MetaflowException(
                "Flow '%s' already has a class member '%s' -- "
                "set overwrite=True in add_parameter to overwrite it."
                % (self._flow_cls.__name__, name)
            )
        if not isinstance(value, Parameter) or value.IS_CONFIG_PARAMETER:
            raise MetaflowException(
                "Only a Parameter or an IncludeFile can be added using `add_parameter`"
                "; got %s" % type(value)
            )
        debug.userconf_exec("Mutable flow decorator adding parameter %s to flow" % name)
        setattr(self._flow_cls, name, value)

    def remove_parameter(self, parameter_name: str) -> bool:
        """
        Remove a parameter from the flow.

        The name given should match the name of the parameter (can be different
        from the name of the parameter in the flow. You can not remove config parameters.

        Parameters
        ----------
        parameter_name : str
            Name of the parameter

        Returns
        -------
        bool
            Returns True if the parameter was removed
        """
        from metaflow.flowspec import _FlowState

        for var, param in self._flow_cls._get_parameters():
            if param.IS_CONFIG_PARAMETER:
                continue
            if param.name == parameter_name:
                delattr(self._flow_cls, var)
                debug.userconf_exec(
                    "Mutable flow decorator removing parameter %s from flow" % var
                )
                # Reset so that we don't list it again
                del self._flow_cls._flow_state[_FlowState.CACHED_PARAMETERS]
                return True
        return False

    def add_decorator(self, deco_type: partial, **kwargs) -> None:
        """
        Add a Metaflow decorator to a flow.

        Parameters
        ----------
        deco_type : partial
            The decorator class to add to this flow
        """
        # Prevent circular import
        from metaflow.decorators import DuplicateFlowDecoratorException, FlowDecorator

        # Validate deco_type
        if (
            not isinstance(deco_type, partial)
            or len(deco_type.args) != 1
            or not issubclass(deco_type.args[0], FlowDecorator)
        ):
            raise TypeError("add_decorator takes a FlowDecorator")

        deco_type = deco_type.args[0]
        if (
            deco_type.name in self._flow_cls._flow_decorators
            and not deco_type.allow_multiple
        ):
            raise DuplicateFlowDecoratorException(deco_type.name)

        self._flow_cls._flow_decorators.setdefault(deco_type.name, []).append(
            deco_type(attributes=kwargs, statically_defined=True)
        )
        debug.userconf_exec(
            "Mutable flow decorator adding decorator %s to flow" % deco_type.name
        )

    def remove_decorator(self, deco_name: str, all: bool = True, **kwargs) -> bool:
        """
        Remove one or more Metaflow decorators from a flow.

        Some decorators can be applied multiple times to a flow. This method allows you
        to choose which decorator to remove or just remove all of them or one of them.

        Parameters
        ----------
        deco_name : str
            Name of the decorator to remove
        all : bool, default True
            If True, remove all instances of the decorator that match the filters
            passed using kwargs (or all the instances of the decorator if no filters are
            passed). If False, removes only the first found instance of the decorator.

        Returns
        -------
        bool
            Returns True if at least one decorator was removed.
        """
        new_deco_list = []
        old_deco_list = self._flow_cls._flow_decorators.get(deco_name)
        if old_deco_list is None:
            return False

        did_remove = False
        for deco in old_deco_list:
            # Evaluate all the configuration values if any
            deco.init()

            # Check filters
            match_ok = True
            if kwargs:
                for k, v in kwargs.items():
                    match_ok = k in deco.attributes and deco.attributes[k] == v
                    if match_ok is False:
                        break
            if match_ok:
                did_remove = True
                debug.userconf_exec(
                    "Mutable flow decorator removing decorator %s from flow" % deco.name
                )
            else:
                new_deco_list.append(deco)
            if did_remove and not all:
                break

        if new_deco_list:
            self._flow_cls._flow_decorators[deco_name] = new_deco_list
        else:
            del self._flow_cls._flow_decorators[deco_name]
        return did_remove

    def __getattr__(self, name):
        # We allow direct access to the steps, configs and parameters but nothing else
        from metaflow.parameters import Parameter

        attr = getattr(self._flow_cls, name)
        if attr:
            # Steps
            if callable(attr) and hasattr(attr, "is_step"):
                return MutableStep(self._flow_cls, attr)
            if name[0] == "_" or name in self._flow_cls._NON_PARAMETERS:
                raise AttributeError(self, name)
            if isinstance(attr, (Parameter, ConfigValue)):
                return attr
        raise AttributeError(self, name)


class CustomFlowDecorator:
    def __init__(self, *args, **kwargs):
        from ..flowspec import FlowSpecMeta

        if args and isinstance(args[0], (CustomFlowDecorator, FlowSpecMeta)):
            # This means the decorator is bare like @MyDecorator
            # and the first argument is the FlowSpec or another decorator (they
            # can be stacked)

            # If we have a init function, we call it with no arguments -- this can
            # happen if the user defines a function with default parameters for example
            try:
                self.init()
            except NotImplementedError:
                pass

            # Now set the flow class we apply to
            if isinstance(args[0], FlowSpecMeta):
                self._set_flow_cls(args[0])
            else:
                self._set_flow_cls(args[0]._flow_cls)
        else:
            # The arguments are actually passed to the init function for this decorator
            self._args = args
            self._kwargs = kwargs

    def __call__(
        self, flow_spec: Optional["metaflow.flowspec.FlowSpecMeta"] = None
    ) -> "metaflow.flowspec.FlowSpecMeta":
        if flow_spec:
            # This is the case of a decorator @MyDecorator(foo=1, bar=2) and so
            # we already called __init__ and saved foo and bar in self._args and
            # self._kwargs and are now calling this on the flow itself.

            # You can use config values in the arguments to a CustomFlowDecorator
            # so we resolve those as well
            new_args = [resolve_delayed_evaluator(arg) for arg in self._args]
            new_kwargs = {
                k: resolve_delayed_evaluator(v) for k, v in self._kwargs.items()
            }
            self.init(*new_args, **new_kwargs)
            if hasattr(self, "_empty_init"):
                raise MetaflowException(
                    "CustomFlowDecorator '%s' is used with arguments "
                    "but does not implement init" % str(self.__class__)
                )

            return self._set_flow_cls(flow_spec)
        elif not self._flow_cls:
            # This means that somehow the initialization did not happen properly
            # so this may have been applied to a non flow
            raise MetaflowException(
                "A CustomFlowDecorator can only be applied to a FlowSpec"
            )
        # NOTA: This returns self._flow_cls() because the object in the case of
        # @FlowDecorator
        # class MyFlow(FlowSpec):
        #     pass
        # the object is a FlowDecorator and when the main function calls it, we end up
        # here and need to actually call the FlowSpec. This is not the case when using
        # a decorator with arguments because in the line above, we will have returned a
        # FlowSpec object. Previous solution was to use __get__ but this does not seem
        # to work properly.
        return self._flow_cls()

    def _set_flow_cls(
        self, flow_spec: "metaflow.flowspec.FlowSpecMeta"
    ) -> "metaflow.flowspec.FlowSpecMeta":
        from ..flowspec import _FlowState

        flow_spec._flow_state.setdefault(_FlowState.CONFIG_DECORATORS, []).append(self)
        self._flow_cls = flow_spec
        return flow_spec

    def init(self, *args, **kwargs):
        """
        This method is intended to be optionally overridden if you need to
        have an initializer.
        """
        self._empty_init = True

    def evaluate(self, mutable_flow: MutableFlow) -> None:
        """
        Implement this method to act on the flow and modify it as needed.

        Parameters
        ----------
        mutable_flow : MutableFlow
            Flow

        Raises
        ------
        NotImplementedError
            _description_
        """
        raise NotImplementedError()


class CustomStepDecorator:
    def __init__(self, *args, **kwargs):
        arg = None
        if args:
            if isinstance(args[0], CustomStepDecorator):
                arg = args[0]._my_step
            else:
                arg = args[0]
        if arg and callable(arg) and hasattr(arg, "is_step"):
            # This means the decorator is bare like @MyDecorator
            # and the first argument is the step
            self._set_my_step(arg)
        else:
            # The arguments are actually passed to the init function for this decorator
            self._args = args
            self._kwargs = kwargs

    def __get__(self, instance, owner):
        # Required so that we "present" as a step when the step decorator is
        # of the form
        # @MyStepDecorator
        # @step
        # def my_step(self):
        #     pass
        #
        # In that case, if we don't have __get__, the object is a CustomStepDecorator
        # and not a step. Other parts of the code rely on steps having is_step. There are
        # other ways to solve this but this allowed for minimal changes going forward.
        return self()

    def __call__(
        self,
        step: Optional[
            Union[
                Callable[["metaflow.decorators.FlowSpecDerived"], None],
                Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
            ]
        ] = None,
    ) -> Union[
        Callable[["metaflow.decorators.FlowSpecDerived"], None],
        Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
    ]:
        if step:
            # You can use config values in the arguments to a CustomFlowDecorator
            # so we resolve those as well
            new_args = [resolve_delayed_evaluator(arg) for arg in self._args]
            new_kwargs = {
                k: resolve_delayed_evaluator(v) for k, v in self._kwargs.items()
            }
            self.init(*new_args, **new_kwargs)
            if hasattr(self, "_empty_init"):
                raise MetaflowException(
                    "CustomStepDecorator '%s' is used with arguments "
                    "but does not implement init" % str(self.__class__)
                )
            return self._set_my_step(step)
        elif not self._my_step:
            # This means that somehow the initialization did not happen properly
            # so this may have been applied to a non step
            raise MetaflowException(
                "A CustomStepDecorator can only be applied to a step function"
            )
        return self._my_step

    def _set_my_step(
        self,
        step: Union[
            Callable[["metaflow.decorators.FlowSpecDerived"], None],
            Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
        ],
    ) -> Union[
        Callable[["metaflow.decorators.FlowSpecDerived"], None],
        Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
    ]:

        self._my_step = step
        self._my_step.config_decorators.append(self)
        return self._my_step

    def init(self, *args, **kwargs):
        """
        This method is intended to be optionally overridden if you need to
        have an initializer.
        """
        self._empty_init = True

    def evaluate(self, mutable_step: MutableStep) -> None:
        raise NotImplementedError()
