from yta_editor_frame.frame import Frame
from yta_validation import PythonValidator
from yta_validation.parameter import ParameterValidator
from typing import Union
from functools import wraps
from collections.abc import Mapping


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

    A `RenderContext` instance is needed as
    `render_context` to be able to make the
    transformation.

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

    The code must be like this below:
    ```
    @force_frame('inputs')
    def test(
        inputs: dict[str, any]
    ):
        # Here the 'inputs' is now a dict with 'Frame'
        pass
    ```
    """
    def decorator(
        func
    ):
        @wraps(func)
        def wrapper(
            *args,
            **kwargs
        ):
            def convert_to_frame(
                input: Union['Frame', 'RenderTarget', 'moderngl.Texture', 'np.ndarray', 'av.VideoFrame', Mapping, None],
                render_context: 'RenderContext'
            ):
                """
                Transform all the inputs provided in the `input`
                parameter to a `Frame` 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 or
                    PythonValidator.is_instance_of(input, 'Frame')
                ):
                    return input

                if not PythonValidator.is_instance_of(input, ['Frame', 'RenderTarget', 'Texture', 'ndarray', 'VideoFrame', Mapping]):
                    raise Exception(f'We cannot convert "{type(input)}" to Frame.')
                
                return (
                    {
                        key: convert_to_frame(
                            input = input_item,
                            render_context = render_context
                        )
                        for key, input_item in input.items()
                    }
                    if PythonValidator.is_instance_of(input, Mapping) else
                    Frame.from_frame(
                        frame = input,
                        render_context = render_context
                    )
                )
            
            # args[0] is self
            if (
                len(args) == 0 or
                args[0]._render_context is None
            ):
                raise Exception('The "._render_context" attribute was not found.')
            
            render_context = args[0]._render_context

            ParameterValidator.validate_mandatory_instance_of('render_context', render_context, 'RenderContext')

            for name in arg_names:
                if name in kwargs:
                    kwargs[name] = convert_to_frame(
                        input = kwargs[name],
                        render_context = render_context
                    )

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