# Copyright (c) 2020 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

from functools import wraps
from typing import Type


import boto3

from aws_lambda_powertools.logging import Logger
from aws_lambda_powertools.tracing import Tracer
from pydantic import BaseModel

from czc.internal_api.handler import iapi_test_key_aware
from czc.logging import log_context_from_event
from czc.parser import parsed_event


events_client = boto3.client('events')
logger = Logger()
tracer = Tracer()


def state_handler(model: Type[BaseModel]):
    def wrapper(handler):

        @wraps(handler)
        @tracer.capture_lambda_handler
        @log_context_from_event(logger)
        @parsed_event(model=model)
        @iapi_test_key_aware()
        def inner(*args, **kwargs):
            return handler(*args, **kwargs)
        return inner
    return wrapper
