from yta_editor_common.frame.core.domain import FrameBufferDomain
from yta_editor_common.frame.helper import FrameHelper
from yta_constants.enum import YTAEnum as Enum
from yta_programming.decorators.requires_dependency import requires_dependency
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.
    """

    RGB = 'rgb24'
    """
    The `rgb24` format.
    """
    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.
    """
    FLOAT32 = 'float32'
    """
    The `float32` dtype.
    """


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

    @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).
        """
        return [
            domain
            for condition, domain in (
                (self._moderngl_texture is not None, FrameBufferDomain.GPU),
                (self._numpy_ndarray is not None, FrameBufferDomain.CPU),
                (self._av_videoframe is not None, FrameBufferDomain.MEDIA),
            )
            if condition
        ]
    
    size: tuple[int, int]
    """
    The size of the frame, in pixels, expressed as a
    `(width, height)` tuple.
    """
    format: FrameBufferFormat
    """
    The format of the frame, that must be a 
    `FrameBufferFormat`.
    """
    dtype: FrameBufferDType
    """
    The `dtype` of the frame (based on numpy), that
    must be a `FrameBufferDType`.
    """
    is_data_normalized: bool
    """
    A flag to indicate if the values are set in the
    `[0, 1]` range or in the `[0, 255]`.
    """
    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)
    """
    The frame as a `av.VideoFrame`.
    """
    _numpy_ndarray: Union['np.ndarray', None] = field(default = None, init = False)
    """
    The frame as a `numpy.ndarray`.
    """
    _moderngl_texture: Union['moderngl.Texture', None] = field(default = None, init = False)
    """
    The frame as a `moderngl.Texture`.
    """

    def __post_init__(
        self
    ):
        self.format = FrameBufferFormat.to_enum(self.format)
        self.dtype = FrameBufferDType.to_enum(self.dtype)

    @classmethod
    @requires_dependency('av', 'yta_editor_common', 'av')
    def from_av_videoframe(
        cls,
        frame: 'av.VideoFrame'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `av.VideoFrame`.
        """
        format = frame.format.name
        # Special cases, force 'rgba'
        if frame.format.name == 'yuv420p':
            frame = frame.reformat(format = 'rgba')
            format = 'rgba'

        frame_buffer = cls(
            size = (frame.width, frame.height),
            format = FrameBufferFormat.to_enum(format),
            # dtype = FrameBufferDType.to_enum(frame.to_ndarray().dtype),
            # 99% of the time is 'uint8' and '.to_ndarray()' is expensive
            dtype = FrameBufferDType.UINT8,
            # TODO: I don't know much about this parameter
            is_data_normalized = False
            
        )

        """
        It would be nice to force this:
        - format = RGBA & dtype = UINT8
        because we will end up soon doing
        `frame.to_ndarray(format="rgba")` to work with
        numpy or moderngl
        """
        frame_buffer._av_videoframe = frame
        frame_buffer.source_domain = FrameBufferDomain.MEDIA

        return frame_buffer
    
    @classmethod
    @requires_dependency('numpy', 'yta_editor_common', 'numpy')
    def from_numpy_ndarray(
        cls,
        frame: 'np.ndarray'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `numpy.ndarray`.
        """
        height, width = frame.shape[:2]

        if frame.ndim != 2:
            raise Exception('Invalid ndarray shape for frame buffer')

        # TODO: I'm ignoring the GRAY by now
        # Infer format from channel dimension
        # if frame.ndim == 2:
        #     fmt = FrameBufferFormat.GRAY
        channels = frame.shape[2]

        if channels not in [3, 4]:
            raise Exception(f'The amount of channels ({channels}) is not supported.')
        
        format = (
            FrameBufferFormat.RGB
            if channels == 3 else
            FrameBufferFormat.RGBA
        )

        frame_buffer = cls(
            size = (width, height),
            format = format,
            dtype = FrameBufferDType.to_enum(frame.dtype),
            # TODO: I don't know much about this parameter
            is_data_normalized = False
        )

        frame_buffer._numpy_ndarray = frame
        frame_buffer.source_domain = FrameBufferDomain.CPU

        return frame_buffer
    
    @classmethod
    @requires_dependency('moderngl', 'yta_editor_common', 'moderngl')
    def from_moderngl_texture(
        cls,
        texture: 'moderngl.Texture'
    ) -> 'FrameBuffer':
        """
        Instantiate a `FrameBuffer` from a `moderngl.Texture`.
        """
        width, height = texture.size
        components = texture.components

        format_map = {
            # 1: FrameBufferFormat.R,
            # 2: FrameBufferFormat.RG,
            3: FrameBufferFormat.RGB,
            4: FrameBufferFormat.RGBA,
        }

        # If exception it is because we don't accept that format
        format = format_map.get(components, FrameBufferFormat.RGBA)

        dtype_map = {
            # 'f1': FrameBufferDType.FLOAT8,
            # 'f2': FrameBufferDType.FLOAT16,
            'f4': FrameBufferDType.FLOAT32,
            'u1': FrameBufferDType.UINT8,
        }

        # If exception it is because we don't accept that dtype
        dtype = dtype_map.get(texture.dtype, FrameBufferDType.UINT8)

        frame_buffer = cls(
            size = (width, height),
            format = format,
            dtype = dtype,
            is_data_normalized = False
        )

        frame_buffer._moderngl_texture = texture
        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.')
    
    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.
        """
        self._validate_source_available()

        if self._numpy_ndarray is None:
            if self._av_videoframe is not None:
                self._numpy_ndarray = FrameHelper.videoframe_to_ndarray(
                    frame = self._av_videoframe,
                    do_include_alpha = 'a' in self.format.value
                )
            elif self._moderngl_texture is not None:
                self._numpy_ndarray = FrameHelper.texture_to_ndarray(
                    frame = self._moderngl_texture
                )[0]

        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:
                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.value
                )
            elif self._av_videoframe is not None:
                # TODO: If 'self._numpy_ndarray' is None we could set it here as
                # we are actually calculating it, right (?)
                self._moderngl_texture = FrameHelper.videoframe_to_texture(
                    frame = self._av_videoframe,
                    moderngl_context = moderngl_context,
                    do_include_alpha = 'a' in self.format.value
                )

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

        return self._moderngl_texture
    
    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.
        """
        self._validate_source_available()

        if self._av_videoframe is None:
            if self._numpy_ndarray is not None:
                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:
                # TODO: Should we be able to force alpha if the
                # texture doesn't include it (?)
                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

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