"""
Module to transform one resource into another
one to be able to do CPU -> GPU, GPU -> CPU,
etc.
"""
from yta_editor_frame.frame.resource.identifier import FrameResourceIdentifier
from yta_editor_frame.frame.resource.cpu_frame_resource import CPUNumpyArrayFrameResource
from yta_editor_frame.frame.resource.gpu_frame_resource import GPUTextureFrameResource
from yta_editor_frame.frame.resource.av_video_frame_resource import AvVideoFrameResource
from yta_editor_frame.frame.resource.render_target_frame_resource import RenderTargetFrameResource
from yta_editor_frame.utils import copy_texture
from yta_programming.decorators.requires_dependency import requires_dependency
from typing import Union


FORMAT = 'rgba'

class FrameResourceConverter:
    """
    Class able to transform frames from one type to
    another.

    Singleton class for a specific `RenderContext`.
    """

    _instances: dict[int, 'FrameResourceConverter'] = {}

    @classmethod
    def get(
        cls,
        render_context: 'RenderContext'
    ) -> 'FrameResourceConverter':
        """
        Get the singleton instance for the specific
        `render_context` provided as parameter.
        """
        key = id(render_context)
        if key not in cls._instances:
            cls._instances[key] = cls(render_context)

        return cls._instances[key]
    
    def __init__(
        self,
        render_context: 'RenderContext'
    ):
        self.render_context: 'RenderContext' = render_context
        """
        The `RenderContext` associated to it.
        """

    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def to_numpy(
        self,
        resource: Union[CPUNumpyArrayFrameResource, GPUTextureFrameResource]
    ) -> CPUNumpyArrayFrameResource:
        """
        Transform the `resource` provided into a CPU
        `CPUNumpyArrayFrameResource` that holds a
        `np.ndarray`.
        """
        from yta_editor_frame.utils import _texture_to_numpy_frame

        if FrameResourceIdentifier.is_cpu(resource):
            # [1] `np.ndarray` -> `np.ndarray`
            return resource
        
        frame = None

        if FrameResourceIdentifier.is_cpu_av_videoframe(resource):
            # [2] `av.VideoFrame` -> `np.ndarray`
            import numpy as np

            numpy_array = resource.frame.to_ndarray(format = FORMAT)

            frame = numpy_array
        else:
            # [3] `moderngl.Texture` -> `np.ndarray`
            # [4] `render_target` -> `np.ndarray`

            # [3] `moderngl.Texture` -> `np.ndarray`
            texture = resource.frame
            if FrameResourceIdentifier.is_gpu_render_target(resource):
                # [4] `render_target` -> `np.ndarray`
                texture = texture.texture

            # `moderngl.Texture` -> `np.ndarray`
            numpy_array = _texture_to_numpy_frame(
                texture = texture,
                do_flip_texture = True,
                do_force_contiguous = False
            )

            frame = numpy_array

            # 2. Force 'rgba' format
            # TODO: Is this needed (?)
            # frame = _ensure_rgba_numpy_array(
            #     ndarray = numpy_array,
            #     # TODO: What do we do with the flip (?)
            #     # do_force_contiguous = do_force_contiguous
            #     do_force_contiguous = False
            # )

        return CPUNumpyArrayFrameResource(
            frame = frame,
            render_context = self.render_context
        )
        
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def to_texture(
        self,
        resource: 'FrameResource'
    ) -> GPUTextureFrameResource:
        """
        Transform the `resource` provided into a GPU
        `GPUTextureFrameResource` that holds a
        `moderngl.Texture`.
        """
        from yta_editor_frame.utils import _numpy_frame_to_texture

        if FrameResourceIdentifier.is_gpu(resource):
            # [1] `moderngl.Texture` -> `moderngl.Texture`
            return resource
    
        frame = None
        
        if FrameResourceIdentifier.is_gpu_render_target(resource):
            # [2] `render_target` -> `moderngl.Texture`
            frame = copy_texture(
                texture = resource.frame.texture,
                render_context = self.render_context,
                do_get_from_pool = True
            )
        elif (
            FrameResourceIdentifier.is_cpu(resource) or
            FrameResourceIdentifier.is_cpu_av_videoframe(resource)
        ):
            # [3] `np.ndarray` -> `moderngl.Texture`
            # [4] `av.VideoFrame` -> `moderngl.Texture`
            numpy_array = resource.frame
            if FrameResourceIdentifier.is_cpu_av_videoframe(resource):
                # [4] `av.VideoFrame` -> `moderngl.Texture`
                numpy_array = numpy_array.to_ndarray(format = FORMAT)

            texture = _numpy_frame_to_texture(
                numpy_frame = numpy_array,
                render_context = self.render_context,
                do_flip = True,
                do_from_pool = True
            )

            frame = texture

            # We preserve the information about the alpha
            # channel that we had when numpy array

        return GPUTextureFrameResource(
            frame = frame,
            render_context = self.render_context
        )
    
    @requires_dependency('av', 'yta_editor_frame', 'av')
    def to_videoframe(
        self,
        resource: 'FrameResource'
    ) -> AvVideoFrameResource:
        """
        Transform the `resource` provided into a CPU
        `AvVideoFrameResource` that holds a
        `av.VideoFrame`.
        """
        from yta_editor_frame.utils import _texture_to_numpy_frame

        import av
        
        if FrameResourceIdentifier.is_cpu_av_videoframe(resource):
            # [1] `av.VideoFrame` -> `av.VideoFrame`
            return resource
        
        frame = None

        if FrameResourceIdentifier.is_cpu(resource):
            # [2] `np.ndarray` -> `av.VideoFrame`
            from yta_editor_frame.utils import _ensure_rgba_numpy_array

            # 1. Force 'rgba' format
            numpy_array = _ensure_rgba_numpy_array(
                ndarray = resource.frame,
                do_force_contiguous = True
            )

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

            # We preserve the information about the alpha
            # channel that we had when numpy array
        else:
            # [3] `moderngl.Texture` -> `av.VideoFrame`
            # [4] `render_target` -> `av.VideoFrame`
            frame = resource.frame
            if FrameResourceIdentifier.is_gpu_render_target(resource):
                # [4] `render_target` -> `av.VideoFrame`
                frame = frame.texture

            # [3] `moderngl.Texture` -> `av.VideoFrame`
            # 1. `moderngl.Texture` -> `np.ndarray`
            numpy_array = _texture_to_numpy_frame(
                texture = frame,
                do_flip_texture = True,
                do_force_contiguous = True
            )

            # 2. `np.ndarray` -> `av.VideoFrame`
            frame = av.VideoFrame.from_ndarray(numpy_array, format = FORMAT)
        
        return AvVideoFrameResource(
            frame = frame,
            render_context = self.render_context
        )
    
    @requires_dependency('numpy', 'yta_editor_frame', 'numpy')
    def to_render_target(
        self,
        resource: 'FrameResource'
    ) -> RenderTargetFrameResource:
        """
        Transform the `resource` provided into a GPU
        `RenderTargetFrameResource` that holds a
        `RenderTarget`.
        """
        from yta_editor_frame.utils import _numpy_frame_to_texture

        if FrameResourceIdentifier.is_gpu_render_target(resource):
            # [1] `RenderTarget` -> `RenderTarget`
            return resource
        
        frame = None

        if FrameResourceIdentifier.is_gpu(resource):
            # [2] `moderngl.Texture` -> `RenderTarget`
            # Create a copy of the texture so the previous
            # owner can still have it and use it if needed
            texture = copy_texture(
                texture = resource.frame,
                render_context = self.render_context,
                do_get_from_pool = True
            )

            render_target = self.render_context.gpu.render_target.acquire(
                width = texture.width,
                height = texture.height,
                number_of_components = texture.components,
                dtype = texture.dtype
            )
            render_target.texture = frame
            
            frame = render_target
        elif (
            FrameResourceIdentifier.is_cpu(resource) or
            FrameResourceIdentifier.is_cpu_av_videoframe(resource)
        ):
            # [3] `np.ndarray` -> `RenderTarget`
            # [4] `av.VideoFrame` -> `RenderTarget`

            numpy_array = resource.frame
            if FrameResourceIdentifier.is_cpu_av_videoframe(resource):
                # [4] `av.VideoFrame` -> `RenderTarget`
                numpy_array = numpy_array.to_ndarray(format = FORMAT)

            # [3] `np.ndarray` -> `RenderTarget`
            # 1. Force 'rgba' format
            # TODO: Maybe this is necessary
            # numpy_array = _ensure_rgba_numpy_array(
            #     ndarray = numpy_array,
            #     do_force_contiguous = False
            # )

            # 2. `np.ndarray` -> `RenderTarget`
            frame = _numpy_frame_to_texture(
                numpy_frame = numpy_array,
                render_context = self.render_context,
                do_flip = True,
                do_from_pool = True
            )

            render_target = self.render_context.gpu.render_target.acquire(
                width = texture.width,
                height = texture.height,
                number_of_components = texture.components,
                dtype = texture.dtype
            )
            render_target.texture = frame

            # render_target.write(frame)

            frame = render_target

            # We preserve the information about the alpha
            # channel that we had when numpy array

        return RenderTargetFrameResource(
            frame = frame,
            render_context = self.render_context
        )


