# 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
from collections.abc import Mapping
from copy import deepcopy
from typing import Any, Optional

import pytest
from pytest_mock import MockerFixture

from czc import json
from czc.logging import Logger, SlimFormatter, get_logger, log_context_from_event

from tests.unittest_helpers import LambdaContext


def get_modified_log_data(data: dict) -> dict:
    data_copy = deepcopy(data)
    for field in [
        'level',
        'location',
        'sampling_rate',
        'timestamp',
        'service',
        'cold_start',
        'function_arn',
        'function_memory_size',
        'function_name',
        'function_request_id',
        'cz_organization_id',
    ]:
        data_copy.pop(field, None)
    if not data_copy.get('message'):
        data_copy.pop('message', None)
    return data_copy


@pytest.mark.unit
def test_log_context_nested(logger: Logger):
    inner_event = {'organization_id': '5678', 'user_id': '5678'}

    @log_context_from_event(logger)
    def test_inner_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format), test_inner_fn(inner_event)

    event = {'organization_id': '1234', 'user_id': '1234'}

    context, inner_context = test_fn(event)

    assert context == {
        'metadata': {
            'organization_id': '1234',
            'user_id': '1234',
        }
    }
    assert inner_context == {
        'metadata': {
            'organization_id': '1234',
            'user_id': '1234',
        }
    }


@pytest.mark.unit
@pytest.mark.parametrize(
    'event',
    [
        {'requestContext': {'authorizer': {'cz_organization_id': '1234'}}},
        {'requestContext': {'authorizer': {'organization_id': '1234'}}},
        {'requestContext': {'authorizer': {'organizationId': '1234'}}},
        {'cz_organization_id': '1234'},
        {'organization_id': '1234'},
        {'organizationId': '1234'},
        {'detail': {'cz_organization_id': '1234'}},
        {'detail': {'organization_id': '1234'}},
        {'detail': {'organizationId': '1234'}},
        {'detail': json.dumps({'cz_organization_id': '1234'})},
        {'detail': json.dumps({'organization_id': '1234'})},
        {'detail': json.dumps({'organizationId': '1234'})},
    ],
)
@pytest.mark.parametrize(
    'context,is_valid_context',
    [
        (LambdaContext('111-222-333', None, {}), True),
        ({'aws_request_id': '111-222-333'}, False),
        (None, False),
        ('', False),
        ({}, False),
        (1234, False),
        (True, False),
    ],
)
def test_organization_id(logger: Logger, event, context, is_valid_context):
    @log_context_from_event(logger)
    def test_fn(event, *args, **kwargs):
        return get_modified_log_data(logger.logger_formatter.log_format)

    assert test_fn(event) == {
        'metadata': {
            'organization_id': '1234',
        }
    }

    if is_valid_context:
        assert test_fn(event, context) == {
            'metadata': {
                'organization_id': '1234',
                'request_id': '111-222-333',
            }
        }
    else:
        assert test_fn(event, context) == {
            'metadata': {
                'organization_id': '1234',
            }
        }


@pytest.mark.unit
@pytest.mark.parametrize(
    'event',
    [
        {'user_id': '1234'},
        {'detail': {'user_id': '1234'}},
        {'detail': json.dumps({'user_id': '1234'})},
    ],
)
def test_user_id(logger: Logger, event):
    @log_context_from_event(logger)
    def test_fn(event, *args, **kwargs):
        return get_modified_log_data(logger.logger_formatter.log_format)

    assert test_fn(event) == {
        'metadata': {
            'user_id': '1234',
        }
    }


@pytest.mark.unit
@pytest.mark.parametrize(
    'event',
    [
        {'requestContext': {'requestId': '1234'}},
        {'request': {'headers': {'x-amzn-trace-id': '1234'}}},
        {'headers': {'x-amzn-trace-id': '1234'}},
        {'id': '1234'},
    ],
)
def test_correlation_id(logger: Logger, event):
    @log_context_from_event(logger)
    def test_fn(event, *args, **kwargs):
        return get_modified_log_data(logger.logger_formatter.log_format)

    assert test_fn(event) == {
        'correlation_id': '1234',
    }


@pytest.mark.unit
def test_handler_name(logger: Logger):
    @log_context_from_event(logger, handler_name='Delete Connection')
    def test_fn(event, *args, **kwargs):
        return get_modified_log_data(logger.logger_formatter.log_format)

    assert test_fn({}) == {
        'metadata': {
            'handler': 'Delete Connection',
        },
    }


