from yta_editor.tracks.items.transition.transition_node_registry import TransitionNodeRegistry
from yta_editor.tracks.items.transition.transition_specification import TransitionSpecification
from yta_editor.media.video import VideoFileMedia
from yta_editor_common.frame.helper import FrameHelper
from yta_editor_utils.texture import TextureUtils
from yta_validation import PythonValidator
from yta_math_easings.abstract import EasingFunction
from av.video.frame import VideoFrame
from quicktions import Fraction
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,
        transition_item: 'TransitionTrackItem',
        duration: float,
        transition_specification: Union['TransitionSpecification', None] = None,
        do_use_gpu: bool = True,
    ):
        self._transition_item: 'TransitionTrackItem' = transition_item
        """
        The transition item this helper was instantiated for.
        """
        self.duration: float = duration
        """
        The duration of the transition.
        """
        self._do_use_gpu: bool = do_use_gpu
        """
        *For internal useo only*

        The boolean flag to indicate if the GPU must be used
        or not when applying the node associated to this
        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}"')
        
        # 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.
        """
        self._node_processor = TransitionNodeRegistry.get(self._transition_specification.node_name)
        # I think it should be initialized here...
        # (
        #     # TODO: What do I do with the 'opengl_context' (?)
        #     opengl_context = None
        # )
        """
        *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: Union['moderngl.Texture', 'np.ndarray'],
        frame_b: Union['moderngl.Texture', 'np.ndarray'],
        evaluation_context: 'EvaluationContext'
    ) -> VideoFrame:
        """
        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'
        )

        processed_frame = frame_to_uint8_ndarray(
            frame = processed_frame,
            do_include_alpha = False
        )

        # if PythonValidator.is_numpy_array(processed_frame):
        #     processed_frame = processed_frame.astype(np.uint8)

        #     # frame = VideoFrame.from_ndarray(
        #     #     # TODO: Force this with the 'yta-numpy' utils to 'np.uint8'
        #     #     array = processed_frame,
        #     #     format = 'rgb24'
        #     # )
        # elif PythonValidator.is_instance_of(processed_frame, 'Texture'):
        #     processed_frame, has_alpha = FrameHelper.texture_to_ndarray(
        #         frame = processed_frame
        #     )
            
        #     # TODO: Removing the alpha channel here but, should I (?)
        #     if has_alpha:
        #         processed_frame = processed_frame[:, :, :3]

        # else:
        #     # TODO: Wtf? This should not happen
        #     raise Exception('Wowowowooo, no texture, no np.ndarray... wtf?')
        
        frame = FrameHelper.ndarray_to_videoframe(
            frame = processed_frame,
            do_include_alpha = False,
            pts = None,
            time_base = Fraction(1, self.fps)
        )
        
        # TODO: If I return this as a numpy and make the
        # transformations and then we want to use that
        # same frame within other GPU processor... we are
        # f***
        return frame
    
    # TODO: Overwrite this method to make it custom (?)
    def _get_process_frame_array(
        self,
        frame_a: Union['moderngl.Texture', 'np.ndarray'],
        frame_b: Union['moderngl.Texture', 'np.ndarray'],
        evaluation_context: 'EvaluationContext'
    ) -> Union[np.ndarray, 'moderngl.Texture']:
        """
        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: Maybe transform it in other way for optimization
            inputs['mask_input'] = self._mask_video.get_video_frame_at(
                t_media = t_transition_media
            ).to_ndarray(format = 'rgba')

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

        return self._node_processor.process(
            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)),
            output_size = self.output_size,
            do_use_gpu = self._do_use_gpu,
            # Parameters will be validated in
            parameters = parameters_dynamic
        )
    
    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


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'):
        frame, has_alpha = FrameHelper.texture_to_ndarray(frame)

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

    return TextureUtils.numpy_to_uint8(frame)