from yta_programming.decorators.requires_dependency import requires_dependency
from yta_editor_utils.texture import TextureUtils
from typing import Union


class FrameHelper:
    """
    Class to simplify the way we handle frames from
    their different sources and according to the
    different and possible types.

    Possible frame types:
    - `av.VideoFrame`
    - `np.ndarray`
    - `moderngl.Texture`
    """

    @staticmethod
    @requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def ndarray_to_texture(
        frame: 'np.ndarray',
        moderngl_context: Union['moderngl.Context', None],
        do_include_alpha: bool = True,
        do_flip: bool = True
    ) -> 'moderngl.Texture':
        """
        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform the `numpy.ndarray` `frame` provided into a
        `moderngl.Texture` frame.

        The `do_flip` must be `True` normally because the
        origin of ModernGL and the origin of numpy are
        inverted.
        """
        # TODO: This is creating a new texture everytime,
        # and not just writing the data to a pre-created one
        return TextureUtils.numpy_to_texture(
            input = frame,
            opengl_context = moderngl_context,
            # TODO: Can I detect it (?)
            dtype = 'f1',
            do_flip = True
        )

        # TODO: Below is outdated, remove it soon
        import moderngl
        import numpy as np

        # TODO: Create 'moderngl_context' if None. I have a
        # class called 'OpenGLContext' to get it

        if frame.ndim != 3:
            raise ValueError('Expected HWC ndarray')

        dtype_map = {
            np.float32: 'f4',
            np.uint8: 'u1'
        }
        dtype = dtype_map.get(frame.dtype, 'u1')

        # We consider HWC format to obtain the size
        height, width = frame.shape[:2]

        # TODO: Is this flip here ok (?)
        frame = np.flip(frame, axis = 0)
        # This is needed because we are creating a texture
        frame = np.ascontiguousarray(frame)

        return moderngl_context.texture(
            (width, height),
            # This was originally 4 only but I added the bool
            # components = (
            #     4
            #     if do_include_alpha else
            #     3
            # ),
            # TODO: Apparently this above is not a good thing
            components = frame.shape[2],
            # data = frame.tobytes(),
            # TODO: Apparently 'frame' is accepted directly
            data = frame,
            dtype = dtype
        )
    
    @staticmethod
    @requires_dependency('av', 'yta_editor_frame', 'av')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def ndarray_to_videoframe(
        frame: 'np.ndarray',
        do_include_alpha: bool = True,
        pts: Union[int, None] = None,
        time_base: Union['Fraction', None] = None
    ) -> 'av.VideoFrame':
        """
        *Optional dependency `av` (imported as `av`) required*

        Transform the `numpy.ndarray` `frame` provided into a
        `av.VideoFrame` frame.
        """
        video_frame = TextureUtils.numpy_frame_to_pyav_frame(
            frame = frame
        )

        if pts is not None:
            video_frame.pts = pts

        if time_base is not None:
            video_frame.time_base = time_base

        return video_frame

        # TODO: Below is outdated, remove it soon
        import av

        format = (
            'rgba'
            if do_include_alpha else
            'rgb24'
        )

        frame = av.VideoFrame.from_ndarray(
            array = frame,
            format = format
        )

        if pts is not None:
            frame.pts = pts

        if time_base is not None:
            frame.time_base = time_base

        return frame

    @staticmethod
    @requires_dependency('av', 'yta_editor_frame', 'av')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def videoframe_to_ndarray(
        frame: 'av.VideoFrame',
        do_include_alpha: bool = True
    ) -> 'np.ndarray':
        """
        *Optional dependency `av` (imported as `av`) required*

        Transform the `av.VideoFrame` `frame` provided into a
        `numpy.ndarray` frame.
        """
        format = (
            'rgba'
            if do_include_alpha else
            'rgb24'
        )

        return frame.to_ndarray(format = format)
    
    @staticmethod
    @requires_dependency('av', 'yta_editor_frame', 'av')
    @requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def videoframe_to_texture(
        frame: 'av.VideoFrame',
        moderngl_context: 'moderngl.Context',
        do_include_alpha: bool = True,
        do_flip: bool = True
    ) -> 'moderngl.Texture':
        """
        *Optional dependency `av` (imported as `av`) required*

        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform the `av.VideoFrame` `frame` provided into a
        `moderngl.Texture` frame.

        The `do_flip` must be `True` normally because the
        origin of ModernGL and the origin of numpy are
        inverted.
        """
        # 1. 'av.VideoFrame' to 'np.ndarray'
        frame = FrameHelper.videoframe_to_ndarray(
            frame = frame,
            do_include_alpha = do_include_alpha
        )

        # 2. 'np.ndarray' to 'moderngl.Texture'
        return FrameHelper.ndarray_to_texture(
            frame = frame,
            moderngl_context = moderngl_context,
            do_include_alpha = do_include_alpha,
            do_flip = do_flip
        )
    
    @staticmethod
    @requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def texture_to_ndarray(
        frame: 'moderngl.Texture',
        do_include_alpha: bool = True,
        do_flip: bool = True,
        do_force_contiguous: bool = False
    ) -> 'np.ndarray':
        """
        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform the `moderngl.Texture` `frame` provided into
        a `np.ndarray` frame.

        This method does a `flipud` in the content.

        (!) This method returns a tuple that includes the array
        but also a boolean indicating if the texture included
        alpha or not.

        The `do_flip` must be `True` normally because the
        origin of ModernGL and the origin of numpy are
        inverted.

        The `do_force_contiguous` must be `True` only if the
        numpy will be transformed inmediately into a pyav
        videoframe (that needs it contiguous). If you will
        operate with it, set it to `False`.
        """
        return TextureUtils.texture_to_numpy(
            texture = frame,
            # TODO: Can I adjust this (?)
            do_include_alpha = do_include_alpha,
            do_flip = do_flip,
            do_force_contiguous = do_force_contiguous
        )

        # TODO: Below is outdated, remove it soon
        import numpy as np

        data = frame.read()
        width, height = frame.size

        dtype_map = {
            'f4': np.float32,
            # 'f2': np.float16,
            'f1': np.uint8,
            'u1': np.uint8
        }
        dtype = dtype_map.get(frame.dtype, np.uint8)

        data = np.frombuffer(data, dtype = dtype)

        if frame.dtype == 'f1':
            data = data * 255.0

        data = data.reshape((height, width, frame.components))

        data = np.flipud(data)
        data = np.ascontiguousarray(data)
        # Similar to 'np.flipud(frame)'
        # data = np.flip(data, axis = 0)

        """
        The 'frame.components defines if RGBA or RGB:
        - `frame.components=4` means RGBA
        - `frame.components=3` means RGB
        """
        return (data, frame.components == 4)

    @staticmethod
    @requires_dependency('av', 'yta_editor_frame', 'av')
    @requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def texture_to_videoframe(
        frame: 'moderngl.Texture',
        pts: Union[int, None] = None,
        time_base: Union['Fraction', None] = None,
        do_include_alpha: bool = True,
        do_flip: bool = True
    ) -> 'av.VideoFrame':
        """
        *Optional dependency `av` (imported as `av`) required*

        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform the `moderngl.Texture` `frame` provided into
        a `av.VideoFrame` frame.

        The `do_flip` must be `True` normally because the
        origin of ModernGL and the origin of numpy are
        inverted.
        """
        video_frame = TextureUtils.texture_to_pyav_frame(
            texture = frame,
            # TODO: Can I adjust this (?)
            do_include_alpha = do_include_alpha,
            do_flip = do_flip,
            do_force_contiguous = True
        )

        if pts is not None:
            video_frame.pts = pts

        if time_base is not None:
            video_frame.time_base = time_base

        return video_frame

        # TODO: Below is outdated, remove it soon
        # 1. 'moderngl.Texture' to 'np.ndarray'
        frame, do_include_alpha = FrameHelper.texture_to_ndarray(
            frame = frame
        )

        # 2. 'np.ndarray' to 'av.VideoFrame'
        frame = FrameHelper.ndarray_to_videoframe(
            frame = frame,
            do_include_alpha = do_include_alpha
        )

        if pts is not None:
            frame.pts = pts

        if time_base is not None:
            frame.time_base = time_base

        return frame