from collections.abc import Callable, Hashable
from functools import partial, wraps

__all__ = ["multimethod"]

"""
Provides multiple dispatch or multimethods as described in
https://en.wikipedia.org/wiki/Multiple_dispatch

multimethods can be used to effectively solve the Expression Problem
described in https://en.wikipedia.org/wiki/Expression_problem

It can find uses in a lot of cases. Some examples are given below, but the uses are not limited
to only those

* Functions where we process inputs differently based on isinstance checks of one or more input arguments
* The basic case of processing differently based on the type of one instance is what is used in
Object oriented programming based on the self parameter. This basic methodology is what is called
Single Dispatch
* Functions where we process inputs differently based on the type of a field.
* Example: Different handler for different versions of payload based on a "version" field in a dict, etc

More important than those is the ability that multimethods provide of NOT having to change the dispatcher
for every single type of 'new' type that we come up with.

Usage is best described with an example:

Let's say we handle a payload based on a version number and type in the dict.

payload = {
    "type": "init"
            "version": 3,
                       ...
}

In the payload above we want to dispatch to a handler based on "type" and "version"

Typical two ways to do this are:

1. def handle(payload):
    type, version = payload["type"], payload["version"]
    if ("init", 3) == (type, version):
        return handle_init_v3(payload)
    elif ...

def handle_init_v3(payload):
    ...

2. Better that above
HANDLERS = {
    ("init", 3): handle_init_v3
        ...
}

def handle(payload):
    type, version = payload["type"], payload["version"]
    return HANDLERS[(type, version)](payload)

3. The above two suffer from the problem of having to change existing files when a new version or handler
is added. What if we could do this?

* In file named handle.py
# Register a handler function using a @multimethod decorator. Arg to the decorator
# is the dispatch key generation function. Body of the function decorated is
# the default handler if a handler wasn't registered for the dispatch key generated
# by the dispatch key generation function
@multimethod(lambda payload: (payload["type"], payload["version"]))
def handle_message(payload):
    # A optional default handler function to be called if a handler wasn't registered for the
    # dispatch key generated by the above dispatch key generation function
    pass

* In file 2
@handle_message.register(("init", 3))
def _(payload):
    # payload handler for init v3
    ...

* In file 2
@handle_message.register(("init", 4))
def _(payload):
    # payload handler for init v4
    ...

As can be seen from above, there is no need to change the dispatcher function or any other existing file

This implementation of multimethods is inspired by the Clojure language's implementation of the concept

Usage:
------
The same as mentioned in point 3 above.
The @multimethod(<dispatch_key_gen_fn>) decorator is used to define a multimethod dispatch function.
In the example above, the dispatcher is based on the payload type and payload version.
It should be noted that the assumptions is that whatever input is given
to this function is the same that is given to the dispatched functions.
The dispatched functions are defined on this multimethod using the <fn>.register(<key>) decorator, where <fn> is
the function decorated using @multimethod() decorator and <key> is the key for which we are registering <fn> as a
handler. The <key> is generated 
The dispatch_key is the result of calling <dispatch_key_gen_fn> of the @multimethod decorator
"""


def __fun_fqn(fun):
    module = getattr(fun, "__module__", "")
    qualname = getattr(fun, "__qualname__", "")
    return ".".join([val for val in [module, qualname] if val])


def __register(key, registry=None, registering_for_name="", overwrite=False):
    """
    Internal decorator to register the non-default handlers for multimethod
    registry and key_fn_name are keyword arguments with defaults to make it easy to apply
    functools.partial on them
    """
    if registry is None:
        registry = {}

    def decorator(handler):
        if not isinstance(key, Hashable):
            raise TypeError(
                f"Cannot register handler for function {registering_for_name} with key {repr(key)} that is not Hashable"
            )
        if key in registry and not overwrite:
            raise KeyError(
                f"Duplicate registration for key {repr(key)} for function {registering_for_name}"
            )
        registry[key] = handler

        @wraps(handler)
        def wrapper(*args, **kwargs):
            handler(*args, **kwargs)

        return wrapper

    return decorator


def multimethod(key_fn):
    key_fn_name = __fun_fqn(key_fn)

    if not isinstance(key_fn, Callable):
        raise TypeError(f"key_fn {key_fn_name} of type {type(key_fn)} must be Callable")

    def decorator(default_handler):
        default_handler_name = __fun_fqn(default_handler)
        registry = {}

        def dispatch(key, *args, **kwargs):
            handler = registry.get(key, default_handler)
            return handler(*args, **kwargs)

        @wraps(default_handler)
        def wrapper(*args, **kwargs):
            key = key_fn(*args, **kwargs)
            if not isinstance(key, Hashable):
                raise TypeError(
                    f"Key generator function {key_fn_name} returned a key {repr(key)} that is not Hashable"
                )
            return dispatch(key, *args, **kwargs)

        wrapper.register = partial(
            __register, registry=registry, registering_for_name=default_handler_name
        )
        wrapper.registered_keys = lambda: registry.keys()
        return wrapper

    return decorator
