# 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
from enum import Enum
from typing import NamedTuple
from unittest.mock import call

import pytest

from czc import xray
from czc.type_util import deserializable
from czc.xray import capture_function
from czt.unittest import Context


@pytest.fixture(scope='function')
def context(mocker):
    with Context(mocker, xray, ['tracer']) as context:

        def passthrough(func):
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            return wrapper

        context.mock_tracer.capture_method.side_effect = lambda *a, **ka: passthrough
        yield context


class ObjSize(Enum):
    big = 'big'
    small = 'small'


@deserializable()
class SomeObj(NamedTuple):
    size: ObjSize = ObjSize.big
    name: str = 'NAME'
    num: int = 123


TEST_ARG1 = 999
TEST_ARG2 = 'SOMETHING'


@pytest.mark.unit
@pytest.mark.parametrize(
    'use_function_annotation,annotations,arg_annotations,expected_annotations',
    [
        (
            True,
            {'foo': 'bar'},
            {'arg1', 'arg2'},
            [('FUNCTION_NAME', 'stub_func1'), ('arg1', TEST_ARG1), ('arg2', TEST_ARG2), ('foo', 'bar')],
        ),
        (False, {'foo': 'bar'}, {'arg1', 'arg2'}, [('arg1', TEST_ARG1), ('arg2', TEST_ARG2), ('foo', 'bar')]),
        (False, {}, {'arg1'}, [('arg1', TEST_ARG1)]),
        (False, None, None, []),
    ],
)
def test_capture_function_should_add_expected_annotations(
    context, use_function_annotation, annotations, arg_annotations, expected_annotations
):
    @capture_function(
        use_function_name_annotation=use_function_annotation, arg_annotations=arg_annotations, annotations=annotations
    )
    def stub_func1(arg1: int, arg2: str, arg3: SomeObj):
        pass

    stub_func1(TEST_ARG1, TEST_ARG2, SomeObj())

    assert context.mock_tracer.put_annotation.call_args_list == [call(k, v) for k, v in expected_annotations]


@pytest.mark.unit
@pytest.mark.parametrize(
    'skip_args, skip_arg_types, expected_metadata',
    [
        (
            set(),
            set(),
            {'arg1': 999, 'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}},
        ),
        ({'arg1'}, set(), {'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}}),
        (set(), {SomeObj}, {'arg1': 999, 'arg2': 'SOMETHING'}),
    ],
)
def test_capture_function_should_add_metadata(context, skip_args, skip_arg_types, expected_metadata):
    @capture_function(skip_args=skip_args, skip_arg_types=skip_arg_types)
    def stub_func1(arg1: int, arg2: str, arg3: SomeObj):
        pass

    stub_func1(TEST_ARG1, TEST_ARG2, SomeObj())

    context.mock_tracer.put_metadata.assert_called_with('arguments', expected_metadata)


@pytest.mark.unit
@pytest.mark.parametrize(
    'skip_args, skip_arg_types, expected_metadata',
    [
        (
            set(),
            set(),
            {'arg1': 999, 'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}},
        ),
        ({'arg1'}, set(), {'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}}),
        (set(), {SomeObj}, {'arg1': 999, 'arg2': 'SOMETHING'}),
    ],
)
def test_capture_function_should_add_metadata_for_named_args(context, skip_args, skip_arg_types, expected_metadata):
    @capture_function(skip_args=skip_args, skip_arg_types=skip_arg_types)
    def stub_func1(arg1: int, arg2: str, arg3: SomeObj):
        pass

    stub_func1(arg1=TEST_ARG1, arg2=TEST_ARG2, arg3=SomeObj())

    context.mock_tracer.put_metadata.assert_called_with('arguments', expected_metadata)


@pytest.mark.unit
@pytest.mark.parametrize(
    'skip_args, skip_arg_types, expected_metadata',
    [
        (
            set(),
            set(),
            {'arg1': 999, 'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}},
        ),
        ({'arg1'}, set(), {'arg2': 'SOMETHING', 'arg3': {'size': 'big', 'name': 'NAME', 'num': 123}}),
        (set(), {SomeObj}, {'arg1': 999, 'arg2': 'SOMETHING'}),
    ],
)
def test_capture_function_should_add_metadata_for_default_args(context, skip_args, skip_arg_types, expected_metadata):
    @capture_function(skip_args=skip_args, skip_arg_types=skip_arg_types)
    def stub_func1(arg1: int = TEST_ARG1, arg2: str = TEST_ARG2, arg3: SomeObj = SomeObj()):
        pass

    stub_func1()

    context.mock_tracer.put_metadata.assert_called_with('arguments', expected_metadata)


class BadClass:
    value: str

    def __init__(self, v):
        self.value = v

    def serializable(self):
        raise ValueError('CANNOT SERIALIZE ME')

    def __repr__(self):
        return f'BadClass: {self.value}'


@pytest.mark.unit
@pytest.mark.parametrize(
    'bad_value,expected_error',
    [
        ({('no', 'tuple', 'keys'): 'bad'}, 'TypeError: keys must be str, int, float, bool or None, not tuple'),
        ({ObjSize.big: 'no enum keys either'}, 'TypeError: keys must be str, int, float, bool or None, not ObjSize'),
        (BadClass('should fail'), 'ValueError: CANNOT SERIALIZE ME'),
    ],
)
def test_capture_function_should_handle_bad_serialization(context, bad_value, expected_error):
    @capture_function()
    def stub_func_with_test_arg(bad_arg):
        pass

    stub_func_with_test_arg(bad_value)

    context.mock_tracer.put_metadata.assert_any_call('arguments', f"{{'bad_arg': {str(bad_value)}}}")

    error_calls = [
        args for args, _ in context.mock_tracer.put_metadata.call_args_list if args[0] == 'metadata_collection_error'
    ]
    assert len(error_calls) == 1, f'Expected exactly one metadata_collection_error call, got {len(error_calls)}'
    actual_error = error_calls[0][1]
    assert actual_error.startswith(f'{expected_error}\n'), (
        f'Expected error starting with {expected_error!r}, got {actual_error!r}'
    )
