# 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 collections.abc import Callable
from typing import Generic, TypeVar

from czc.exceptions import InvalidValueError, ValuesHasMultipleDispatchesError

T = TypeVar('T')


class value_dispatch(Generic[T]):
    """
    This decorator can be used to create a single dispatch function that will dispatch
    the call to the appropriately designated function. For example:

    @value_dispatch[str]
    def perform_operation(operation: str, value_a, value_b):
        return f'{operation}: {value_a} {value_b}'


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


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


    >> print(perform_operation('subtract', 5, 3))
    2
    >> print(perform_operation('multiply', 5, 3))
    15
    >> print(perform_operation('something_else', 5, 3))
    something_else: 5 3

    """

    _registry: dict[T, Callable]
    _base_func: Callable

    @property
    def value_type(self):
        return self.__orig_class__.__args__[0]  # type: ignore[attr-defined]

    def __init__(self, base_func: Callable) -> None:
        super().__init__()

        self._base_func = base_func
        self._registry = {}

    def __call__(self, value: T, *args, **kwargs):
        if not args:
            return TypeError('At least one argument required')

        if not isinstance(value, self.value_type):
            raise InvalidValueError(f'Cannot dispatch. The argument {value} must be a value of {self.value_type}')

        if value not in self._registry:
            return self._base_func(value, *args, **kwargs)

        return self._registry[value](*args, **kwargs)

    def register(self, value: T):
        if not isinstance(value, self.value_type):
            raise InvalidValueError(f'Value {value} must be a value of {self.value_type}')

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

        return handler_func
