# 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
from typing import Type
from functools import update_wrapper

from czc.exceptions import ValuesHasMultipleDispatchesError, InvalidValueError


def enum_dispatch(enum_type: Type[Enum]):
    '''
    This decorator can be used to create a single dispatch function that will dispatch
    the call to the appropriately designated function. For example:

    class Operation(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.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


    >> print(perform_operation(Operation.add, 5, 3))
    8
    >> print(perform_operation(Operation.subtract, 5, 3))
    2
    >> print(perform_operation(Operation.multiply, 5, 3))
    15

    '''
    def base_func_handler(base_func):
        registry = {}

        def register(value: Enum):
            if not isinstance(value, enum_type):
                raise InvalidValueError(f'Value {value} must be a value of {enum_type}')

            def handler_func(func):
                if value in registry:
                    raise ValuesHasMultipleDispatchesError(f'{value} cannot have multiple dispatches')
                registry[value] = func
                return func
            return handler_func

        def wrapper(*args, **kwargs):
            if not args:
                return TypeError('At least one argument required')

            if not isinstance(args[0], enum_type):
                raise InvalidValueError(f'Cannot dispatch. The argument {args[0]} must be a value of {enum_type}')

            if args[0] not in registry:
                return NotImplementedError(f'Dispatch for {args[0]} not implemented')

            return registry[args[0]](*args[1:], **kwargs)

        wrapper.register = register
        update_wrapper(wrapper, base_func)
        return wrapper
    return base_func_handler
