"""
Two magic tricks for classes:

    class X(metaclass=extendabletype):
        ...

    # in some other file...
    class __extend__(X):
        ...      # add new methods and class attributes to X

Mostly useful together with the second trick, which lets you build
methods whose 'self' is a pair of objects instead of just one:

    class __extend__(pairtype(X, Y)):
        attribute = 42
        def method(self, other, arguments):
            x, y = self
            ...

    pair(x, y).attribute
    pair(x, y).method(other, arguments)

This finds methods and class attributes based on the actual
class of both objects that go into the pair(), with the usual
rules of method/attribute overriding in (pairs of) subclasses.

For more information, see test_pairtype.
"""

from typing import Any, Dict, Tuple, Type, TypeVar, Generic, Callable, Optional, Iterable

T1 = TypeVar('T1')
T2 = TypeVar('T2')

class extendabletype(type):
    """A type with a syntax trick: 'class __extend__(t)' actually extends
    the definition of 't' instead of creating a new subclass."""
    def __new__(mcs, name: str, bases: Tuple[type, ...], dct: Dict[str, Any]) -> Optional[type]:
        if name == '__extend__':
            for cls in bases:
                for key, value in dct.items():
                    if key == '__module__':
                        continue
                    # Add attributes to the base class
                    setattr(cls, key, value)
            return None
        else:
            return super().__new__(mcs, name, bases, dct)

def pair(a: T1, b: T2) -> Tuple[T1, T2]:
    """Return a pair object with dynamic type dispatch."""
    tp = pairtype(type(a), type(b))
    return tp((a, b))   # tp is a subclass of tuple

pairtypecache: Dict[Tuple[type, type], type] = {}

def pairtype(cls1: type, cls2: type) -> type:
    """type(pair(a,b)) is pairtype(a.__class__, b.__class__)."""
    if (cls1, cls2) in pairtypecache:
        return pairtypecache[(cls1, cls2)]
    
    # Generate a meaningful name for the new pair type
    name = f'pairtype({cls1.__name__}, {cls2.__name__})'
    
    # Create base types for the new pair type
    bases1 = [pairtype(base1, cls2) for base1 in cls1.__bases__]
    bases2 = [pairtype(cls1, base2) for base2 in cls2.__bases__]
    bases = tuple(bases1 + bases2) or (tuple,)  # 'tuple': ultimate base
    
    # Create the new pair type
    pair = pairtypecache[(cls1, cls2)] = extendabletype(name, bases, {})
    return pair

def pairmro(cls1: type, cls2: type) -> Iterable[Tuple[type, type]]:
    """
    Return the resolution order on pairs of types for double dispatch.

    This order is compatible with the mro of pairtype(cls1, cls2).
    """
    for base2 in cls2.__mro__:
        for base1 in cls1.__mro__:
            yield base1, base2

class DoubleDispatchRegistry:
    """
    A mapping of pairs of types to arbitrary objects respecting inheritance
    """
    def __init__(self):
        self._registry: Dict[Tuple[type, type], Callable] = {}
        self._cache: Dict[Tuple[type, type], Callable] = {}

    def __getitem__(self, clspair: Tuple[type, type]) -> Callable:
        cls1, cls2 = clspair
        # Check cache first
        if clspair in self._cache:
            return self._cache[clspair]
        
        # Traverse the MRO to find the closest match
        for c1, c2 in pairmro(cls1, cls2):
            if (c1, c2) in self._cache:
                return self._cache[(c1, c2)]
        
        # If no match found, use the default implementation
        return self._registry.get(clspair, None)

    def __setitem__(self, clspair: Tuple[type, type], value: Callable):
        self._registry[clspair] = value
        self._cache = self._registry.copy()

class DoubleDispatchFunction:
    def __init__(self, default_func: Callable):
        self._registry = DoubleDispatchRegistry()
        self._default = default_func

    def __call__(self, arg1: Any, arg2: Any, *args, **kwargs) -> Any:
        func = self._registry[(type(arg1), type(arg2))]
        if func is None:
            func = self._default
        return func(arg1, arg2, *args, **kwargs)

    def register(self, cls1: Type, cls2: Type) -> Callable:
        def decorator(func: Callable) -> Callable:
            self._registry[(cls1, cls2)] = func
            return func
        return decorator

def doubledispatch(default_func: Optional[Callable] = None) -> Callable:
    """
    Decorator returning a double-dispatch function

    Usage
    -----
        @doubledispatch
        def func(x, y):
            return 0

        @func.register(str, str)
        def func_string_string(x, y):
            return 42

        func(1, 2)        # returns 0
        func('x', 'y')     # returns 42
    """
    def decorator(func):
        return DoubleDispatchFunction(func)
    
    if default_func is None:
        return decorator
    return decorator(default_func)