from yta_editor_frame.frame import Frame

import numpy as np


class _FrameOperation:
    """
    *For internal use only*

    Class to indicate an operation that must be done
    within a frame.
    """

    def apply_gpu(
        self,
        frame: 'moderngl.Texture',
        render_context: 'RenderContext'
    ) -> 'moderngl.Texture':
        """
        Perform this frame operation by using the GPU.
        """
        raise NotImplementedError

    def apply_cpu(
        self,
        frame: 'np.ndarray'
    ) -> 'np.ndarray':
        """
        Perform this frame operation by using the CPU.
        """
        raise NotImplementedError
    
class ResizeWithPaddingOperation(_FrameOperation):
    """
    Operation to resize a frame with padding.
    """

    def __init__(
        self,
        size: tuple[int, int]
    ):
        self.size: tuple[int, int] = size
        """
        The expected size of the frame after adapting it.
        """

    def apply_gpu(
        self,
        texture: 'moderngl.Texture',
        render_context: 'RenderContext'
    ):
        width_output, height_output = self.size
        width_input, height_input = texture.size

        scale_x = width_input / width_output
        scale_y = height_input / height_output

        opengl_context = render_context._render_resources.opengl_context

        # TODO: Maybe a 'RenderTarget' is better (?)
        # render_target = render_context._render_resources._render_target_provider.acquire(
        #     width = width_output,
        #     height = height_output,
        #     number_of_components = 4,
        #     dtype = 'f1'
        # )
        output_texture = render_context._render_resources.texture.acquire(
            width = width_output,
            height = height_output,
            number_of_components = 4,
            dtype = 'f1'
        )
        # output_tex = ctx.texture((width_output, height_output), 4)
        fbo = opengl_context.framebuffer([output_texture])

        fbo.use()
        opengl_context.clear(0, 0, 0, 0)

        shader = get_cached_shader(ctx, "resize_pad")

        shader["scale"].value = (scale_x, scale_y)

        texture.use(0)
        render_fullscreen_quad(shader)

        return output_texture
    
    def apply_cpu(
        self,
        frame: 'Frame'
    ):
        numpy_array = _adapt_with_alpha(
            frame = frame.as_numpy_ndarray(),
            size = self.size
        )

        frame = Frame.from_frame(numpy_array)
        """
        We added and alpha background, so we have transparent
        pixels and the alpha layer.
        """
        frame.metadata.has_transparent_pixels = True
        frame.metadata.is_fully_transparent = False
        frame.metadata.is_fully_opaque = False

        return frame

        return super().apply_cpu(frame)
    
"""
Utils below
"""
def _adapt_with_alpha(
    frame: np.ndarray,
    size: tuple[int, int]
) -> np.ndarray:
    """
    Creates a fully alpha background with the `size`
    given and puts the `frame` provided in the middle of
    it.

    This is the functionality to be done when using the
    `ALPHA_BACKGROUND` option.
    """
    if frame.ndim != 3:
        raise Exception('The frame must be HxWxC.')

    src_h, src_w, src_c = frame.shape

    if src_c not in (3, 4):
        raise Exception('The frame must be RGB or GBA.')

    # Final transparent frame
    dst = np.zeros((size[1], size[0], 4), dtype = frame.dtype)

    # RGBA if necessary
    if src_c == 3:
        src_rgba = np.empty((src_h, src_w, 4), dtype = frame.dtype)
        src_rgba[..., :3] = frame
        src_rgba[..., 3] = np.iinfo(frame.dtype).max
    else:
        src_rgba = frame

    # Offsets to center
    offset_x = (size[0] - src_w) // 2
    offset_y = (size[1] - src_h) // 2

    if (
        offset_x < 0 or
        offset_y < 0
    ):
        # TODO: Do rescale and crop if needed
        raise Exception('Original frame is bigger than the size requested.')

    # Copy it
    dst[
        offset_y:offset_y + src_h,
        offset_x:offset_x + src_w
    ] = src_rgba

    # Force it to be continuous to avoid PyAv errors
    dst = np.ascontiguousarray(dst)

    return dst

# TODO: Create the GPU node that does this
