# 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 Optional

from toolz.dicttoolz import get_in

from czc.api.environment import TestKeyEnvironment
from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czc.constants import CLOUDZERO_TEST_KEY_EVENT_PARAMETER
from czc.logging import get_logger

logger = get_logger()


def event_handler_test_key_aware(*, detail_path: Optional[list[str]] = None):
    f"""Make this iAPI Lambda Handler Test Key Aware

    If {CLOUDZERO_TEST_KEY_EVENT_PARAMETER} key is present in lambda event, set {CLOUDZERO_TEST_KEY} environment
    variable so that downstream processing (usually DynamoDB) honors it.

    Parameters
    ----------
    detail_path: str - A list of element names to get to where the test key parameter would be found.
        If not specified it defaults to ['detail'].

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """

    def _event_handler(handler):
        @functools.wraps(handler)
        def handler_wrapper(event, context, **kwargs):
            keys = [*(detail_path if detail_path is not None else ['detail'])]
            keys.append(CLOUDZERO_TEST_KEY_EVENT_PARAMETER)
            with TestKeyEnvironment(get_in(keys, event)):
                return handler(event, context, **kwargs)

        return handler_wrapper

    return _event_handler