@pytest.mark.unit
def test_log_context_unhandled_error(logger: Logger, mocker: MockerFixture):
    mocker.patch('czc.logging._set_correlation_id', side_effect=Exception('Test Exception'))

    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    test_fn({})


@pytest.mark.unit
def test_trace_id(logger: Logger, monkeypatch):
    monkeypatch.setenv('AWS_LAMBDA_EXEC_WRAPPER', '/opt/otel-instrument')

    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'detail': {'organization_id': '1234'}}

    assert test_fn(event) == {'metadata': {'organization_id': '1234', 'trace_id': '00000000000000000000000000000000'}}


def test_log_context_empty_event(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event: dict[Any, Any] = {}

    assert test_fn(event) == {}


@pytest.mark.unit
def test_log_secrets(logger: Logger):
    logger.mask('NOT_SECRET')

    @log_context_from_event(logger)
    def test_fn(event: dict):
        record = logging.LogRecord(
            name='test',
            level=logging.INFO,
            pathname='test',
            lineno=1,
            msg=f'NOT_SECRET {event}',
            args=(),
            exc_info=None,
            func='test',
        )
        extra = {'password': 'NOT_SECRET'}
        logger.append_keys(**extra)
        masked_str = logger.logger_formatter.format(record)
        logger.append_keys(message=masked_str)
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'headers': {'Authorization': 'Basic pass1', 'authorization': 'pass2'},
        'multiValueHeaders': {'Authorization': ['Bearer pass3'], 'authorization': ['pass4']},
        'requestContext': {
            'authorizer': {'email': 'example@example.com', 'first_name': 'test_first_name', 'last_name': 'test_last_name'},
            'identity': {'accessKey': 'passkey'},
        },
    }
    res = test_fn(event)
    assert {k: v for k, v in json.loads(res['message']).items() if k not in ['timestamp', 'location']} == {
        'level': 'INFO',
        'message': "<HIDDEN> {'headers': {'Authorization': 'Basic <HIDDEN>', 'authorization': "
        "'<HIDDEN>'}, 'multiValueHeaders': {'Authorization': ['Bearer <HIDDEN>'], "
        "'authorization': ['<HIDDEN>']}, 'requestContext': {'authorizer': "
        "{'email': 'example@example.com', 'first_name': 'test_first_name', 'last_name': 'test_last_name'}, "
        "'identity': {'accessKey': '<HIDDEN>'}}}",
        'password': '<HIDDEN>',
    }


@pytest.mark.unit
def test_logger_functions(capsys):
    logger = get_logger()
    secret = 'A' * 5
    logger.mask(secret)
    logger.exception(f'This is an exception: {secret}')
    cap = capsys.readouterr()
    assert '<HIDDEN>' in str(cap)
    assert secret not in cap.err


@pytest.mark.unit
def test_start_and_success_logs(caplog):
    logger = get_logger()

    @log_context_from_event(logger)
    def test_fn(event):
        logger.info('Doing a thing in the handler')
        return event

    test_fn({})
    assert caplog.record_tuples == [
        ('service_undefined', logging.INFO, 'Entering handler'),
        ('service_undefined', logging.INFO, 'Doing a thing in the handler'),
        ('service_undefined', logging.INFO, 'Handler completed successfully'),
    ]


@pytest.mark.unit
def test_start_and_exception_logs(caplog):
    logger = get_logger()

    @log_context_from_event(logger)
    def test_fn(event):
        raise Exception('uh oh')

    try:
        test_fn({})
    except Exception:
        pass
    assert caplog.record_tuples == [
        ('service_undefined', logging.INFO, 'Entering handler'),
        ('service_undefined', logging.ERROR, 'Unhandled exception encountered'),
    ]


@pytest.mark.unit
def test_add_remove():
    logger = get_logger()

    @log_context_from_event(
        logger,
        event_jmespaths={
            'test': 'headers.test',
            'organization_id': 'this_should_not_overwrite',
            'not_found_field': 'some.where.not_found',
        },
    )
    def test_fn(event):
        logger.append_keys(test_1=True, metadata={'test_field': '1234', 'organization_id': 'this_should_not_overwrite'})
        res_1 = get_modified_log_data(logger.logger_formatter.log_format)
        logger.remove_keys(['test_1', 'metadata.test_field'])
        res_2 = get_modified_log_data(logger.logger_formatter.log_format)
        return res_1, res_2

    def sanitize_result(res: dict):
        if res.get('message'):
            del res['message']
        if res.get('password'):
            del res['password']
        return res

    event = {
        'id': '1234',
        'organization_id': '1234',
        'headers': {
            'test': 'test',
        },
    }

    res_1, res_2 = test_fn(event)
    assert sanitize_result(res_1) == {
        'correlation_id': '1234',
        'metadata': {
            'test_field': '1234',
            'organization_id': '1234',
            'test': 'test',
            'not_found_field': None,
        },
        'test_1': True,
    }
    assert sanitize_result(res_2) == {
        'correlation_id': '1234',
        'metadata': {'organization_id': '1234', 'test': 'test', 'not_found_field': None},
    }


