# 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 json
import logging
from copy import deepcopy

import pytest

from czc.logging import Logger, get_logger, log_context_from_event


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
def test_log_context_authorizer_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'requestContext': {'authorizer': {'cz_organization_id': '1234'}}}

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


@pytest.mark.unit
def test_log_context_authorizer_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'requestContext': {'authorizer': {'organization_id': '1234'}}}

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


@pytest.mark.unit
def test_log_context_root_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'cz_organization_id': '1234'}

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


@pytest.mark.unit
def test_log_context_root_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'organization_id': '1234'}

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


@pytest.mark.unit
def test_log_context_detail_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

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

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


@pytest.mark.unit
def test_log_context_detail_organization_id(logger: Logger):
    @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'}}


@pytest.mark.unit
def test_log_context_root_user_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {'user_id': '1234'}

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


@pytest.mark.unit
def test_log_context_detail_user_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

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

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


@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 = {}

    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': 'first_name', 'last_name': 'lname'},
            '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': '<HIDDEN>', '<HIDDEN>': '<HIDDEN>', 'last_name': '<HIDDEN>'}, "
        "'identity': {'accessKey': '<HIDDEN>'}}}",
        'password': '<HIDDEN>',
    }


@pytest.mark.unit
def test_logger_exception_error(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_correlation_id_from_request_context_request_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'requestContext': {
            'requestId': '1234',
        }
    }

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


@pytest.mark.unit
def test_correlation_id_from_request_headers(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'request': {
            'headers': {
                'x-amzn-trace-id': '1234',
            }
        }
    }

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


@pytest.mark.unit
def test_correlation_id_from_headers(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'headers': {
            'x-amzn-trace-id': '1234',
        }
    }

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


@pytest.mark.unit
def test_correlation_id_from_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'id': '1234',
    }

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


@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'})
    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

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

    res_1, res_2 = test_fn(event)

    assert res_1 == {
        'correlation_id': '1234',
        'metadata': {
            'test_field': '1234',
            'organization_id': '1234',
            'test': 'test',
        },
        'test_1': True,
    }
    assert res_2 == {
        'correlation_id': '1234',
        'metadata': {'organization_id': '1234', 'test': 'test'},
    }
