from yta_editor_frame.domain import FrameBufferDomain
from yta_constants.enum import YTAEnum as Enum
from yta_programming.decorators.requires_dependency import requires_dependency
from yta_validation.parameter import ParameterValidator
from yta_validation import PythonValidator
from dataclasses import dataclass, field
from typing import Union


# TODO: Move these Enums, please
class FrameBufferFormat(Enum):
    """
    The format of the frame buffer that will identify
    if we have alpha channel or not.
    """

    RGBA = 'rgba'
    """
    The `rgba` format.
    """

class FrameBufferDType(Enum):
    """
    The numpy dtype of the frame buffer that will be
    helpful to operate with it.
    """

    UINT8 = 'uint8'
    """
    The `uint8` dtype.
    """

FORMAT_TO_FORCE = 'rgba'

@dataclass
class FrameBuffer:
    """
    *Dataclass*

    Dataclass to store the information about a buffer
    that we will need during the processing.

    This library could need one of these optional
    dependencies depending on what method you call to
    use it in the context you need it.
    - `av`
    - `numpy`
    - `moderngl`

    We will always store the frames as they are 
    provided, but forced to be 'RGBA' when returned
    with the `as_...` method.
    """

    # Common attributes
    size: tuple[int, int]
    """
    The size of the frame, in pixels, expressed as a
    `(width, height)` tuple.
    """

    # Attributes if 'av.VideoFrame'
    @property
    def av_format(
        self
    ) -> Union[str, None]:
        """
        The `format` of the frame as a `av.VideoFrame`, if
        existing.
        """
        return (
            self._av_videoframe.format.name
            if self._av_videoframe is not None else
            None
        )
    
    @property
    def av_planes(
        self
    ) -> Union[int, None]:
        """
        The number of `planes` of the frame as a
        `av.VideoFrame`, if existing.
        """
        return (
            len(self._av_videoframe.planes)
            if self._av_videoframe is not None else
            None
        )

    # Attributes if 'np.ndarray'
    @property
    def np_dtype(
        self
    ) -> Union['np.dtype', None]:
        """
        The `dtype` of the frame as a `np.ndarray`, if existing.
        """
        return (
            self._numpy_ndarray.dtype
            if self._numpy_ndarray is not None else
            None
        )
    
    @property
    def np_shape(
        self
    ) -> Union[tuple, None]:
        """
        The `shape` of the frame as a `np.ndarray`, if existing.
        """
        return (
            self._numpy_ndarray.shape
            if self._numpy_ndarray is not None else
            None
        )
    
    # Attributes if 'moderngl.Texture'
    @property
    def gl_dtype(
        self
    ) -> Union[str, None]:
        """
        The `dtype` of the frame as a `moderngl.Texture`, if
        existing.
        """
        return (
            self._moderngl_texture.dtype
            if self._moderngl_texture is not None else
            None
        )
    
    @property
    def gl_components(
        self
    ) -> Union[str, None]:
        """
        The number of `components` of the frame as a 
        `moderngl.Texture`, if existing.
        """
        return (
            self._moderngl_texture.components
            if self._moderngl_texture is not None else
            None
        )

    source_domain: FrameBufferDomain = field(default = None, init = False)
    """
    The domain that has been used as the source of
    the frame, which means that is the one that was
    used to create the instance and the rest, if
    existing, have been created from it.

    This source domain will also set when we modify
    the frame and we write on it, so it becomes the
    new source.
    """

    _av_videoframe: Union['av.VideoFrame', None] = field(default = None, init = False)
    """
    *For internal use only*

    The frame as a `av.VideoFrame`.
    """
    _numpy_ndarray: Union['np.ndarray', None] = field(default = None, init = False)
    """
    *For internal use only*

    The frame as a `numpy.ndarray`.
    """
    _moderngl_texture: Union['moderngl.Texture', None] = field(default = None, init = False)
    """
    *For internal use only*

    The frame as a `moderngl.Texture`.
    """

    @property
    def width(
        self
    ) -> int:
        """
        The width of the frame.
        """
        return self.size[0]
    
    @property
    def height(
        self
    ) -> int:
        """
        The height of the frame.
        """
        return self.size[1]
    
    @property
    def domains(
        self
    ) -> list[FrameBufferDomain]:
        """
        Get a list containing the domains available based
        on the information we have in the different
        domains (as numpy, texture or pyav frame).
        """
        domains = []

        if self._moderngl_texture is not None:
            domains.append(FrameBufferDomain.GPU)

        if self._numpy_ndarray is not None:
            domains.append(FrameBufferDomain.CPU)

        if self._av_videoframe is not None:
            domains.append(FrameBufferDomain.MEDIA)

        return domains

    @classmethod
    def from_frame(
        cls,
        frame: Union['av.VideoFrame', 'np.ndarray', 'moderngl.Texture', 'FrameBuffer']
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from one of the following
        frame types:
        - `av.VideoFrame`
        - `numpy.ndarray`
        - `moderngl.Texture`
        - `FrameBuffer`
        """
        ParameterValidator.validate_mandatory_instance_of('frame', frame, ['VideoFrame', 'ndarray', 'Texture', 'FrameBuffer'])

        return (
            FrameBuffer.from_av_videoframe(frame)
            if PythonValidator.is_instance_of(frame, 'VideoFrame') else
            FrameBuffer.from_numpy_ndarray(frame)
            if PythonValidator.is_numpy_array(frame) else
            FrameBuffer.from_moderngl_texture(frame)
            if PythonValidator.is_instance_of(frame, 'Texture') else
            frame
        )

    @classmethod
    @requires_dependency('av', 'yta_editor_frame', 'av')
    def from_av_videoframe(
        cls,
        frame: 'av.VideoFrame'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `av.VideoFrame`.

        This method will store the original frame as it is,
        modifying not the type nor the format.
        """
        frame_buffer = cls(
            size = (frame.width, frame.height)
        )
        frame_buffer._av_videoframe = frame
        frame_buffer.source_domain = FrameBufferDomain.MEDIA

        return frame_buffer
    
    @classmethod
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def from_numpy_ndarray(
        cls,
        frame: 'np.ndarray'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `numpy.ndarray`.

        This method will store the original frame as it is,
        modifying not the type nor the format.
        """
        height, width = frame.shape[:2]

        frame_buffer = cls(
            size = (width, height)
        )
        frame_buffer._numpy_ndarray = frame
        frame_buffer.source_domain = FrameBufferDomain.CPU

        return frame_buffer
    
        # To force 'rgba'
        frame = frame.reformat(format = 'rgba')

    @classmethod
    @requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
    def from_moderngl_texture(
        cls,
        frame: 'moderngl.Texture'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `moderngl.Texture`.
        
        This method will store the original frame as it is,
        modifying not the type nor the format.
        """
        width, height = frame.size

        # components = arr.shape[2] if arr.ndim == 3 else 1

        frame_buffer = cls(
            size = (width, height)
        )
        frame_buffer._moderngl_texture = frame
        frame_buffer.source_domain = FrameBufferDomain.GPU

        return frame_buffer
    
    def _validate_source_available(
        self
    ):
        """
        Validate that there is at least one source to be
        used to transform into the one we need for the
        method that calls this validation.

        This method will raise an exception if there is
        not a single domain available.
        """
        if len(self.domains) == 0:
            raise RuntimeError('No source available in this FrameBuffer instance.')
        
    """
    TODO: Maybe I need the 'requires_dependency' in some
    of these 'as_...' methods
    """
    
    def as_numpy_ndarray(
        self,
        do_write: bool = False
    ) -> 'np.ndarray':
        """
        Get the frame but as a `numpy.ndarray`. This is
        meant for CPU use.

        This method will lazy load the value from the
        one that is available in the moment it is
        requested.

        If you will modify the return, set the `do_write`
        as `True`.

        The `do_write` parameter is to indicate that
        when you call the method with `do_write=True`
        is because you'll modify the instance and
        that one will become the new absolute truth,
        so the other types will be removed and this one,
        as returned as a reference, will be modified by
        your code and become the new valid one.

        This method will always force the `rgba` format.
        """
        self._validate_source_available()

        if self._numpy_ndarray is None:
            if self._av_videoframe is not None:
                # 1. `av.VideoFrame` -> `np.ndarray`
                # 2. Force 'rgba' format
                self._numpy_ndarray = self._av_videoframe.to_ndarray(format = FORMAT_TO_FORCE)

                # self._numpy_ndarray = FrameHelper.videoframe_to_ndarray(
                #     frame = self._av_videoframe,
                #     do_include_alpha = 'a' in self.format
                # )
            elif self._moderngl_texture is not None:
                # 1. `moderngl.Texture` -> `np.ndarray`
                frame = _texture_to_numpy_frame(
                    texture = self._moderngl_texture,
                    size = (self.width, self.height)
                )

                # 2. Force 'rgba' format
                frame = frame.reformat(format = FORMAT_TO_FORCE)

                self._numpy_ndarray = frame

                # self._numpy_ndarray = FrameHelper.texture_to_ndarray(
                #     frame = self._moderngl_texture
                # )

        if do_write:
            self._av_videoframe = None
            self._moderngl_texture = None
            self.source_domain = FrameBufferDomain.CPU
            
        return self._numpy_ndarray
    
    def as_moderngl_texture(
        self,
        moderngl_context: 'moderngl.Context',
        do_write: bool = False
    ) -> 'moderngl.Texture':
        """
        Get the frame but as a `moderngl.Texture`. This
        is meant for GPU use.

        This method will lazy load the value from the
        one that is available in the moment it is
        requested.

        If you will modify the return, set the `do_write`
        as `True`.

        The `do_write` parameter is to indicate that
        when you call the method with `do_write=True`
        is because you'll modify the instance and
        that one will become the new absolute truth,
        so the other types will be removed and this one,
        as returned as a reference, will be modified by
        your code and become the new valid one.

        Remember that when obtaining this parameter
        """
        self._validate_source_available()

        if self._moderngl_texture is None:
            if self._numpy_ndarray is not None:
                # 1. Force 'rgba' format
                numpy_frame = self._numpy_ndarray.reformat(format = FORMAT_TO_FORCE)

                # 3. `np.ndarray` -> `moderngl.Texture`
                self._moderngl_texture = _numpy_frame_to_texture(
                    numpy_frame = numpy_frame,
                    moderngl_context = moderngl_context,
                    size = (self.width, self.height)
                )
                # self.gl_components = number_of_components


                # self._moderngl_texture = FrameHelper.ndarray_to_texture(
                #     frame = self._numpy_ndarray,
                #     # TODO: What do we do with the context? I think we should
                #     # create one standalone context if this is None...
                #     moderngl_context = moderngl_context,
                #     do_include_alpha = 'a' in self.format
                # )
            elif self._av_videoframe is not None:
                # 1. `av.VideoFrame` -> `np.ndarray`
                # 2. Force 'rgba' format
                numpy_frame = self._av_videoframe.to_ndarray(format = FORMAT_TO_FORCE)

                # 3. `np.ndarray` -> `moderngl.Texture`
                self._moderngl_texture = _numpy_frame_to_texture(
                    numpy_frame = numpy_frame,
                    moderngl_context = moderngl_context,
                    size = (self.width, self.height)
                )

                # self._moderngl_texture = FrameHelper.videoframe_to_texture(
                #     frame = self._av_videoframe,
                #     moderngl_context = moderngl_context,
                #     do_include_alpha = 'a' in self.format
                # )

        if do_write:
            self._av_videoframe = None
            self._numpy_ndarray = None
            self.source_domain = FrameBufferDomain.GPU

        return self._moderngl_texture
    
    @requires_dependency('av', 'yta_editor_frame', 'av')
    def as_av_videoframe(
        self,
        do_write: bool = False
    ) -> 'av.VideoFrame':
        """
        Get the frame but as a `av.VideoFrame`.

        This method will lazy load the value from the
        one that is available in the moment it is
        requested.
        """
        import av

        self._validate_source_available()

        if self._av_videoframe is None:
            if self._numpy_ndarray is not None:
                # 1. Force 'rgba' format
                numpy_array = self._numpy_ndarray.reformat(format = FORMAT_TO_FORCE)

                # 2. `np.ndarray` -> `av.VideoFrame`
                frame = av.VideoFrame.from_ndarray(numpy_array, format = FORMAT_TO_FORCE)

                self._av_videoframe = frame

                # self._av_videoframe = FrameHelper.ndarray_to_videoframe(
                #     frame = self._numpy_ndarray,
                #     do_include_alpha = 'a' in self.format
                # )
            elif self._moderngl_texture is not None:
                # 1. `moderngl.Texture` -> `np.ndarray`
                frame = _texture_to_numpy_frame(
                    texture = self._moderngl_texture,
                    size = (self.width, self.height)
                )

                # 2. `np.ndarray` -> `av.VideoFrame`
                frame = av.VideoFrame.from_ndarray(numpy_array, format = FORMAT_TO_FORCE)

                self._av_videoframe = frame
                
                # self._av_videoframe = FrameHelper.texture_to_videoframe(
                #     frame = self._moderngl_texture
                # )

        if do_write:
            self._numpy_ndarray = None
            self._moderngl_texture = None
            self.source_domain = FrameBufferDomain.MEDIA

        return self._av_videoframe

def _numpy_frame_to_texture(
    numpy_frame: 'np.ndarray',
    moderngl_context: 'moderngl.Context',
    size: tuple[int, int]
) -> 'moderngl.Texture':
    """
    *For internal use only*

    Create a texture from the given `numpy_frame`
    and return it.

    This method is only meant to be used within
    the `FrameBuffer` class.
    """
    number_of_components = (
        numpy_frame.shape[2]
        if numpy_frame.ndim == 3 else
        1
    )

    return moderngl_context.texture(
        size = size,
        components = number_of_components,
        data = numpy_frame.tobytes(),
        dtype = str(numpy_frame.dtype)
    )

@requires_dependency('numpy', 'yta_editor_frame', 'numpy')
def _texture_to_numpy_frame(
    texture: 'moderngl.Texture',
    size: tuple[int, int]
):
    import numpy as np

    data = texture.read()
    frame = np.frombuffer(data, dtype = texture.dtype)
    # frame = np.frombuffer(data, dtype = dtype)
    # The format is (height, width, number_of_components)
    frame = frame.reshape(size[1], size[0], texture.components)
    # frame = frame.reshape(size[1], size[0], number_of_components)

    # We need to flip due to our shaders. We could avoid
    # it by changing the vertex of the shaders
    frame = np.flipud(frame)

    # We always force 'RGBA'
    frame = frame.reformat(format = FORMAT_TO_FORCE)

    return frame

    
"""
A frame in HWC format (most common) is like this:
- `h, w = frame.shape[:2]`

If grayscale, is like this:
- `h, w = frame.shape`

If CHW, is like this:
- `c, h, w = frame.shape`
"""

"""
Puntos críticos (muy importantes en tu pipeline):

Flip vertical (OpenGL vs imagen)
OpenGL tiene el origen en bottom-left.
NumPy / PyAV esperan top-left.

Debes invertir verticalmente:
- `frame = np.flipud(frame)`

Número de canales:
- `texture.components`

3 → RGB
4 → RGBA

Tipo de dato
Por defecto:
- `dtype = uint8`

Si usas texturas float:
- `frame = np.frombuffer(data, dtype = np.float32)`

Contigüidad
Después del flipud, fuerza contigüidad:
- `frame = np.ascontiguousarray(frame)`
"""