from yta_editor_frame.buffer import FrameBuffer
from yta_validation import PythonValidator
from typing import Union
from functools import wraps
from collections.abc import Mapping


def force_framebuffer(
    *arg_names
):
    """
    Decorator to transform the parameters provided into
    `FrameBuffer` instances, if the parameters are
    one of the following types:
    - `moderngl.Texture`
    - `np.ndarray`
    - `av.VideoFrame`

    This method could need one of the following optional
    libraries
    to work:
    - `av`
    - `moderngl`
    - `numpy`

    The code must be like this below:
    ```
    @force_framebuffer('inputs')
    def test(
        inputs: dict[str, any]
    ):
        # Here the 'inputs' is now a dict with 'FrameBuffer'
        pass
    ```
    """
    def decorator(
        func
    ):
        @wraps(func)
        def wrapper(
            *args,
            **kwargs
        ):
            def convert_to_framebuffer(
                input: Union['FrameBuffer', 'moderngl.Texture', 'np.ndarray', 'av.VideoFrame', Mapping, None]
            ):
                """
                Transform all the inputs provided in the `input`
                parameter to a `FrameBuffer` instance. The `input`
                parameter can be a single input or a dict of them.
                """
                # We can have something like 'base_input': None
                # that we accept for some specific nodes
                if input is None:
                    return input

                if not PythonValidator.is_instance_of(input, ['FrameBuffer', 'Texture', 'ndarray', 'VideoFrame', Mapping]):
                    raise Exception(f'We cannot convert "{type(input)}" to FrameBuffer.')
                
                return (
                    {
                        k: convert_to_framebuffer(v) for k, v in input.items()
                    }
                    if PythonValidator.is_instance_of(input, Mapping) else
                    FrameBuffer.from_frame(input)
                )

            for name in arg_names:
                if name in kwargs:
                    kwargs[name] = convert_to_framebuffer(kwargs[name])

            return func(*args, **kwargs)
        return wrapper
    return decorator