from yta_editor_nodes.utils import validate_render_target
from yta_editor_time.evaluation_context import EvaluationContext
from yta_editor_nodes.utils import render_context_to_gpu_render_context, render_context_to_cpu_render_context
from yta_editor_frame.decorators import force_frame
from yta_editor_frame.frame import Frame
from yta_editor_nodes_common import _Node
from typing import Union


class _ProcessorGPUAndCPU(_Node):
    """
    *For internal use only*

    *Abstract class*

    Abstract class to share the common behaviour of
    being able to handle a process with a GPU and/or
    a CPU processor, chosen by the user.

    Mandatory inputs:
    - `None`

    Optional inputs:
    - `None`
    """

    mandatory_inputs = []
    optional_inputs = []
    parameters_specifications = []

    @property
    def is_gpu_available(
        self
    ) -> bool:
        """
        Boolean flag to indicate if the GPU is available
        or not, that means that the processor that uses
        GPU is set.
        """
        return self._processor_gpu is not None
    
    @property
    def is_cpu_available(
        self
    ) -> bool:
        """
        Boolean flag to indicate if the CPU is available
        or not, that means that the processor that uses
        CPU is set.
        """
        return self._processor_cpu is not None
    
    def __init__(
        self,
        render_context: 'RenderContext',
        processor_cpu: Union['_ProcessorGPU', None] = None,
        processor_gpu: Union['_ProcessorCPU', None] = None
    ):
        """
        The `processor_cpu` and `processor_gpu` have to be
        set by the developer when building the specific
        classes, but the `do_use_gpu` boolean flag will be
        set by the user when instantiating the class to
        choose between GPU and CPU.

        The `opengl_context` is needed to instantiate the
        GPU unit but also to transform the input into a
        `moderngl.Texture` if needed.
        """
        if (
            processor_cpu is None and
            processor_gpu is None
        ):
            raise Exception('No node processor provided. At least one node processor is needed.')

        self._processor_cpu: Union['_ProcessorCPU', None] = processor_cpu
        """
        *For internal use only*

        The transition processor that is able to do the
        processing by using the CPU. If it is None we cannot
        process it with CPU.
        """
        self._processor_gpu: Union['_ProcessorGPU', None] = processor_gpu
        """
        *For internal use only*

        The transition processor that is able to do the
        processing by using the GPU. If it is None we cannot
        process it with GPU.
        """
        self._render_context: 'RenderContext' = render_context
        """
        *For internal use only*

        The `RenderContext` instance that stores all the
        configuration and resources we need to render,
        including the `opengl_context`.
        """

    def _get_processor(
        self
    ) -> Union['ProcessorGPU', 'ProcessorCPU']:
        """
        *For internal use only*

        Get the processor according to the configuration
        set in the `RenderContext` and if the one
        requested is available or not.
        """
        return (
            (
                # Prefer GPU if available
                self._processor_gpu or
                self._processor_cpu
            ) if self._render_context.render_config.do_use_gpu else (
                # Prefer CPU if available
                self._processor_cpu or
                self._processor_gpu
            )
        )
    
    def _do_use_gpu(
        self
    ) -> bool:
        """
        *For internal use only*

        Check if we should use GPU or not, according
        to what the user is requesting within the
        `RenderContext`, but checking if it is
        available in the node itself.

        This method is useful to know if we have to
        transform a frame to a specific format (for
        GPU or CPU) before sending it as an argument.
        """
        from yta_editor_nodes.utils import _do_use_gpu

        return _do_use_gpu(
            node_processor = self,
            do_use_gpu = self._render_context.render_config.do_use_gpu
        )

    # TODO: Maybe move to a more general thing (?)
    def _prepare_input(
        self,
        input: 'Frame',
        do_use_gpu: bool
    ) -> Union['moderngl.Texture', 'np.ndarray']:
        """
        *For internal use only*

        Prepare the `input` provided based on the `do_use_gpu`
        parameter given to be used with the processer,
        according to the availability.

        This method will transform a `texture` into a `numpy`
        array if they are asking for CPU and it is available,
        or a `numpy` into a `texture` if they are asking for
        GPU and it is available.
        """
        # TODO: We need to transform it
        from yta_editor_frame.frame.resource.converter import FrameResourceConverter

        """
        This modifies the instance, but I think it is
        the correct way to do it. The code below is
        creating a new one that is passed through...
        """
        return (
            # Return `moderngl.Texture`
            input.as_texture().resource.frame
            if do_use_gpu else
            # Return `np.ndarray`
            input.as_numpy().resource.frame
        )

        # TODO: Remove this below when above is working
        frame_resource = (
            FrameResourceConverter.to_texture(
                frame = input.resource,
                moderngl_context = self._render_context.render_resources.opengl_context)
            if self._do_use_gpu(do_use_gpu) else
            FrameResourceConverter.to_numpy(input.resource)
        )

        frame = Frame.from_frame(
            frame = frame_resource,
            render_context = self._render_context,
            metadata = {}
        )
        # Copy metadata from the previous one
        frame.metadata = input.metadata.copy(frame)

        # TODO: Now maybe we have to release the `input`,
        # but the transformed data should be stored inside
        # the same `frame` instance to not lose the info...

        return frame

        # Old code below (when FrameBuffer was being applied)
        return (
            input.as_moderngl_texture(
                moderngl_context = self._opengl_context
            )
            if self._do_use_gpu(do_use_gpu) else
            input.as_numpy_ndarray()
        )
    
    def _validate_and_prepare_inputs_for_node(
        self,
        inputs: dict[str, 'Frame'],
        do_use_gpu: bool
    ) -> dict[str, Union['np.ndarray', 'moderngl.Texture']]:
        """
        *For internal use only*

        Validate the `inputs` provided and prepare and
        convert them (if needed) to be used in CPU or GPU
        according to the `do_use_gpu` parameter.
        
        This must be done before sending the `inputs` to
        the node processor.
        """
        # 1. Validate
        super()._validate_inputs(inputs, 'Frame')

        """
        2. Prepare for the node (GPU or CPU)
        This will obtain the `np.ndarray` or the
        `moderngl.Texture` that will be send directly
        (unwrapped) to the specific nodes that accept
        only that raw input.
        """
        # 2. Prepare for the node (GPU or CPU), obtaining
        # the `np.ndarray` or `moderngl.Texture` directly
        for input_name in inputs:
            inputs[input_name] = (
                self._prepare_input(
                    input = inputs[input_name],
                    do_use_gpu = do_use_gpu
                )
                if inputs[input_name] is not None else
                None
            )

        return inputs
    
    def _format_parameters(
        self,
        parameters: dict[str, any]
    ) -> dict[str, any]:
        """
        *For internal use only*

        Format the parameters to be sent to the node.
        This method must be called after the parameters
        have been validated and prepared, and before
        sending them to the node so we can format some
        of them.

        For example, we use `x` and `y` parameters that
        we transform into a `position` tuple when being
        sent to the node. This is what must be done here.
        """
        return parameters
    
    def use_gpu(
        self
    ) -> '_Blender':
        """
        Activate the desire to use GPU (and no CPU) by
        setting the `do_use_gpu` parameter to `True`.
        """
        self._render_context.use_gpu()

        return self
    
    def use_cpu(
        self
    ) -> '_Blender':
        """
        Activate the desire to use CPU (and not GPU) by
        setting the `do_use_gpu` parameter to `False`.
        """
        self._render_context.use_cpu()

        return self

    @force_frame('inputs')
    def process(
        self,
        render_target: Union['RenderTarget', None],
        inputs: dict[str, 'Frame'],
        evaluation_context: EvaluationContext,
        parameters: dict[str, any] = {}
    ) -> Frame:
        """
        Prepare the `inputs` provided to be processed by
        this processor node according to if the GPU or
        CPU is requested and available, returning the
        result as a `moderngl.Texture` or a `np.ndarray`.
        """
        # Check if GPU is requested and available
        do_use_gpu = self._do_use_gpu()

        # We cannot GPU without 'render_target'
        validate_render_target(
            render_target = render_target,
            do_use_gpu = do_use_gpu
        )
        
        inputs = self._validate_and_prepare_inputs_for_node(
            inputs = inputs,
            do_use_gpu = do_use_gpu
        )

        parameters = self._validate_and_prepare_parameters(
            parameters = parameters
        )
        parameters = self._format_parameters(
            parameters = parameters
        )

        # The 'output_size' goes in the 'cpu_render_target'
        # when instantiated

        processor = self._get_processor()

        # This will be 'np.ndarray' (CPU) or 'RenderTarget' (GPU)
        processed_frame = processor.process(
            render_target = render_target,
            inputs = inputs,
            render_context = self._render_context,
            evaluation_context = evaluation_context,
            # TODO: Should we pass 'dict' and parse again instead (?)
            **parameters
        )

        return Frame.from_frame(
            frame = processed_frame,
            render_context = self._render_context,
            metadata = {}
        )
        # If it is a special node and we know some
        # metadata we could store it in the `process`
        # method we use to overwrite in the subclass

