# 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 os
import threading
from collections.abc import Callable, Iterable, Mapping
from types import TracebackType
from typing import IO, Any, Optional, Union

from aws_lambda_powertools import Tracer
from aws_lambda_powertools.logging import Logger as AwsLogger
from aws_lambda_powertools.logging.formatter import LambdaPowertoolsFormatter
from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope

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


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:
        if 'metadata' in record.__dict__:
            record.__dict__['metadata'].update(self.get_current_keys().get('metadata', {}))
        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:
                _masks.add(s)

    def append_keys(self, **additional_keys: object) -> None:
        if 'metadata' in additional_keys:
            additional_keys['metadata'].update(self.get_current_keys().get('metadata', {}))  # type: ignore

        super().append_keys(**additional_keys)

    def remove_keys(self, keys: Iterable[str]) -> None:
        for key in keys:
            parts = key.split('.')
            current_data = self.logger_formatter.log_format
            # Navigate to the parent of the key to remove
            for part in parts[:-1]:
                current_data = current_data.get(part)  # type: ignore
                if current_data is None:
                    break  # Key path does not exist, exit early
            else:
                # Remove the key if it exists, do nothing if it doesn't
                current_data.pop(parts[-1], None)


class LogContext:
    _logger: Logger
    _context_keys: Mapping

    def __init__(self, logger: Logger, **additional_keys) -> None:
        self._logger = logger
        self._context_keys = additional_keys
        self._logger.get_current_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: Mapping) -> 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 _is_tracing_enabled() -> bool:
    # AWS_LAMBDA_EXEC_WRAPPER must match "opt/otel-instrument" to ensure the Lambda is instrumented by OpenTelemetry.
    # This makes it a safe environment variable to check for to determine if tracing should be enabled.
    return os.environ.get('AWS_LAMBDA_EXEC_WRAPPER') == '/opt/otel-instrument'


def _get_trace_id() -> str:
    """Pulls the trace ID from the current span context. This is necessary to handle X-Ray sampling.
    Per the docs (https://docs.aws.amazon.com/lambda/latest/dg/services-xray.html):

    "X-Ray doesn't trace all requests to your application. X-Ray applies a sampling algorithm to
    ensure that tracing is efficient, while still providing a representative sample of all requests.
    The sampling rate is 1 request per second and 5 percent of additional requests."

    In the cases where X-Ray does not sample the Lambda, OpenTemetry will still generate a trace ID.
    This handles getting the trace ID, regardless of whether the Lambda is sampled by X-Ray.

    Returns:
        str: Current trace ID
    """
    from opentelemetry import trace

    return format(trace.get_current_span().get_span_context().trace_id, '032x')


def _get_metadata(event: dict, event_jmespaths: Optional[Mapping[str, str]] = None) -> dict:
    organization_id = _get_organization_id(event)
    user_id = _get_user_id(event)
    trace_id = _get_trace_id() if _is_tracing_enabled() else ''

    return {
        **({k: extract_data_from_envelope(event, v) for k, v in event_jmespaths.items()} if event_jmespaths else {}),
        **({'organization_id': organization_id} if organization_id else {}),
        **({'user_id': user_id} if user_id else {}),
        **({'trace_id': trace_id} if trace_id else {}),
    }


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 _set_correlation_id(logger: Logger, event: dict):
    """Set correlation id from event. Values taken from
    https://docs.powertools.aws.dev/lambda/python/latest/core/logger/#built-in-correlation-id-expressions

    Args:
        logger (Logger): Logger instance
        event (dict): Lambda event
    """
    if correlation_id := event.get('requestContext', {}).get('requestId'):
        logger.set_correlation_id(correlation_id)
    elif correlation_id := event.get('request', {}).get('headers', {}).get('x-amzn-trace-id'):
        logger.set_correlation_id(correlation_id)
    elif correlation_id := event.get('headers', {}).get('x-amzn-trace-id'):
        logger.set_correlation_id(correlation_id)
    elif correlation_id := event.get('id'):
        logger.set_correlation_id(correlation_id)


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, event_jmespaths: Optional[dict[str, str]] = None):
    """Decorator to extract event and context fields and add them to the log metadata.
    Sensitive values in logs are masked to prevent them from being leaked.

    Args:
        logger (Logger): Logger instance
        event_jmespaths (list[str]): JMESPath expressions to extract from the event and add to the log metadata
    """

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

        return function_wrapper

    return log_context_from_event_handler
