from functools import wraps
from typing import Optional, Type

from maps_adv.common.protomallow import ProtobufSchema

__all__ = ["with_schemas"]


def with_schemas(
    input_schema: Optional[Type[ProtobufSchema]] = None,
    output_schema: Optional[Type[ProtobufSchema]] = None,
):
    """Decorator for serializing input arguments and deserializing returned value."""

    def decorator(func):
        @wraps(func)
        async def wrapper(s, data=None, **kwargs):
            if input_schema is not None:
                if data is None:
                    raise RuntimeError(
                        "This api_provider requires protobuf message input"
                    )
                dumped = input_schema().from_bytes(data)
                kwargs.update(dumped)
            got = await func(s, **kwargs)
            if output_schema is not None:
                return output_schema().to_bytes(got)

        return wrapper

    return decorator
