from yta_editor.media.video import VideoFileMedia
from yta_editor.transformations.effects.abstract import Effect
from yta_editor_frame.frame import Frame
from yta_editor_time.evaluation_context import EvaluationContext
from yta_editor_frame.frame.resource.identifier import FrameResourceIdentifier
from yta_editor_frame.decorators import force_frame
from yta_editor.utils import try_release_frame
from yta_validation.parameter import ParameterValidator
from yta_validation import PythonValidator
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.
    """

    @property
    @abstractmethod
    def do_need_render_target(
        self
    ) -> bool:
        """
        Boolean flag to indicate if the effect needs a
        `RenderTarget` instance to write on it because it
        will use at least one node that works in GPU.
        """
        pass

    def __init__(
        self,
        # TODO: Set the type (VideoEditorParameter)
        render_context: 'RenderContext',
        parameters: dict[str, any]
    ):
        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.render_context: 'Rendercontext' = render_context
        """
        The `RenderContext` instance.
        """

    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: 'Frame',
        render_target: Union['RenderTarget', None],
        parameters: dict[str, any],
        evaluation_context: EvaluationContext,
    ) -> Frame:
        """
        *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_frame('frame')
    def apply(
        self,
        frame: 'Frame',
        render_target: Union['RenderTarget', None],
        evaluation_context: EvaluationContext,
    ) -> 'Frame':
        """
        Apply the effect to the given `frame` according to
        the `evaluation_context`.
        """
        # 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_modified = self._apply_effect(
            frame = frame,
            render_target = render_target,
            parameters = parameters,
            evaluation_context = evaluation_context
        )

        return frame_modified
    
    def copy(
        self
    ) -> 'VideoEffect':
        """
        Get a copy of this `VideoEffect` instance.
        """
        return VideoEffect(
            render_context = self.render_context,
            parameters = self._parameters.copy()
        )
    
