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

import pytest

from czc.enum_dispatch import enum_dispatch
from czc.exceptions import InvalidValueError, ValuesHasMultipleDispatchesError


class Operation(str, Enum):
    add = 'add'
    subtract = 'subtract'
    multiply = 'multiply'


@enum_dispatch(Operation)
def perform_operation(enum_value: Operation, value_a, value_b):
    raise NotImplementedError()


@perform_operation.register(Operation.add)
def _add(value_a, value_b):
    return value_a + value_b


@perform_operation.register(Operation.subtract)
def _subtract(value_a, value_b):
    return value_a - value_b


@perform_operation.register(Operation.multiply)
def _multiply(value_a, value_b):
    return value_a * value_b


@pytest.mark.unit
@pytest.mark.parametrize(
    'operation, a, b, expected_result',
    [
        (Operation.add, 5, 3, 8),
        (Operation.subtract, 5.3, 3.4, 1.9),
        (Operation.multiply, 5, 3, 15),
    ],
)
def test_enum_dispatch(operation: Operation, a, b, expected_result):
    assert expected_result == perform_operation(operation, a, b)


@pytest.mark.unit
def test_enum_dispatch_catches_overlaps():
    with pytest.raises(ValuesHasMultipleDispatchesError):

        @perform_operation.register(Operation.add)
        def _another_add(a, b):
            return b + a


class AnotherEnum(str, Enum):
    a = 'a'


@pytest.mark.unit
def test_enum_dispatch_catches_invalid_enum():
    with pytest.raises(InvalidValueError):

        @perform_operation.register(AnotherEnum.a)
        def _another_add(a, b):
            return b + a


@pytest.mark.unit
def test_enum_dispatch_catches_dispatch_of_invalid_enum():
    with pytest.raises(InvalidValueError):
        perform_operation(AnotherEnum.a, 5, 3)  # ignore: type
