from yta_editor.media.video import VideoFileMedia
from yta_editor.transformations.effects.abstract import Effect
from yta_editor_frame.buffer import FrameBuffer
from yta_editor_time.evaluation_context import EvaluationContext
from yta_editor_frame.decorators import force_framebuffer
from yta_validation.parameter import ParameterValidator
from yta_validation import PythonValidator
# from av.video import VideoFrame
from typing import Union
from abc import abstractmethod


class VideoEffect(Effect):
    """
    An effect for a video input.

    This class is just to be able to identify the
    hierarchy and which effect is designed for 
    the video.
    """

    _parameter_values_to_skip_process: dict[str, any] = {}
    """
    The dict that contains the parameter keys and
    the values that must have to skip the
    processing if all the parameters included here
    have the expected values. A single value that
    has a different value or if the dict is empty,
    the processing must be done and not skipped.
    """

    def __init__(
        self,
        # TODO: Set the type (VideoEditorParameter)
        parameters: dict[str, any],
        do_use_gpu: bool = True
    ):
        self._parameters: dict[str, any] = parameters
        """
        The parameters that must be evaluated to be able to
        apply the effect in the video, and the keys must be
        exactly the values the nodes are expecting.
        """
        self.do_use_gpu: bool = do_use_gpu
        """
        Flag to indicate if using GPU or not, based on what
        the user requested.
        """

    def _get_params_at(
        self,
        evaluation_context: EvaluationContext
    ) -> dict:
        """
        *For internal use only*

        Get the parameters that must be applied at the given
        `evaluation_context`.
        """
        ParameterValidator.validate_mandatory_instance_of('evaluation_context', evaluation_context, EvaluationContext)

        parameters = {}
        for parameter_name, parameter_value in self._parameters.items():
            parameters[parameter_name] = parameter_value.evaluate(
                evaluation_context = evaluation_context
            )

        return parameters
    
    def _do_skip_process(
        self,
        parameters: dict[str, any]
    ) -> bool:
        """
        *For internal use only*

        Check if the `parameters` provided have the
        values that make the frame change not, so we
        can avoid applying the effect and returning it
        as it is directly. If a single parameter has a
        different value that the one to skip
        processing the frame, we must process it and
        this method will return `False`.

        This method must be set for every specific
        effect to avoid calculations if we don't need
        them.
        
        Here you have on example:

        ```
        parameters = {
            'factor': 1.0
        }
        ```
        """
        if len(self._parameter_values_to_skip_process.keys()) == 0:
            return False
        
        for parameter_key, parameter_value in self._parameter_values_to_skip_process.items():
            if parameters[parameter_key] != parameter_value:
                return False

        return True

    @abstractmethod
    def _apply_effect(
        self,
        frame: 'FrameBuffer',
        parameters: dict[str, any],
        output_size: Union[tuple[int, int], None],
        evaluation_context: EvaluationContext,
    ) -> Union['np.ndarray', 'moderngl.Texture']:
        """
        *For internal use only*

        Execute the node processors we need to generate
        the desired effect, using the `inputs` and the
        `parameters` provided to obtain a processed and
        final frame.

        This method must be overwritten in the specific
        effects to do the changes we need.
        """
        pass

    @force_framebuffer('frame')
    def apply(
        self,
        frame: 'FrameBuffer',
        evaluation_context: EvaluationContext,
        output_size: Union[tuple[int, int], None]
    ) -> 'FrameBuffer':
        """
        Apply the effect to the given `frame` according to
        the `evaluation_context` provided and forcing the
        `output_size` if given.
        """
        # 1. Evaluate the parameters
        parameters = self._get_params_at(
            evaluation_context = evaluation_context
        )

        # 2. Skip if parameters doesn't make any change
        if self._do_skip_process(parameters):
            return frame
        
        # 3. Process the effect
        frame = self._apply_effect(
            frame = frame,
            parameters = parameters,
            output_size = output_size,
            evaluation_context = evaluation_context
        )

        # 4. Transform to `FrameBuffer` again to return
        return (
            FrameBuffer.from_numpy_ndarray(frame)
            if PythonValidator.is_numpy_array(frame) else
            FrameBuffer.from_moderngl_texture(frame)
        )
    
class VideoEffectSingleNode(VideoEffect):
    """
    A VideoEffect that is applied by using a
    single node processor for a single frame.
    """

    def __init__(
        self,
        # TODO: Set type
        node_processor: any,
        parameters: dict[str, any],
        do_use_gpu: bool = True
    ):
        # TODO: Validate that is a NodeProcessor (?)
        self._node_processor: any = node_processor(
            # TODO: What do I do with the 'opengl_context' (?)
            opengl_context = None
        )
        """
        *For internal use only*

        The node processor that will modify the frame.
        """

        """
        We will force the value to be the one according
        to what the user requested and the availability
        in the node processor
        """
        do_use_gpu = _should_use_gpu(
            node_processor = self._node_processor,
            do_use_gpu = do_use_gpu
        )

        # TODO: Validate that are ok for the specific node
        super().__init__(
            parameters = parameters,
            do_use_gpu = do_use_gpu
        )