@pytest.mark.unit
@pytest.mark.parametrize(
    'response,expected',
    [
        ({'statusCode': 200}, {'response_metadata': {'status_code': 200}}),
        (
            {'statusCode': 403, 'headers': {'x-amzn-trace-id': '1234'}, 'body': {'test': 'test'}},
            {'response_metadata': {'status_code': 403}},
        ),
        ({'test': 'test'}, {}),
        ({}, {}),
        (None, {}),
        ('test', {}),
        (True, {}),
        (1234, {}),
    ],
)
def test_response_metadata_no_jmespaths(mocker: MockerFixture, logger: Logger, response, expected):
    def mock_logger_info(
        msg: object,
        *args: object,
        exc_info: Any = None,
        stack_info: bool = False,
        stacklevel: int = 2,
        extra: Optional[Mapping[str, object]] = None,
        **kwargs: object,
    ):
        if not extra:
            extra = {}
        logger.append_keys(**extra)

    mocker.patch.object(logger, 'info', side_effect=mock_logger_info)

    @log_context_from_event(logger)
    def test_fn(event, *args, **kwargs):
        return response

    test_fn({})
    assert get_modified_log_data(logger.logger_formatter.log_format) == expected


@pytest.mark.unit
@pytest.mark.parametrize(
    'response,expected',
    [
        ({'statusCode': 200}, {'response_metadata': {'status_code': 200, 'x-amzn-trace-id': None}}),
        (
            {'statusCode': 403, 'headers': {'x-amzn-trace-id': '1234'}, 'body': {'test': 'test'}},
            {'response_metadata': {'status_code': 403, 'x-amzn-trace-id': '1234'}},
        ),
        ({'test': 'test'}, {'response_metadata': {'x-amzn-trace-id': None}}),
        ({}, {'response_metadata': {'x-amzn-trace-id': None}}),
        (None, {}),
        ('test', {}),
        (True, {}),
        (1234, {}),
    ],
)
def test_response_metadata_jmespaths(mocker: MockerFixture, logger: Logger, response, expected):
    def mock_logger_info(
        msg: object,
        *args: object,
        exc_info: Any = None,
        stack_info: bool = False,
        stacklevel: int = 2,
        extra: Optional[Mapping[str, object]] = None,
        **kwargs: object,
    ):
        if not extra:
            extra = {}
        logger.append_keys(**extra)

    mocker.patch.object(logger, 'info', side_effect=mock_logger_info)

    @log_context_from_event(logger, response_jmespaths={'x-amzn-trace-id': 'headers."x-amzn-trace-id"'})
    def test_fn(event, *args, **kwargs):
        return response

    test_fn({})
    assert get_modified_log_data(logger.logger_formatter.log_format) == expected


@pytest.mark.unit
def test_slim_formatter():
    logger = get_logger(logger_formatter=SlimFormatter())

    # Test basic message
    logger.info('simple message')
    assert (
        logger.logger_formatter.format(logger._logger.makeRecord('test', logging.INFO, 'test', 1, 'simple message', (), None))
        == 'simple message'
    )

    # Test with extra fields
    logger.info('message with extra', extra={'field': 'value'})
    assert (
        logger.logger_formatter.format(
            logger._logger.makeRecord('test', logging.INFO, 'test', 1, 'message with extra', (), None, extra={'field': 'value'})
        )
        == 'message with extra'
    )

    # Test with structured logging
    logger.structure_logs(append=True)
    logger.info('structured message', extra={'metadata': {'test': 'value'}})
    assert (
        logger.logger_formatter.format(
            logger._logger.makeRecord(
                'test', logging.INFO, 'test', 1, 'structured message', (), None, extra={'metadata': {'test': 'value'}}
            )
        )
        == 'structured message'
    )