class _ProcessorGPUAndCPUConfigured(
    _ProcessorGPUAndCPU
):
    """
    Class to define the behaviour of our nodes that
    are able to handle the inputs and process them
    with GPU and/or CPU, depending on how they are
    defined.
    
    This class has been created to be able to
    separate behaviour and configuration.
    """

    mandatory_inputs = []
    optional_inputs = []
    parameters_specifications = []
    # Nodes
    cpu_class: Union['NodeProcessorCPU', None] = None
    gpu_class: Union['NodeProcessorGPU', None] = None

    def __init_subclass__(
        cls,
        **kwargs
    ):
        super().__init_subclass__(**kwargs)

        # Children must define them
        if not hasattr(cls, 'mandatory_inputs'):
            raise Exception('The "mandatory_inputs" is not defined.')
        
        if not hasattr(cls, 'optional_inputs'):
            raise Exception('The "optional_inputs" is not defined.')
        
        if not hasattr(cls, 'parameters_specifications'):
            raise Exception('The "parameters_specifications" is not defined.')

        if (
            len(cls.mandatory_inputs) == 0 and
            len(cls.optional_inputs) == 0
        ):
            raise Exception('The mandatory and optional inputs lists are empty.')
        
        cpu_cls = getattr(cls, 'cpu_class', None)
        gpu_cls = getattr(cls, 'gpu_class', None)

        if (
            not cls._is_marker and
            cpu_cls is None and
            gpu_cls is None
        ):
            raise Exception('No CPU nor GPU node classes define. At least one is needed.')

    def __init__(
        self,
        render_context: 'RenderContext'
    ):
        node_processor_cpu = (
            self.cpu_class(
                cpu_render_context = render_context_to_cpu_render_context(render_context)
            )
            if self.cpu_class is not None else
            None
        )

        node_processor_gpu = (
            self.gpu_class(
                gpu_render_context = render_context_to_gpu_render_context(render_context)
            )
            if self.gpu_class is not None else
            None
        )

        super().__init__(
            render_context = render_context,
            processor_cpu = node_processor_cpu,
            processor_gpu = node_processor_gpu
        )