# 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

from typing import List, Mapping

from aws_lambda_powertools.logging import Logger
from toolz.dicttoolz import keyfilter


class LogContext():
    _logger: Logger
    _context_keys: Mapping

    def __init__(self, logger: Logger, **additional_keys) -> None:
        self._logger = logger
        self._context_keys = additional_keys

    def __enter__(self):
        self._logger.append_keys(**self._context_keys)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._logger.remove_keys(self._context_keys)


def log_context_from_event(logger: Logger, exclude_fields: List[str] = None):
    def log_context_from_event_handler(handler):
        def function_wrapper(event, *args, **kwargs):
            additional_keys = keyfilter(lambda f: f not in (exclude_fields or []), event)
            with LogContext(logger, **additional_keys):
                return handler(event, *args, **kwargs)

        return function_wrapper
    return log_context_from_event_handler
