from yta_editor.tracks.items.transition.transition_node_registry import TransitionNodeRegistry
from yta_editor.tracks.items.transition.transition_specification import TransitionSpecification
from yta_editor_frame.frame import Frame
from yta_editor.media.video import VideoFileMedia
from yta_editor_frame.helper import FrameHelper
from yta_editor_utils.texture import TextureUtils
from yta_validation import PythonValidator
from yta_math_easings.abstract import EasingFunction
from typing import Union

import numpy as np


# TODO: This class below could be repeated and I just
# created because I needed to (at least by now)
# TODO: The progress is not being calculated correctly
class _TransitionCalculator:
    """
    *For internal use only*

    A class to simplify the way we make the calculations
    related to a transition. This class must be used within
    a transition item that needs these kind of calculations.
    """

    @property
    def fps(
        self
    ) -> float:
        """
        The `fps` associated to the `transition_item` that
        is attached to this instance.
        """
        # TODO: Is this the 'output_fps' (?)
        return self._transition_item.fps
    
    @property
    def output_size(
        self
    ) -> float:
        """
        The `output_size` associated to the `transition_item`
        that is attached to this instance.
        """
        return self._transition_item.output_size

    def __init__(
        self,
        render_context: 'RenderContext',
        transition_item: 'TransitionTrackItem',
        duration: float,
        transition_specification: Union['TransitionSpecification', None] = None
    ):
        self.render_context: 'RenderContext' = render_context
        """
        The `RenderContext` instance.
        """
        self._transition_item: 'TransitionTrackItem' = transition_item
        """
        The transition item this helper was instantiated for.
        """
        self.duration: float = duration
        """
        The duration of the transition.
        """

        transition_specification = (
            TransitionSpecification(
                node_name = TransitionNodeRegistry.get_default_key(),
                # TODO: This must be dynamic according to the default
                parameters = {
                    'border_smoothness': 0.02
                },
                easing_function = 'linear'
            )
            if transition_specification is None else
            transition_specification
        )

        if not TransitionNodeRegistry.is_registered(transition_specification.node_name):
            raise Exception(f'There is no transition node processor registered with the given "{transition_specification.node_name}" name.')
        
        # TODO: The parameters of the specification must be
        # validated here according to the node processor
        # specification (which doesn't exist yet). By now
        # we are not doing this parameters validation
        self._transition_specification: 'TransitionSpecification' = transition_specification
        """
        The specification that defines the node processor
        and the parameters we should use.
        """

        # The transition could need the 'mask_input'
        mask_input_filename = self._transition_specification.parameters.get('mask_input', None)
        self._mask_video: Union[VideoFileMedia, None] = (
            VideoFileMedia(
                filename = mask_input_filename
            )
            if mask_input_filename is not None else
            None
        )
        """
        *For internal use only*

        The alpha video that we will use to read the
        frames and apply the mask.
        """
        # TODO: What about the GPU or CPU (?)
        self._node_processor = TransitionNodeRegistry.get(self._transition_specification.node_name)
        """
        *For internal use only*

        The node processor that we will use to process
        the input frame with the corresponding mask
        frames.
        """

    # TODO: This method can be overwritten to change
    # its behaviour if the specific transition needs
    # it
    def process_frame(
        self,
        frame_a: Frame,
        frame_b: Frame,
        evaluation_context: 'EvaluationContext'
    ) -> Frame:
        """
        Process the `frame_a` and the `frame_b` provided and
        apply the transition at the `t_progress` provided by
        using the node of the `TransitionTrackItem` associated
        to this instance.
        """
        # TODO: Maybe this can be placed in the general
        # class if it doesn't change

        # TODO: This can be a texture
        processed_frame = self._get_process_frame_array(
            frame_a = frame_a,
            frame_b = frame_b,
            evaluation_context = evaluation_context
        # TODO: Force this with the 'yta-numpy' utils to 'np.uint8' (?)
        )

        # frame_a.release()
        # frame_b.release()

        # I think it is uint8 always...
        return processed_frame
    
    # TODO: Overwrite this method to make it custom (?)
    def _get_process_frame_array(
        self,
        frame_a: Frame,
        frame_b: Frame,
        evaluation_context: 'EvaluationContext'
    ) -> 'Frame':
        """
        Get the numpy array that will be used to build the
        pyav VideoFrame that will be returned.
        """
        inputs = {
            'first_input': frame_a,
            'second_input': frame_b
        }

        # TODO: Check this below to improve it if possible
        parameters_dynamic = self._transition_specification.parameters.copy()

        # Alpha mask transition, we need to handle it in
        # a special way
        if self._mask_video is not None:
            # TODO: Make it a property (?)
            easing_function = EasingFunction.get(self._transition_specification.easing_function)()
            # We calculate the transition t_media with the progress
            t_transition_media = easing_function.ease(evaluation_context.progress.value) * (self._mask_video.number_of_frames - 1) / self._mask_video.fps

            # We get a VideoFrame, so we transform it
            """
            TODO: This below is working but maybe it is not
            the best way. An alpha video should be short and
            maybe we can load it in memory directly or 
            something. Try to improve it, but remember: it
            is working :)
            """
            # TODO: Maybe transform it in other way for optimization
            inputs['mask_input'] = Frame.from_frame(
                frame = self._mask_video.get_video_frame_at(
                    t_media = t_transition_media
                ).to_ndarray(format = 'rgba'),
                render_context = self.render_context
            )

            # We don't want it as a dynamic parameter but as input
            del parameters_dynamic['mask_input']

        render_target = self.render_context.gpu.render_target.acquire(
            width = self.render_context.config.size[0],
            height = self.render_context.config.size[1],
            number_of_components = 4,
            dtype = 'f1'
        )

        frame_processed = self._node_processor.process(
            render_target = render_target,
            inputs = inputs,
            evaluation_context = evaluation_context,
            # TODO: This is how we were using the easing function before
            # progress = self.easing_function.ease(float(evaluation_context.progress)),
            # TODO: Remove when everything is working
            # output_size = self.output_size,
            # Parameters will be validated in
            parameters = parameters_dynamic
        )

        return frame_processed
    
    def _get_process_audio_frames(
        self,
        audio_frames_a: list['AudioFrameWrapped'],
        audio_frames_b: list['AudioFrameWrapped'],
        evaluation_context: 'EvaluationContext'
    ) -> list['AudioFrame']:
        """
        Get the numpy array that will be used to build the
        pyav AudioFrame that will be returned.
        """

        """
        TODO: I don't know how to handle the audio within the
        transitions, because each transition works in a specific
        way and making the audio behave like the video movement
        is not easy... It could be based on the 't_progress',
        using something like this below (based on S-Curve Crossfade
        weights):

        weight_a = 0.5 * (1 + math.cos(t_progress * math.pi))
        weight_b = 1 - weight_a
        mixed_audio_frames = audio_frames_a * weight_a + audio_frames_b * weight_b

        But, as I don't know actually how to do it, by now we are
        just mixing both of them during the transition time in the
        same way as the video frames are accessed.
        """
        # TODO: Move this to top (?)
        from yta_editor.utils.frame_combinator import AudioFrameCombinator
        from yta_editor.timeline import concatenate_audio_frames

        collapsed_frames = [
            concatenate_audio_frames(frames)
            for frames in [audio_frames_a, audio_frames_b]
        ]

        # We keep only the non-silent frames because
        # we will sum them after and keeping them will
        # change the results.
        non_empty_collapsed_frames = [
            frame._frame
            for frame in collapsed_frames
            if not frame.is_from_gap
        ]

        if len(non_empty_collapsed_frames) == 0:
            # If they were all silent, just keep one
            non_empty_collapsed_frames = [collapsed_frames[0]._frame]

        frames = [
            AudioFrameCombinator.sum_tracks_frames(
                tracks_frames = non_empty_collapsed_frames,
                sample_rate = self._transition_item._track.audio_fps,
                # TODO: This was not being sent before
                layout = self._transition_item._track.audio_layout,
                format = self._transition_item._track.audio_format
            )
        ]

        # TODO: Do we need the '**self._transition_specification.parameters'
        # for the audio frames (?)

        # for audio_frame in frames:
        #     yield audio_frame
        # TODO: Should I yield (?)
        return frames

# TODO: No longer used, remove it (?)
def frame_to_uint8_ndarray(
    frame: Union['moderngl.Texture', 'np.ndarray'],
    do_include_alpha: bool = True
):
    """
    Force the `frame` provided to be transformed into
    a uint8 numpy ndarray, removing the alpha channel
    if existing and `do_include_alpha` is `False`.
    """
    if (
        not PythonValidator.is_instance_of(frame, 'Texture') and
        not PythonValidator.is_numpy_array(frame)
    ):
        # TODO: Wtf? This should not happen
        raise Exception('Wowowowooo, no texture, no np.ndarray... wtf?')
    
    if PythonValidator.is_instance_of(frame, 'Texture'):
        # Old way below
        # frame, has_alpha = FrameHelper.texture_to_ndarray(frame)

        # if (
        #     has_alpha and
        #     not do_include_alpha
        # ):
        #     frame = frame[:, :, :3]

        frame = FrameHelper.texture_to_ndarray(frame)

        # if (
        #     not do_include_alpha and
        #     frame.ndim == 3 and
        #     frame.shape[2] > 3
        # ):
        #     frame = frame[:, :, :3]

    return TextureUtils.numpy_to_uint8(frame)