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.
    """

    """
    TODO: Maybe we don't need this one, if it is
    fully opaque we treat it as non-transparent,
    and if it is not opaque nor transparent, it
    means that it has some transparency so we have
    to handle it and this 'has_alpha_pixels' would
    be automatically yes.
    """
    has_alpha_pixels: Union[bool, None] = field(default = None)
    """
    Flag to indicate if the frame has alpha pixels.

    The value could be `None` if we don't actually
    know if yes or no (maybe because it came from an
    unknown `moderngl.Texture`).
    """
    is_fully_opaque: Union[bool, None] = field(default = None)
    """
    Flag to indicate if the frame is fully opaque,
    which means that the alpha channel is giving no
    transparency to this frame. This is useful for
    the blending process because we can stop blending
    when this frame appears (the next ones will not be
    shown due to the opacity of this one).

    The value could be `None` if we don't actually
    know if yes or no (maybe because it came from an
    unknown `moderngl.Texture`).
    """
    is_fully_transparent: Union[bool, None] = field(default = None)
    """
    Flag to indicate if the frame is fully transparent,
    which means that the alpha channel is giving full
    transparency to this frame. This is ueful for the
    blending process because we can omit it when its
    not the only frame to combine (to blend).

    The value could be `None` if we don't actually
    know if yes or no (maybe because it came from an
    unknown `moderngl.Texture`).
    """

    # 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.has_gpu:
            domains.append(FrameBufferDomain.GPU)

        if self.has_cpu:
            domains.append(FrameBufferDomain.CPU)

        if self.has_media:
            domains.append(FrameBufferDomain.MEDIA)

        return domains
    
    @property
    def has_gpu(
        self
    ) -> bool:
        """
        Boolean flag to indicate if there is GPU
        information stored, which means that the
        `self._moderngl_texture` is not `None`.
        """
        return self._moderngl_texture is not None
    
    @property
    def has_cpu(
        self
    ) -> bool:
        """
        Boolean flag to indicate if there is CPU
        information stored, which means that the
        `self._numpy_ndarray` is not `None`.
        """
        return self._numpy_ndarray is not None
    
    @property
    def has_media(
        self
    ) -> bool:
        """
        Boolean flag to indicate if there is Media
        information stored, which means that the
        `self._av_videoframe` is not `None`.
        """
        return self._av_videoframe is not None
    
    @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
        )
    
    def _autodetect_transparency(
        self
    ) -> 'FrameBuffer':
        """
        *For internal use only*

        Internal method to detect the transparency from
        the frame that is available, using this priority:
        - `cpu (np.ndarray)`
        - `media (av.VideoFrame)`
        - `gpu (moderngl.Texture)`
        """
        if self.has_cpu:
            if (
                self._numpy_ndarray.ndim == 3 and
                self._numpy_ndarray.shape[2] == 4
            ):
                alpha = self._numpy_ndarray[..., 3]
                a_min = alpha.min()
                a_max = alpha.max()
                self.has_alpha_pixels = a_min < 255
                self.is_fully_opaque = a_min == 255
                self.is_fully_transparent = a_max == 0
            else:
                self.has_alpha_pixels = False
                self.is_fully_opaque = True
                self.is_fully_transparent = False
        elif self.has_media:
            """
            TODO: We could force the '.reformat("rgba")' and store
            the corresponding numpy array because we will need it
            after to be processed (as CPU or to transform it to GPU).
            """
            self.has_alpha_pixels = 'a' in self._av_videoframe.format.name
            self.is_fully_opaque = (
                True
                if not self.has_alpha_pixels else
                # We don't know yet
                None
            )
            self.is_fully_transparent = (
                False
                if not self.has_alpha_pixels else
                # We don't know yet
                None
            )
        elif self.has_gpu:
            # TODO: By now we don't know it...
            self.has_alpha_pixels = None
            self.is_fully_opaque = None
            self.is_fully_transparent = None

        return self

    def _autocalculate_numpy_from_av(
        self,
        frame: 'av.VideoFrame'
    ) -> 'FrameBuffer':
        """
        *For internal use only*

        Autocalculate the `np.ndarray` from the `frame`
        given that must be an `av.VideoFrame`. We
        pre-calculate it because we will need it soon.
        """
        self._numpy_ndarray = frame.to_ndarray(format = FORMAT_TO_FORCE)
        # Detect the real transparency from the numpy
        self._autodetect_transparency()

        return self

    @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.
        """
        ParameterValidator.validate_mandatory_instance_of('frame', frame, 'VideoFrame')

        # TODO: We could reformat and transform to 'ndarray'
        # and read. We will need the 'ndarray' when processing
        frame_buffer = cls(
            size = (frame.width, frame.height)
        )
        frame_buffer._av_videoframe = frame
        frame_buffer.source_domain = FrameBufferDomain.MEDIA

        # Calculate and store the 'np.ndarray' because we
        # will need it later
        frame_buffer._autocalculate_numpy_from_av(frame)

        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.
        """
        ParameterValidator.validate_mandatory_numpy_array('frame', frame)

        height, width = frame.shape[:2]

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

        return frame_buffer

    @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.
        """
        ParameterValidator.validate_mandatory_instance_of('frame', frame, 'Texture')

        width, height = frame.size

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

        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 update_frame(
        self,
        frame: Union['np.ndarray', 'moderngl.Texture', 'av.VideoFrame']
    ) -> 'FrameBuffer':
        """
        Update the internal frame, setting it as the only
        source and removing the previous ones.
        """
        ParameterValidator.validate_mandatory_instance_of('frame', frame, ['ndarray', 'Texture', 'VideoFrame'])

        if PythonValidator.is_instance_of(frame, 'ndarray'):
            self._numpy_ndarray = frame
            self._av_videoframe = None
            self._moderngl_texture = None
        elif PythonValidator.is_instance_of(frame, 'Texture'):
            self._numpy_ndarray = None
            self._av_videoframe = None
            self._moderngl_texture = frame
        elif PythonValidator.is_instance_of(frame, 'VideoFrame'):
            self._numpy_ndarray = None
            self._av_videoframe = frame
            self._moderngl_texture = None

            self._autocalculate_numpy_from_av(frame)
        
        self._autodetect_transparency()
    
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def as_numpy_ndarray(
        self,
        do_flip_texture: bool = True,
        do_force_contiguous: 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.

        This method will always force the `rgba` format.
        """
        import numpy as np

        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
                frame = self._av_videoframe.to_ndarray(format = FORMAT_TO_FORCE)

                # Flipping it would be unexpected, so we don't
                frame = (
                    np.ascontiguousarray(frame)
                    if do_force_contiguous else
                    frame
                )

                self._numpy_ndarray = frame
                self._autodetect_transparency()
            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),
                    do_flip_texture = do_flip_texture,
                    do_force_contiguous = False
                )

                # 2. Force 'rgba' format
                frame = _ensure_rgba_numpy_array(
                    ndarray = frame,
                    do_force_contiguous = do_force_contiguous
                )

                self._numpy_ndarray = frame
                self._autodetect_transparency()

        return self._numpy_ndarray
    
    def as_moderngl_texture(
        self,
        moderngl_context: 'moderngl.Context',
        do_flip_texture: bool = True
    ) -> '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.

        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
                frame = _ensure_rgba_numpy_array(
                    ndarray = self._numpy_ndarray,
                    do_force_contiguous = False
                )

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

                # We preserve the information about the alpha
                # channel that we had when numpy array
            elif self._av_videoframe is not None:
                # 1. `av.VideoFrame` -> `np.ndarray`
                # 2. Force 'rgba' format
                frame = self._av_videoframe.to_ndarray(format = FORMAT_TO_FORCE)

                # As we obtained it as numpy array, preserve it and
                # the information about the alpha channel from it
                self._numpy_ndarray = frame
                self._autodetect_transparency()

                # 3. `np.ndarray` -> `moderngl.Texture`
                self._moderngl_texture = _numpy_frame_to_texture(
                    numpy_frame = frame,
                    moderngl_context = moderngl_context,
                    size = (self.width, self.height),
                    do_flip_texture = do_flip_texture
                )
                
        return self._moderngl_texture
    
    @requires_dependency('av', 'yta_editor_frame', 'av')
    def as_av_videoframe(
        self,
        do_flip_texture: bool = True
    ) -> '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
                frame = _ensure_rgba_numpy_array(
                    ndarray = self._numpy_ndarray,
                    do_force_contiguous = True
                )

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

                # We preserve the information about the alpha
                # channel that we had when numpy array
            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),
                    do_flip_texture = do_flip_texture,
                    do_force_contiguous = True
                )

                # As we obtained it as numpy array, preserve it and
                # the information about the alpha channel from it
                self._numpy_ndarray = frame
                self._autodetect_transparency()

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

        return self._av_videoframe

@requires_dependency('numpy', 'yta_editor_frame', 'numpy')
@requires_dependency('moderngl', 'yta_editor_frame', 'moderngl')
def _numpy_frame_to_texture(
    numpy_frame: 'np.ndarray',
    moderngl_context: 'moderngl.Context',
    size: tuple[int, int],
    do_flip_texture: bool = True
) -> 'moderngl.Texture':
    """
    *For internal use only*

    Create a `moderngl.Texture` from the given
    `numpy_frame` and return it.

    This method is only meant to be used within
    the `FrameBuffer` class.
    """
    import moderngl
    import numpy as np

    numpy_frame = (
        np.flipud(numpy_frame)
        if do_flip_texture else
        numpy_frame
    )

    number_of_components = (
        numpy_frame.shape[2]
        if numpy_frame.ndim == 3 else
        1
    )

    # We force contiguous for moderngl
    numpy_frame = np.ascontiguousarray(numpy_frame)

    """
    We are always using 'np.uint8' because we are also
    forcing 'f1' textures dtype.
    """
    # TODO: Sometimes we can receive a 'float32', why (?)
    # TODO: Transform dtype to texture dtype with utils
    dtype = (
        'f1'
        if numpy_frame.dtype.name == 'uint8' else
        'f4'
        if numpy_frame.dtype.name == 'float32' else
        # This is not ok, but by now...
        'f1'
    )

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

    # Default filter
    texture.filter = (moderngl.NEAREST, moderngl.NEAREST)

    return texture

@requires_dependency('numpy', 'yta_editor_frame', 'numpy')
def _texture_to_numpy_frame(
    texture: 'moderngl.Texture',
    size: tuple[int, int],
    do_flip_texture: bool = True,
    do_force_contiguous: bool = False
):
    """
    *For internal use only*

    Create a `np.ndarray` from the `texture` provided
    and return it.

    This method is only meant to be used within
    the `FrameBuffer` class.

    This frame is not forced to be contiguous.
    """
    import numpy as np

    data = texture.read()
    """
    We are always using 'np.uint8' because we are also
    forcing 'f1' textures dtype.
    """
    frame = np.frombuffer(data, dtype = np.uint8)
    # The format is (height, width, number_of_components)
    frame = frame.reshape(size[1], size[0], texture.components)

    # We need to flip it depending on our moderngl shaders
    frame = (
        np.flipud(frame)
        if do_flip_texture else
        frame
    )

    # We force it as contiguous
    frame = (
        np.ascontiguousarray(frame)
        if do_force_contiguous else
        frame
    )

    # # We always force 'RGBA'
    # frame = _ensure_rgba_numpy_array(frame)

    return frame

@requires_dependency('numpy', 'yta_editor_frame', 'numpy')
def _ensure_rgba_numpy_array(
    ndarray: 'np.ndarray',
    do_force_contiguous: bool = False
) -> 'np.ndarray':
    """
    *For internal use only*

    Transform the `ndarray` provided to make sure
    that it is representing a 'rgba' frame,
    compatible with textures like the following one:
    - `dtype = uint8`
    - `shape = (H, W, 4)`
    - `layout = RGBA`

    This method is only meant to be used within
    the `FrameBuffer` class.

    This array is forced to be contiguous if the
    `do_force_contiguous` parameter is `True`, that
    must be set as `True` when we are about to use
    that frame to build a texture o to write with
    the pyav library.
    """
    import numpy as np

    if ndarray.ndim == 2:
        ndarray = np.repeat(ndarray[:, :, None], 3, axis = 2)

    if ndarray.shape[2] == 3:
        alpha = np.full((ndarray.shape[0], ndarray.shape[1], 1), 255, dtype = ndarray.dtype)
        ndarray = np.concatenate((ndarray, alpha), axis = 2)

    if ndarray.shape[2] != 4:
        raise Exception('The format of this numpy array is not accepted.')
    
    ndarray = (
        np.ascontiguousarray(ndarray)
        if do_force_contiguous else
        ndarray
    )

    return ndarray

    
"""
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)`
"""