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

from czc.logging import Logger, get_logger, log_context_from_event


@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 {**logger.context}

    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}, 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',
            'trace_id': '00000000000000000000000000000000',
        }
    }
    assert inner_context == {
        'metadata': {
            'organization_id': '5678',
            'user_id': '5678',
            'trace_id': '00000000000000000000000000000000',
        }
    }
    assert logger.context == {}


@pytest.mark.unit
def test_log_context_authorizer_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_log_context_authorizer_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_log_context_root_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

    event = {'cz_organization_id': '1234'}

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


@pytest.mark.unit
def test_log_context_root_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

    event = {'organization_id': '1234'}

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


@pytest.mark.unit
def test_log_context_detail_cz_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_log_context_detail_organization_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_log_context_root_user_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

    event = {'user_id': '1234'}

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


@pytest.mark.unit
def test_log_context_detail_user_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


def test_log_context_empty_event(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

    event = {}

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


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

    @log_context_from_event(logger)
    def test_fn(event: dict):
        return logger.info(f'NOT_SECRET {event}', extra={**event, 'password': 'NOT_SECRET'})

    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'},
        },
    }
    log = test_fn(event)
    assert (
        log
        == "<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>'}}} {'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 {**logger.context}

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

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


@pytest.mark.unit
def test_correlation_id_from_request_headers(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_correlation_id_from_headers(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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


@pytest.mark.unit
def test_correlation_id_from_id(logger: Logger):
    @log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

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

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