# 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 logging
import threading
from types import TracebackType
from typing import IO, Any, Callable, Dict, Mapping, List, Optional, Tuple, Type, Union

from aws_lambda_powertools import Tracer
from aws_lambda_powertools.logging import Logger as AwsLogger
from aws_lambda_powertools.logging.formatter import LambdaPowertoolsFormatter

MASKS_LOCK = threading.Lock()
MASKED_VALUE = "<HIDDEN>"
_masks: List[str] = []


class Formatter(LambdaPowertoolsFormatter):

    def format_masks(self, s: str):
        for mask in _masks:
            if mask in s:
                s = s.replace(mask, MASKED_VALUE)
        return s

    def format(self, record: logging.LogRecord) -> str:
        return self.format_masks(super().format(record))


class Logger(AwsLogger):
    def __init__(self, service: Optional[str] = None, level: Optional[Union[str, int]] = None, child: bool = False, sampling_rate: Optional[float] = None, stream: Optional[IO[str]] = None, logger_formatter: Formatter = Formatter(), logger_handler: Optional[logging.Handler] = None, log_uncaught_exceptions: bool = False, json_serializer: Optional[Callable[[Dict], str]] = None, json_deserializer: Optional[Callable[[Union[Dict, str, bool, int, float]], str]] = None, json_default: Optional[Callable[[Any], Any]] = None, datefmt: Optional[str] = None, use_datetime_directive: bool = False, log_record_order: Optional[List[str]] = None, utc: bool = False, use_rfc3339: bool = False, serialize_stacktrace: bool = True, **kwargs) -> None:
        self.logger_formatter: Formatter = logger_formatter  # type: ignore
        super().__init__(service, level, child, sampling_rate, stream, logger_formatter, logger_handler, log_uncaught_exceptions, json_serializer, json_deserializer, json_default, datefmt, use_datetime_directive, log_record_order, utc, use_rfc3339, serialize_stacktrace, **kwargs)

    def exception(self, msg: object, *args: object, exc_info: Optional[Union[bool, Tuple[Type[BaseException], BaseException, Optional[TracebackType]], Tuple[None, None, None], BaseException]] = True, stack_info: bool = False, stacklevel: int = 2, extra: Optional[Mapping[str, object]] = None, **kwargs: object) -> None:
        """_summary_
        When using logger.exception outside of a try/except block, it will raise an AttributeError that __name__ is missing. If there is sensitive information in the exception message, it will not be masked in the console. This override will mask the exception message before logging it.

        A common use case for this is running `logger.exception(event)` in a lambda handler, where the event may contain sensitive information, like an authorization header.
        """
        return super().exception(self.logger_formatter.format_masks(str(msg)), *args, exc_info=exc_info, stack_info=stack_info, stacklevel=stacklevel, extra=extra, **kwargs)

    @staticmethod
    def mask(s: str):
        with MASKS_LOCK:
            if s and s not in _masks:
                _masks.append(s)


class LogContext():
    _logger: AwsLogger
    _context_keys: Mapping

    def __init__(self, logger: AwsLogger, **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 _get_organization_id(event: dict) -> str:
    def check_variants(item: dict) -> str:
        return item.get("cz_organization_id", item.get("organization_id"))
    if organization_id := check_variants(event):
        return organization_id
    if organization_id := check_variants(event.get("requestContext", {}).get("authorizer", {})):
        return organization_id
    if organization_id := check_variants(event.get("detail", {})):
        return organization_id
    return ""


def _get_user_id(event: dict) -> str:
    if user_id := event.get("user_id"):
        return user_id
    if user_id := event.get("requestContext", {}).get("authorizer", {}).get("user_id"):
        return user_id
    if user_id := event.get("detail", {}).get("user_id"):
        return user_id
    return ""


def _get_metadata(event: dict) -> dict:
    fields = {}
    organization_id = _get_organization_id(event)
    user_id = _get_user_id(event)
    if organization_id:
        fields["organization_id"] = organization_id
    if user_id:
        fields["user_id"] = user_id
    return fields


def _get_auth_token(header: str):
    return header.replace("Bearer", "").replace("bearer", "").replace("Basic", "").replace("basic", "").strip()


def _mask_authorization_header(logger: Logger, event: dict):
    if headers := event.get("headers"):
        logger.mask(_get_auth_token(headers.get("Authorization", "")))
        logger.mask(_get_auth_token(headers.get("authorization", "")))
    if headers := event.get("multiValueHeaders"):
        if len(headers.get("Authorization", [])):
            logger.mask(_get_auth_token(headers.get("Authorization")[0]))
        if len(headers.get("authorization", [])):
            logger.mask(_get_auth_token(headers.get("authorization")[0]))


def _mask_request_context(logger: Logger, event: dict):
    logger.mask(event.get("requestContext", {}).get("identity", {}).get("accessKey", ""))


def _mask_user_data(logger: Logger, event: dict):
    if authorizer := event.get("requestContext", {}).get("authorizer"):
        logger.mask(authorizer.get("email", ""))
        logger.mask(authorizer.get("first_name", ""))
        logger.mask(authorizer.get("last_name", ""))


def _mask_all(logger: Logger, event: dict):
    _mask_authorization_header(logger, event)
    _mask_request_context(logger, event)
    _mask_user_data(logger, event)


def get_logger(*args, **kwargs) -> Logger:
    return Logger(*args, **kwargs)


def get_tracer(*args, **kwargs) -> Tracer:
    return Tracer(*args, **kwargs)


def log_context_from_event(logger: Logger):
    def log_context_from_event_handler(handler):
        def function_wrapper(event, *args, **kwargs):
            additional_fields = {}
            if metadata := _get_metadata(event):
                additional_fields["metadata"] = metadata
            _mask_all(logger, event)
            with LogContext(logger, **additional_fields):
                return handler(event, *args, **kwargs)
        return function_wrapper
    return log_context_from_event_handler