class VideoEffectSingleNodeBaseInput(VideoEffectSingleNode):
    """
    A VideoEffect that is applied by using a
    single node processor for a single frame.

    This effect will use a single input `base_input`
    with the frame provided, and apply the changes
    based on the `parameters` given.
    """

    def _apply_effect(
        self,
        frame: 'FrameBuffer',
        parameters: dict[str, any],
        evaluation_context: EvaluationContext,
        output_size: Union[tuple[int, int], None]
    ) -> Union['np.ndarray', 'moderngl.Texture']:
        """
        *For internal use only*

        Process the `frame` provided with the `parameters`
        given, according to the specific `node` that this
        specific class is using.

        This method will return a `np.ndarray` or a
        `moderngl.Texture` (according to if it is processed
        with CPU or GPU) that must be transformed into a
        `FrameBuffer` instance.
        """
        return self._node_processor.process(
            inputs = {
                'base_input': frame
            },
            evaluation_context = evaluation_context,
            do_use_gpu = self.do_use_gpu,
            output_size = output_size,
            parameters = parameters
        )
    
# TODO: Maybe we could implement this directly in
# the 'alphapedia_mask_reveal.py' because we don't
# have any other node
class VideoEffectSingleNodeBaseInputAndMaskInput(VideoEffectSingleNode):
    """
    A VideoEffect that is applied by using a
    single node processor for a single frame
    combined with a mask frame.

    This effect will use a single input `base_input`
    with the frame provided and a 'mask_input' to
    combine with, and apply the changes based on the
    `parameters` given.
    """

    def __init__(
        self,
        # TODO: Set type
        node_processor: any,
        mask_input_filename: str,
        parameters: dict[str, any],
        do_use_gpu: bool = True
    ):
        self._mask_video: VideoFileMedia = VideoFileMedia(
            filename = mask_input_filename
        )
        """
        *For internal use only*

        The alpha video that we will use to read the
        frames and apply the mask.
        """

        super().__init__(
            node_processor = node_processor,
            parameters = parameters,
            do_use_gpu = do_use_gpu
        )

    """
    TODO: We are mixing parameters and inputs, so maybe
    I could refactor this a bit to respect more the
    input and parameter concepts.
    """
    def _get_params_at(
        self,
        evaluation_context: EvaluationContext
    ):
        # Get parameters as usual
        parameters = super()._get_params_at(
            evaluation_context = evaluation_context
        )

        # TODO: This is repeated in the 'transition_calculator.py'
        # TODO: What about an easing function (?)
        t_transition_media = evaluation_context.progress.value * (self._mask_video.number_of_frames - 1) / self._mask_video.fps

        video_frame = self._mask_video.get_video_frame_at(
            t_media = t_transition_media
        ).to_ndarray(format = 'rgba')
        
        video_frame = (
            video_frame.frame
            if PythonValidator.is_instance_of(video_frame, 'VideoFrame') else
            video_frame
        )

        # Get extra input
        parameters['mask_frame'] = video_frame

        return parameters
        
    def _apply_effect(
        self,
        frame: 'FrameBuffer',
        parameters: dict[str, any],
        evaluation_context: EvaluationContext,
        output_size: Union[tuple[int, int], None]
    ) -> Union['np.ndarray', 'moderngl.Texture']:
        """
        *For internal use only*

        Process the `frame` provided with the `parameters`
        given, according to the specific `node` that this
        specific class is using.

        This method will return a `np.ndarray` or a
        `moderngl.Texture` (according to if it is processed
        with CPU or GPU) that must be transformed into a
        `FrameBuffer` instance.
        """
        return self._node_processor.process(
            inputs = {
                'base_input': frame,
                # We don't want the input as a parameter => '.pop'
                'mask_input': parameters.pop('mask_frame')
            },
            evaluation_context = evaluation_context,
            do_use_gpu = self.do_use_gpu,
            output_size = output_size,
            parameters = parameters
        )

class VideoEffects:
    """
    A set of effects to apply in a video input.
    """

    # TODO: This is not as simple as a list, it must be
    # managed with the specific nodes we designed that
    # are able to combine and are time-concerned.

    def __init__(
        self,
        effects: list[VideoEffect] = []
    ):
        self.effects: list[VideoEffect] = effects
        """
        The list of the effects to apply in the video.
        """

    @force_framebuffer('frame')
    def apply(
        self,
        # frame: Union['np.ndarray', 'Texture'],
        frame: 'FrameBuffer',
        output_size: Union[tuple[int, int], None],
        evaluation_context: EvaluationContext
    # ) -> Union['np.ndarray', 'Texture']:
    ) -> 'FrameBuffer':
        """
        Apply all the effects that should be applied at the
        given `t` time moment to the provided audio `frame`.
        """
        if len(self.effects) > 0:
            processed_frame = frame

            for effect in self.effects:
                processed_frame = effect.apply(
                    frame = processed_frame,
                    output_size = output_size,
                    evaluation_context = evaluation_context
                )

            return processed_frame

        return frame

# TODO: This should go to a utils I think
def _should_use_gpu(
    # TODO: Set type
    node_processor: any,
    do_use_gpu: bool = True
) -> bool:
    """
    *For internal use only*

    Check if we should use GPU or not, according
    to what the user is requesting with the
    `do_use_gpu` parameter but checking if it is
    available in the `node_processor` or not.

    This method is useful to know if we have to
    transform a frame to a specific format (for
    GPU or CPU) before sending it as an argument.
    """
    # TODO: Validate 'node_processor'
    return (
        (
            do_use_gpu and
            node_processor.is_gpu_available
        )
    ) or (
        (
            not do_use_gpu and
            not node_processor.is_cpu_available
        )
    )