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

from czc import json

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 SlimFormatter(Formatter):
    """Formatter that only returns the message field from the log record.

    This is useful in build systems, CLIs, and other systems that do not need fully JSON formatted logs.
    """

    def format(self, record):
        msg = super().format(record)
        return json.loads(msg).get('message', '')


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: Any) -> str:
        if isinstance(item, str):
            item_obj = json.loads(item)
            return item_obj.get('cz_organization_id', item_obj.get('organization_id', item_obj.get('organizationId')))
        return item.get('cz_organization_id', item.get('organization_id', item.get('organizationId')))

    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:
    def check_variants(item: Any) -> str:
        if isinstance(item, str):
            return json.loads(item).get('user_id')
        return item.get('user_id')

    if user_id := check_variants(event):
        return user_id
    if user_id := check_variants(event.get('requestContext', {}).get('authorizer', {})):
        return user_id
    if user_id := check_variants(event.get('detail', {})):
        return user_id
    return ''


def _get_environment_variables() -> dict[str, str]:
    handler = os.environ.get('AWS_LAMBDA_FUNCTION_NAME')
    feature = os.environ.get('FEATURE')
    namespace = os.environ.get('NAMESPACE')
    return {
        **({'handler': handler} if handler else {}),
        **({'feature': feature} if feature else {}),
        **({'namespace': namespace} if namespace else {}),
    }


def _get_metadata(
    event: dict, context: Any, handler_name: Optional[str], event_jmespaths: Optional[Mapping[str, str]] = None
) -> dict:
    organization_id = _get_organization_id(event)
    user_id = _get_user_id(event)
    trace_id = ''
    request_id = '' if not context else context.aws_request_id

    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 {}),
        **({'request_id': request_id} if request_id else {}),
        **_get_environment_variables(),
        **(
            {'handler': handler_name} if handler_name else {}
        ),  # This will override the handler name from the environment variables
    }


def _get_response_metadata(response: dict, response_jmespaths: Optional[Mapping[str, str]] = None) -> dict:
    status_code = response.get('statusCode')
    return {
        **({k: extract_data_from_envelope(response, v) for k, v in response_jmespaths.items()} if response_jmespaths else {}),
        **({'status_code': status_code} if status_code 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_all(logger: Logger, event: dict):
    _mask_authorization_header(logger, event)
    _mask_request_context(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,
    response_jmespaths: Optional[dict[str, str]] = None,
    handler_name: Optional[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):
        @functools.wraps(handler)
        def function_wrapper(event, *args, **kwargs):
            context = None
            if args and hasattr(args[0], 'aws_request_id'):
                context = args[0]
            additional_fields = {}
            try:
                if metadata := _get_metadata(event, context, handler_name, event_jmespaths):
                    additional_fields['metadata'] = metadata
                _mask_all(logger, event)
                _set_correlation_id(logger, event)
            except Exception as e:
                # Errors with the logger should not prevent the handler from running
                logger.exception(f'Unhandled exception from czc logger: {e}')
            with LogContext(logger, **additional_fields):
                logger.info('Entering handler')
                try:
                    res = handler(event, *args, **kwargs)
                except:
                    logger.exception('Unhandled exception encountered')
                    raise
                logger.info('Exiting handler')
                response_metadata = _get_response_metadata(dict(res), response_jmespaths) if isinstance(res, Mapping) else None
                if response_metadata:
                    logger.info('Response metadata', extra={'response_metadata': response_metadata})
                return res

        return function_wrapper

    return log_context_from_event_handler
