# 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 Dict, NamedTuple

import pytest
from unittest.mock import call

from czc import xray
from czt.unittest import Context, MockEx
from czc.xray import capture_function


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


class SomeObj(NamedTuple):
    size: ObjSize = ObjSize.big
    name: str = 'NAME'
    num: int = 123
    cost: Dict[ObjSize, float] = {
        ObjSize.big: 1000000,
        ObjSize.small: 1
    }


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': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    ({'arg1'}, set(), {
        'arg2': 'SOMETHING',
        'arg3': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    (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())

    assert MockEx(context.mock_tracer.put_metadata).was_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': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    ({'arg1'}, set(), {
        'arg2': 'SOMETHING',
        'arg3': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    (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())

    assert MockEx(context.mock_tracer.put_metadata).was_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': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    ({'arg1'}, set(), {
        'arg2': 'SOMETHING',
        'arg3': ['big', 'NAME', 123, {'ObjSize.big': 1000000, 'ObjSize.small': 1}]}),
    (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()

    assert MockEx(context.mock_tracer.put_metadata).was_called_with('arguments', expected_metadata)
