# Copyright (c) 2022 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import functools
from typing import Dict, Optional, Type

from aws_lambda_powertools.utilities.parser import parse
from pydantic import BaseModel
from voluptuous import Schema


def parsed_event(model: Type[BaseModel]):
    """Parse lambda event into Pydantic `model`, which is added to lambda handler
    kwargs as `parsed_event`.

    Example
    -------

    **Create BaseModel and use in lambda handler**

        class Event(BaseModel):
            id: int
            ...

        @parsed_event(Event)
        def handler(event: Order, context: LambdaContext, parsed_event: Event):
            assert isinstance(parsed_event, Event)


    Parameters
    ----------
    model:    pydantic.BaseModel
        Lambda event model

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    def parsed_event_handler(handler):
        @functools.wraps(handler)
        def handler_wrapper(event: Dict, context, **kwargs):
            parsed_event = parse(event, model=model)
            return handler(event, context, parsed_event=parsed_event, **kwargs)
        return handler_wrapper
    return parsed_event_handler


def uses_event_parameters(schema: Optional[Schema] = None, passthrough_event: bool = False, passthrough_context: bool = False):
    """Parse lambda event with the given `voluptuous.Schema`. The root keys of the event are
    added to the lambda handler kwargs.

    Example
    -------

    **Create Schema and use in lambda handler without passthrough**

        schema = Schema({
            'id': str
        }, required=True, extra=REMOVE_EXTRA)

        @uses_event_parameters(schema=schema)
        def handler(id: str):
            assert isinstance(id, str)

    **Create Schema and use in lambda handler with passthrough**

        schema = Schema({
            'id': str
        }, required=True, extra=REMOVE_EXTRA)

        @uses_event_parameters(schema=schema, passthrough_event=True, passthrough_context=True)
        def handler(event: Mapping, context: LambdaContext, id: str):
            assert isinstance(event, dict)
            assert isinstance(context, LambdaContext)
            assert isinstance(id, str)


    Parameters
    ----------
    schema:    voluptuous.Schema
        Dictionary schema for lambda event
    passthrough_event: bool, optional default to False
        Whether or not to pass raw lambda event through to handler args
    passthrough_context: bool, optional default to False
        Whether or not to pass raw lambda context through to lambda args

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    def _uses_event_parameters(handler):
        @functools.wraps(handler)
        def handler_wrapper(event, context, **kwargs):
            params = schema(event) if schema else event
            args = [
                *([event] if passthrough_event else []),
                *([context] if passthrough_context else [])
            ]
            return handler(*args, **params, **kwargs)

        return handler_wrapper

    return _uses_event_parameters
