# Copyright (c) 2020 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 os

import pytest
from pydantic import BaseModel, Extra

from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czc.logging import get_logger
from czt.unittest import Context

from tests.test_logging import get_modified_log_data


@pytest.fixture(scope='function')
def context(mocker, environ):
    from czc.state_handler import state_handler

    with Context(mocker, state_handler, modules_to_mock=[], extra={'state_handler': state_handler}) as context:
        yield context


class TestEvent(BaseModel, extra=Extra.allow):
    field1: str
    field2: int


@pytest.mark.unit
@pytest.mark.parametrize('with_test_key', [True, False])
def test_state_handler(context, with_test_key):
    test_key_value = 'test-key-value'

    @context.state_handler(TestEvent)
    def handler(event, _, parsed_event: TestEvent):
        if with_test_key:
            assert os.environ[CLOUDZERO_TEST_KEY] == test_key_value
        else:
            assert CLOUDZERO_TEST_KEY not in os.environ

        return parsed_event

    event = {'field1': 'value1', 'field2': 2, **({'__test_key__': test_key_value} if with_test_key else {})}

    result = handler(event, None)

    assert CLOUDZERO_TEST_KEY not in os.environ
    assert result.model_dump() == event


@pytest.mark.unit
@pytest.mark.parametrize('with_jmespath_override', [True, False])
@pytest.mark.parametrize('execution_arn_present', [True, False])
def test_state_handler_logging_context(context, with_jmespath_override, execution_arn_present):
    execution_arn_value = 'aws:execution_arn'
    execution_arn_override_value = 'aws:execution_arn_override'
    logger = get_logger()

    event_jmespaths = (
        {'execution_arn': 'some_field.execution_arn', 'some_other_field': 'field1'} if with_jmespath_override else None
    )

    @context.state_handler(TestEvent, event_jmespaths=event_jmespaths)
    def handler(event, _, parsed_event: TestEvent):
        return get_modified_log_data(logger.logger_formatter.log_format)

    event = {
        'field1': 'value1',
        'field2': 2,
        **(
            {'execution_arn': execution_arn_value, 'some_field': {'execution_arn': execution_arn_override_value}}
            if execution_arn_present
            else {}
        ),
    }

    result = handler(event, None)

    expected_execution_arn = execution_arn_override_value if with_jmespath_override else execution_arn_value

    if execution_arn_present:
        assert result['metadata']['execution_arn'] == expected_execution_arn
    else:
        assert result['metadata']['execution_arn'] is None

    if with_jmespath_override:
        assert result['metadata']['some_other_field'] == 'value1'
    else:
        assert 'some_other_field' not in result['metadata']
