from typing import Any, Callable

from .functions import (
    TypeConstructorVariable,
    TypeVariable,
    TypeVariableApplication_Nullary,
    TypeVariableApplication_Unary,
)
from .typeclasses import Type3ClassArgs
from .types import (
    KindArgument,
    Type3,
    TypeApplication_Type,
    TypeApplication_TypeInt,
    TypeConstructor_Base,
)


class NoRouteForTypeException(Exception):
    pass

class TypeApplicationRouter[S, R]:
    """
    Helper class to find a method based on a constructed type
    """
    __slots__ = ('by_constructor', 'by_type', )

    by_constructor: dict[Any, Callable[[S, Any], R]]
    """
    Contains all the added routing functions for constructed types
    """

    by_type: dict[Type3, Callable[[S], R]]
    """
    Contains all the added routing functions for constructed types
    """

    def __init__(self) -> None:
        self.by_constructor = {}
        self.by_type = {}

    def add_n(self, typ: Type3, helper: Callable[[S], R]) -> None:
        """
        Lets you route to types that were not constructed

        Also known types of kind *
        """
        self.by_type[typ] = helper

    def add[T](self, constructor: TypeConstructor_Base[T], helper: Callable[[S, T], R]) -> None:
        self.by_constructor[constructor] = helper

    def __call__(self, arg0: S, typ: Type3) -> R:
        t_helper = self.by_type.get(typ)
        if t_helper is not None:
            return t_helper(arg0)

        c_helper = self.by_constructor.get(typ.application.constructor)
        if c_helper is not None:
            return c_helper(arg0, typ.application.arguments)

        raise NoRouteForTypeException(arg0, typ)

TypeVariableLookup = tuple[
    dict[TypeVariable, KindArgument],
    dict[TypeConstructorVariable, TypeConstructor_Base[Any]],
]

class TypeClassArgsRouter[S, R]:
    """
    Helper class to find a method based on a type class argument list
    """
    __slots__ = ('args', 'data', )

    args: Type3ClassArgs

    data: dict[tuple[Type3 | TypeConstructor_Base[Any], ...], Callable[[S, TypeVariableLookup], R]]

    def __init__(self, args: Type3ClassArgs) -> None:
        self.args = args
        self.data = {}

    def add(
            self,
            tv_map: dict[TypeVariable, Type3],
            tc_map: dict[TypeConstructorVariable, TypeConstructor_Base[Any]],
            helper: Callable[[S, TypeVariableLookup], R],
        ) -> None:

        key: list[Type3 | TypeConstructor_Base[Any]] = []

        for tc_arg in self.args:
            if isinstance(tc_arg, TypeVariable):
                key.append(tv_map[tc_arg])
            else:
                key.append(tc_map[tc_arg])

        self.data[tuple(key)] = helper

    def __call__(self, arg0: S, tv_map: dict[TypeVariable, Type3]) -> R:
        key: list[Type3 | TypeConstructor_Base[Any]] = []
        arguments: TypeVariableLookup = (dict(tv_map), {}, )

        for tc_arg in self.args:
            if isinstance(tc_arg, TypeVariable):
                key.append(tv_map[tc_arg])
                arguments[0][tc_arg] = tv_map[tc_arg]
                continue

            for tvar, typ in tv_map.items():
                tvar_constructor = tvar.application.constructor
                if tvar_constructor != tc_arg:
                    continue

                key.append(typ.application.constructor)
                arguments[1][tc_arg] = typ.application.constructor

                if isinstance(tvar.application, TypeVariableApplication_Unary):
                    if isinstance(typ.application, TypeApplication_Type):
                        da_type, = typ.application.arguments
                        sa_type_tv = tvar.application.arguments
                        arguments[0][sa_type_tv] = da_type
                        continue

                    # FIXME: This feels sketchy. Shouldn't the type variable
                    # have the exact same number as arguments?
                    if isinstance(typ.application, TypeApplication_TypeInt):
                        sa_type, sa_len = typ.application.arguments
                        sa_type_tv = tvar.application.arguments
                        sa_len_tv = TypeVariable(sa_type_tv.name + '*', TypeVariableApplication_Nullary(None, None))

                        arguments[0][sa_type_tv] = sa_type
                        arguments[0][sa_len_tv] = sa_len
                        continue

                raise NotImplementedError(tvar.application, typ.application)

                continue

        t_helper = self.data.get(tuple(key))
        if t_helper is not None:
            return t_helper(arg0, arguments)

        raise NoRouteForTypeException(arg0, tv_map)