from yta_video_opengl.context.provider import OpenGLContextProvider
from yta_editor_time.evaluation_context import EvaluationContext
from yta_editor_frame.decorators import force_framebuffer
from yta_editor_frame.buffer import FrameBuffer
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,
        processor_cpu: Union['_ProcessorGPU', None] = None,
        processor_gpu: Union['_ProcessorCPU', None] = None,
        opengl_context: Union['moderngl.Context', 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._opengl_context: Union['moderngl.Context', None] = (
            OpenGLContextProvider().context
            if opengl_context is None else
            opengl_context
        )
        """
        *For internal use only*

        The context that we will use for the GPU node.
        """

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

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

        Check if we should use GPU or not, according
        to what the user is requesting with the
        `do_use_gpu` parameter 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 = do_use_gpu
        )

    # TODO: Maybe move to a more general thing (?)
    def _prepare_input(
        self,
        input: 'FrameBuffer',
        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.
        """
        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, 'FrameBuffer'],
        do_use_gpu: bool = True
    ) -> 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, 'FrameBuffer')

        # 2. Prepare for the node (GPU or CPU)
        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

    @force_framebuffer('inputs')
    def process(
        self,
        inputs: dict[str, 'FrameBuffer'],
        evaluation_context: EvaluationContext,
        output_size: Union[tuple[int, int], None],
        do_use_gpu: bool = True,
        parameters: dict[str, any] = {}
    ) -> Union['moderngl.Texture', 'np.ndarray']:
        """
        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`.
        """
        inputs = self._validate_and_prepare_inputs_for_node(
            inputs = inputs,
            do_use_gpu = do_use_gpu
        )

        first_input = next(iter(inputs.values()))

        # Output size of first input if not forced
        output_size = (
            first_input.size
            if output_size is None else
            output_size
        )

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

        processor = self._get_processor(
            do_use_gpu = do_use_gpu
        )

        return processor.process(
            inputs = inputs,
            evaluation_context = evaluation_context,
            output_size = output_size,
            # TODO: Should we pass 'dict' and parse again instead (?)
            **parameters
        )