class VideoEffectSingleNode(VideoEffect):
    """
    A VideoEffect that is applied by using a
    single node processor for a single frame.
    """

    @property
    def do_need_render_target(
        self
    ):
        return self._node_processor._do_use_gpu

    def __init__(
        self,
        # TODO: Set type
        node_processor: any,
        parameters: dict[str, any],
        render_context: 'RenderContext'
    ):
        # TODO: Validate that is a NodeProcessor (?)
        self._node_processor: any = node_processor(
            render_context = render_context
        )
        """
        *For internal use only*

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

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

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: 'Frame',
        render_target: Union['RenderTarget', None],
        parameters: dict[str, any],
        evaluation_context: EvaluationContext
    ) -> Frame:
        """
        *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
        `Frame` instance.
        """
        return self._node_processor.process(
            render_target = render_target,
            inputs = {
                'base_input': frame
            },
            evaluation_context = evaluation_context,
            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],
        render_context: 'RenderContext'
    ):
        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,
            render_context = render_context
        )

    """
    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: 'Frame',
        render_target: Union['RenderTaregt', None],
        parameters: dict[str, any],
        evaluation_context: EvaluationContext
    ) -> Frame:
        """
        *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
        `Frame` instance.
        """
        return self._node_processor.process(
            render_target = render_target,
            inputs = {
                'base_input': frame,
                # We don't want the input as a parameter => '.pop'
                'mask_input': parameters.pop('mask_frame')
            },
            evaluation_context = evaluation_context,
            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.

    @property
    def render_context(
        self
    ) -> 'RenderContext':
        """
        The `RenderContext` instance associated to this,
        extracted from the first effect available.
        """
        return self.effects[0].render_context
    
    def __init__(
        self,
        effects: list[VideoEffect] = []
    ):
        self.effects: list[VideoEffect] = effects
        """
        The list of the effects to apply in the video.
        """

    def __getitem__(
        self,
        index: int
    ):
        return self.effects[index]

    def __len__(
        self
    ):
        return len(self.effects)

    def __iter__(
        self
    ):
        return iter(self.effects)

    @force_frame('frame')
    def apply(
        self,
        frame: 'Frame',
        evaluation_context: EvaluationContext,
        output_size: tuple[int, int]
    ) -> 'Frame':
        """
        Apply all the effects that should be applied at the
        given `t` time moment to the provided audio `frame`.
        """
        if len(self) > 0:
            """
            I need to handle the RenderTarget only if needed,
            because if it is being processed with CPU I don't need
            any RenderTarget, but if processed with GPU I'd
            probably need a ping-pong.

            The problem is that an Effect could be consecutive
            effects (nodes) that handle with GPU/CPU mixed...
            
            Only consecutive GPU effects must handle different RTs.

            If effects are parallel I need more than ping pong, but
            by now this is not happening.
            """
            if len(self.effects) == 1:
                effect = self.effects[0]
                
                if effect.do_need_render_target:
                    # GPU
                    render_target_effect = self.render_context.gpu.render_target.acquire(
                        width = output_size[0],
                        height = output_size[1],
                        number_of_components = self.render_context.config.number_of_components,
                        dtype = self.render_context.config.dtype
                    )

                    """
                    This 'frame_modified' could be exactly the
                    'frame' provided if the effect does not have
                    to be applied actually.
                    """
                    # TODO: Need to detect if effect is with CPU or not
                    frame_modified = self.effects[0].apply(
                        frame = frame,
                        render_target = render_target_effect,
                        evaluation_context = evaluation_context
                    )

                    if frame_modified is not frame:
                        try_release_frame(
                            frame = frame,
                            extra_message = f'in "VideoEffects.apply() with evaluation_context.t={str(float(evaluation_context.t))}'
                        )
                        frame = frame_modified
                else:
                    # CPU
                    frame_modified = self.effects[0].apply(
                        frame = frame,
                        render_target = None,
                        evaluation_context = evaluation_context
                    )

                    frame = frame_modified
            else:
                """
                If at least one effect needs GPU, we need to
                handle the `RenderTarget`
                """
                number_of_gpu_effects: int = sum(
                    effect.do_need_render_target
                    for effect in self.effects
                )

                """
                TODO: This could be simplified and if it is 1
                effect, acquire only one render target, but by
                now I act as I was doing before (that was 
                working correctly).
                """
                if number_of_gpu_effects > 0:
                    # GPU, at least one effect is based on it
                    # len(self.effects) will be always >1 here
                    render_target_1, render_target_2 = self.render_context.gpu.render_target.acquire_ping_pong(
                        width = output_size[0],
                        height = output_size[1],
                        number_of_components = self.render_context.config.number_of_components,
                        dtype = self.render_context.config.dtype
                    )

                    render_target_to_write = render_target_1
                    current = frame

                    for video_effect in self.effects:
                        if video_effect.do_need_render_target:
                            output = video_effect.apply(
                                frame = current,
                                render_target = render_target_to_write,
                                evaluation_context = evaluation_context
                            )

                            if output is current:
                                # Passthrough → no swap, no changes
                                continue

                            current = output

                            # Ping-pong swap
                            render_target_to_write = (
                                render_target_2
                                if render_target_to_write is render_target_1
                                else render_target_1
                            )
                        else:
                            output = video_effect.apply(
                                frame = current,
                                render_target = None,
                                evaluation_context = evaluation_context
                            )

                            if output is current:
                                # Passthrough → no swap, no changes
                                continue

                            current = output

                    frame = current

                    # Lets free resources
                    final_render_target = None
                    if (
                        hasattr(frame, 'resource') and
                        hasattr(frame.resource, 'frame') and
                        FrameResourceIdentifier.is_gpu_render_target(frame.resource)
                    ):
                        if frame.resource.frame in (render_target_1, render_target_2):
                            final_render_target = frame.resource.frame

                    # Free the 'RenderTarget' that is not in the final frame
                    if final_render_target is render_target_1:
                        self.render_context.gpu.render_target.release(render_target_2)
                    elif final_render_target is render_target_2:
                        self.render_context.gpu.render_target.release(render_target_1)
                    else:
                        # No effect applied, no RT used...
                        self.render_context.gpu.render_target.release(render_target_1)
                        self.render_context.gpu.render_target.release(render_target_2)
                else:
                    # CPU (all nodes are CPU-based)
                    current = frame

                    for video_effect in self.effects:
                        output = video_effect.apply(
                            frame = current,
                            render_target = None,
                            evaluation_context = evaluation_context
                        )

                        if output is current:
                            # Passthrough → no swap, no changes
                            continue

                        current = output

                    frame = current

        return frame

    def copy(
        self
    ) -> 'VideoEffects':
        """
        Get a copy of this `VideoEffects` instance.
        """
        return VideoEffects(
            effects = [
                video_effect.copy()
                for video_effect in self.effects
            ]
        )

# 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
        )
    )