"""
This module contains possible constraints generated based on the AST

These need to be resolved before the program can be compiled.
"""
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

from .. import ourlang
from ..build.base import BuildBase
from ..wasm import (
    WasmTypeFloat32,
    WasmTypeFloat64,
    WasmTypeInt32,
    WasmTypeInt64,
)
from .functions import FunctionArgument, TypeVariable
from .placeholders import PlaceholderForType, Type3OrPlaceholder
from .routers import NoRouteForTypeException, TypeApplicationRouter
from .typeclasses import Type3Class
from .types import (
    IntType3,
    Type3,
    TypeApplication_Nullary,
    TypeApplication_Struct,
    TypeApplication_Type,
    TypeApplication_TypeInt,
    TypeApplication_TypeStar,
    TypeConstructor_Base,
    TypeConstructor_Struct,
)


class Error:
    """
    An error returned by the check functions for a contraint

    This means the programmer has to make some kind of chance to the
    typing of their program before the compiler can do its thing.
    """
    def __init__(self, msg: str, *, comment: Optional[str] = None) -> None:
        self.msg = msg
        self.comment = comment

    def __repr__(self) -> str:
        return f'Error({repr(self.msg)}, comment={repr(self.comment)})'

class RequireTypeSubstitutes:
    """
    Returned by the check function for a contraint if they do not have all
    their types substituted yet.

    Hopefully, another constraint will give the right information about the
    typing of the program, so this constraint can be updated.
    """

SubstitutionMap = Dict[PlaceholderForType, Type3]

NewConstraintList = List['ConstraintBase']

CheckResult = Union[None, SubstitutionMap, Error, NewConstraintList, RequireTypeSubstitutes]

HumanReadableRet = Tuple[str, Dict[str, Union[None, int, str, ourlang.Expression, Type3, PlaceholderForType]]]

class Context:
    """
    Context for constraints
    """

    __slots__ = ('build', )

    build: BuildBase[Any]

    def __init__(self, build: BuildBase[Any]) -> None:
        self.build = build

class ConstraintBase:
    """
    Base class for constraints
    """
    __slots__ = ('context', 'comment', )

    context: Context
    """
    Additional information regarding the type environment
    """

    comment: Optional[str]
    """
    A comment to help the programmer with debugging the types in their program
    """

    def __init__(self, context: Context, comment: Optional[str] = None) -> None:
        self.context = context
        self.comment = comment

    def check(self) -> CheckResult:
        """
        Checks if the constraint hold

        This function can return an error, if the constraint does not hold,
        which indicates an error in the typing of the input program.

        This function can return RequireTypeSubstitutes(), if we cannot deduce
        all the types yet.

        This function can return a SubstitutionMap, if during the evaluation
        of the contraint we discovered new types. In this case, the constraint
        is expected to hold.

        This function can return None, if the constraint holds, but no new
        information was deduced from evaluating this constraint.
        """
        raise NotImplementedError(self.__class__, self.check)

    def human_readable(self) -> HumanReadableRet:
        """
        Returns a more human readable form of this constraint
        """
        return repr(self), {}

class SameTypeConstraint(ConstraintBase):
    """
    Verifies that a number of types all are the same type
    """
    __slots__ = ('type_list', )

    type_list: List[Type3OrPlaceholder]

    def __init__(self, context: Context, *type_list: Type3OrPlaceholder, comment: Optional[str] = None) -> None:
        super().__init__(context=context, comment=comment)

        assert len(type_list) > 1
        self.type_list = [*type_list]

    def check(self) -> CheckResult:
        known_types: List[Type3] = []
        phft_list = []
        for typ in self.type_list:
            if isinstance(typ, Type3):
                known_types.append(typ)
                continue

            if isinstance(typ, PlaceholderForType):
                if typ.resolve_as is not None:
                    known_types.append(typ.resolve_as)
                else:
                    phft_list.append(typ)
                continue

            raise NotImplementedError(typ)

        if not known_types:
            return RequireTypeSubstitutes()

        first_type = known_types[0]
        for ktyp in known_types[1:]:
            if ktyp != first_type:
                return Error(f'{ktyp:s} must be {first_type:s} instead', comment=self.comment)

        if not phft_list:
            return None

        for phft in phft_list:
            phft.resolve_as = first_type

        return {
            typ: first_type
            for typ in phft_list
        }

    def human_readable(self) -> HumanReadableRet:
        return (
            ' == '.join('{t' + str(idx) + '}' for idx in range(len(self.type_list))),
            {
                't' + str(idx): typ
                for idx, typ in enumerate(self.type_list)
            },
        )

    def __repr__(self) -> str:
        args = ', '.join(repr(x) for x in self.type_list)

        return f'SameTypeConstraint({args}, comment={repr(self.comment)})'

class SameTypeArgumentConstraint(ConstraintBase):
    __slots__ = ('tc_var', 'arg_var', )

    tc_var: PlaceholderForType
    arg_var: PlaceholderForType

    def __init__(self, context: Context, tc_var: PlaceholderForType, arg_var: PlaceholderForType, *, comment: str) -> None:
        super().__init__(context=context, comment=comment)

        self.tc_var = tc_var
        self.arg_var = arg_var

    def check(self) -> CheckResult:
        if self.tc_var.resolve_as is None:
            return RequireTypeSubstitutes()

        tc_typ = self.tc_var.resolve_as
        arg_typ = self.arg_var.resolve_as

        if isinstance(tc_typ.application, TypeApplication_Nullary):
            return Error(f'{tc_typ:s} must be a constructed type instead')

        if isinstance(tc_typ.application, TypeApplication_TypeStar):
            # Sure, it's a constructed type. But it's like a struct,
            # though without the way to implement type classes
            # Presumably, doing a naked `foo :: t a -> a`
            # doesn't work since you don't have any info on t
            # So we can let the MustImplementTypeClassConstraint handle it.
            return None

        if isinstance(tc_typ.application, TypeApplication_Type):
            return [SameTypeConstraint(
                self.context,
                tc_typ.application.arguments[0],
                self.arg_var,
                comment=self.comment,
            )]

        # FIXME: This feels sketchy. Shouldn't the type variable
        # have the exact same number as arguments?
        if isinstance(tc_typ.application, TypeApplication_TypeInt):
            return [SameTypeConstraint(
                self.context,
                tc_typ.application.arguments[0],
                self.arg_var,
                comment=self.comment,
            )]

        raise NotImplementedError(tc_typ, arg_typ)

    def human_readable(self) -> HumanReadableRet:
        return (
            '{tc_var}` == {arg_var}',
            {
                'tc_var': self.tc_var if self.tc_var.resolve_as is None else self.tc_var,
                'arg_var': self.arg_var if self.arg_var.resolve_as is None else self.arg_var,
            },
        )

class SameFunctionArgumentConstraint(ConstraintBase):
    __slots__ = ('type3', 'func_arg', 'type_var_map', )

    type3: PlaceholderForType
    func_arg: FunctionArgument
    type_var_map: dict[TypeVariable, PlaceholderForType]

    def __init__(self, context: Context, type3: PlaceholderForType, func_arg: FunctionArgument, type_var_map: dict[TypeVariable, PlaceholderForType], *, comment: str) -> None:
        super().__init__(context=context, comment=comment)

        self.type3 = type3
        self.func_arg = func_arg
        self.type_var_map = type_var_map

    def check(self) -> CheckResult:
        if self.type3.resolve_as is None:
            return RequireTypeSubstitutes()

        typ = self.type3.resolve_as

        if isinstance(typ.application, TypeApplication_Nullary):
            return Error(f'{typ:s} must be a function instead')

        if not isinstance(typ.application, TypeApplication_TypeStar):
            return Error(f'{typ:s} must be a function instead')

        type_var_map = {
            x: y.resolve_as
            for x, y in self.type_var_map.items()
            if y.resolve_as is not None
        }

        exp_type_arg_list = [
            tv if isinstance(tv, Type3) else type_var_map[tv]
            for tv in self.func_arg.args
            if isinstance(tv, Type3) or tv in type_var_map
        ]

        if len(exp_type_arg_list) != len(self.func_arg.args):
            return RequireTypeSubstitutes()

        return [
            SameTypeConstraint(
                self.context,
                typ,
                self.context.build.function(*exp_type_arg_list),
                comment=self.comment,
            )
        ]

    def human_readable(self) -> HumanReadableRet:
        return (
            '{type3} == {func_arg}',
            {
                'type3': self.type3,
                'func_arg': self.func_arg.name,
            },
        )

class TupleMatchConstraint(ConstraintBase):
    __slots__ = ('exp_type', 'args', 'generate_router', )

    exp_type: Type3OrPlaceholder
    args: list[Type3OrPlaceholder]
    generate_router: TypeApplicationRouter['TupleMatchConstraint', CheckResult]

    def __init__(self, context: Context, exp_type: Type3OrPlaceholder, args: Iterable[Type3OrPlaceholder], comment: str):
        super().__init__(context=context, comment=comment)

        self.exp_type = exp_type
        self.args = list(args)
        self.generate_router = TypeApplicationRouter()
        self.generate_router.add(context.build.dynamic_array, self.__class__._generate_dynamic_array)
        self.generate_router.add(context.build.static_array, self.__class__._generate_static_array)
        self.generate_router.add(context.build.tuple_, self.__class__._generate_tuple)

    def _generate_dynamic_array(self, sa_args: tuple[Type3]) -> CheckResult:
        sa_type, = sa_args

        return [
            SameTypeConstraint(self.context, arg, sa_type)
            for arg in self.args
        ]

    def _generate_static_array(self, sa_args: tuple[Type3, IntType3]) -> CheckResult:
        sa_type, sa_len = sa_args

        if sa_len.value != len(self.args):
            return Error('Mismatch between applied types argument count', comment=self.comment)

        return [
            SameTypeConstraint(self.context, arg, sa_type)
            for arg in self.args
        ]

    def _generate_tuple(self, tp_args: tuple[Type3, ...]) -> CheckResult:
        if len(tp_args) != len(self.args):
            return Error('Mismatch between applied types argument count', comment=self.comment)

        return [
            SameTypeConstraint(self.context, arg, oth_arg)
            for arg, oth_arg in zip(self.args, tp_args, strict=True)
        ]

    def check(self) -> CheckResult:
        exp_type = self.exp_type
        if isinstance(exp_type, PlaceholderForType):
            if exp_type.resolve_as is None:
                return RequireTypeSubstitutes()

            exp_type = exp_type.resolve_as

        try:
            return self.generate_router(self, exp_type)
        except NoRouteForTypeException:
            raise NotImplementedError(exp_type)

class MustImplementTypeClassConstraint(ConstraintBase):
    """
    A type must implement a given type class
    """
    __slots__ = ('type_class3', 'types', )

    type_class3: Type3Class
    types: list[Type3OrPlaceholder]

    def __init__(self, context: Context, type_class3: Type3Class, typ_list: list[Type3OrPlaceholder], comment: Optional[str] = None) -> None:
        super().__init__(context=context, comment=comment)

        self.type_class3 = type_class3
        self.types = typ_list

    def check(self) -> CheckResult:
        typ_list: list[Type3 | TypeConstructor_Base[Any] | TypeConstructor_Struct] = []
        for typ in self.types:
            if isinstance(typ, PlaceholderForType) and typ.resolve_as is not None:
                typ = typ.resolve_as

            if isinstance(typ, PlaceholderForType):
                return RequireTypeSubstitutes()

            if isinstance(typ.application, (TypeApplication_Nullary, TypeApplication_Struct, )):
                typ_list.append(typ)
                continue

            if isinstance(typ.application, (TypeApplication_Type, TypeApplication_TypeInt, TypeApplication_TypeStar)):
                typ_list.append(typ.application.constructor)
                continue

            raise NotImplementedError(typ, typ.application)

        assert len(typ_list) == len(self.types)

        key = (self.type_class3, tuple(typ_list), )
        if key in self.context.build.type_class_instances:
            return None

        typ_cls_name = self.type_class3 if isinstance(self.type_class3, str) else self.type_class3.name
        typ_name_list = ' '.join(x.name for x in typ_list)
        return Error(f'Missing type class instantation: {typ_cls_name} {typ_name_list}')

    def human_readable(self) -> HumanReadableRet:
        keys = {
            f'type{idx}': typ
            for idx, typ in enumerate(self.types)
        }

        return (
            'Exists instance {type_class3} ' + ' '.join(f'{{{x}}}' for x in keys),
            {
                'type_class3': str(self.type_class3),
                **keys,
            },
        )

    def __repr__(self) -> str:
        return f'MustImplementTypeClassConstraint({repr(self.type_class3)}, {repr(self.types)}, comment={repr(self.comment)})'

class LiteralFitsConstraint(ConstraintBase):
    """
    A literal value fits a given type
    """
    __slots__ = ('type3', 'literal', 'generate_router', )

    type3: Type3OrPlaceholder
    literal: Union[ourlang.ConstantPrimitive, ourlang.ConstantBytes, ourlang.ConstantTuple, ourlang.ConstantStruct]
    generate_router: TypeApplicationRouter['LiteralFitsConstraint', CheckResult]

    def __init__(
            self,
            context: Context,
            type3: Type3OrPlaceholder,
            literal: Union[ourlang.ConstantPrimitive, ourlang.ConstantBytes, ourlang.ConstantTuple, ourlang.ConstantStruct],
            comment: Optional[str] = None,
        ) -> None:
        super().__init__(context=context, comment=comment)

        self.type3 = type3
        self.literal = literal

        self.generate_router = TypeApplicationRouter['LiteralFitsConstraint', CheckResult]()
        self.generate_router.add(context.build.dynamic_array, self.__class__._generate_dynamic_array)
        self.generate_router.add(context.build.static_array, self.__class__._generate_static_array)
        self.generate_router.add(context.build.struct, self.__class__._generate_struct)
        self.generate_router.add(context.build.tuple_, self.__class__._generate_tuple)

    def _generate_dynamic_array(self, da_args: tuple[Type3]) -> CheckResult:
        da_type, = da_args

        if da_type.name == 'u8':
            if not isinstance(self.literal, ourlang.ConstantBytes):
                return Error('Must be bytes', comment=self.comment)

            return None

        if not isinstance(self.literal, ourlang.ConstantTuple):
            return Error('Must be tuple', comment=self.comment)

        res: list[ConstraintBase] = []

        res.extend(
            LiteralFitsConstraint(self.context, da_type, y)
            for y in self.literal.value
        )

        # Generate placeholders so each Literal expression
        # gets updated when we figure out the type of the
        # expression the literal is used in
        res.extend(
            SameTypeConstraint(self.context, da_type, PlaceholderForType([y]))
            for y in self.literal.value
        )

        return res

    def _generate_static_array(self, sa_args: tuple[Type3, IntType3]) -> CheckResult:
        if not isinstance(self.literal, ourlang.ConstantTuple):
            return Error('Must be tuple', comment=self.comment)

        sa_type, sa_len = sa_args

        if sa_len.value != len(self.literal.value):
            return Error('Member count mismatch', comment=self.comment)

        res: list[ConstraintBase] = []

        res.extend(
            LiteralFitsConstraint(self.context, sa_type, y)
            for y in self.literal.value
        )

        # Generate placeholders so each Literal expression
        # gets updated when we figure out the type of the
        # expression the literal is used in
        res.extend(
            SameTypeConstraint(self.context, sa_type, PlaceholderForType([y]))
            for y in self.literal.value
        )

        return res

    def _generate_struct(self, st_args: tuple[tuple[str, Type3], ...]) -> CheckResult:
        if not isinstance(self.literal, ourlang.ConstantStruct):
            return Error('Must be struct')

        if len(st_args) != len(self.literal.value):
            return Error('Struct element count mismatch')

        res: list[ConstraintBase] = []

        res.extend(
            LiteralFitsConstraint(self.context, x, y)
            for (_, x), y in zip(st_args, self.literal.value, strict=True)
        )

        # Generate placeholders so each Literal expression
        # gets updated when we figure out the type of the
        # expression the literal is used in
        res.extend(
            SameTypeConstraint(self.context, x_t, PlaceholderForType([y]), comment=f'{self.literal.struct_type3.name}.{x_n}')
            for (x_n, x_t, ), y in zip(st_args, self.literal.value, strict=True)
        )

        res.append(SameTypeConstraint(
            self.context,
            self.literal.struct_type3,
            self.type3,
            comment='Struct types must match',
        ))

        return res

    def _generate_tuple(self, tp_args: tuple[Type3, ...]) -> CheckResult:
        if not isinstance(self.literal, ourlang.ConstantTuple):
            return Error('Must be tuple', comment=self.comment)

        if len(tp_args) != len(self.literal.value):
            return Error('Tuple element count mismatch', comment=self.comment)

        res: list[ConstraintBase] = []

        res.extend(
            LiteralFitsConstraint(self.context, x, y)
            for x, y in zip(tp_args, self.literal.value, strict=True)
        )

        # Generate placeholders so each Literal expression
        # gets updated when we figure out the type of the
        # expression the literal is used in
        res.extend(
            SameTypeConstraint(self.context, x, PlaceholderForType([y]))
            for x, y in zip(tp_args, self.literal.value, strict=True)
        )

        return res

    def check(self) -> CheckResult:
        if isinstance(self.type3, PlaceholderForType):
            if self.type3.resolve_as is None:
                return RequireTypeSubstitutes()

            self.type3 = self.type3.resolve_as

        type_info = self.context.build.type_info_map.get(self.type3.name)

        if type_info is not None and (type_info.wasm_type is WasmTypeInt32 or type_info.wasm_type is WasmTypeInt64):
            assert type_info.signed is not None

            if isinstance(self.literal.value, int):
                try:
                    self.literal.value.to_bytes(type_info.alloc_size, 'big', signed=type_info.signed)
                except OverflowError:
                    return Error(f'Must fit in {type_info.alloc_size} byte(s)', comment=self.comment) # FIXME: Add line information

                return None

            return Error('Must be integer', comment=self.comment) # FIXME: Add line information

        if type_info is not None and (type_info.wasm_type is WasmTypeFloat32 or type_info.wasm_type is WasmTypeFloat64):
            if isinstance(self.literal.value, float):
                # FIXME: Bit check

                return None

            return Error('Must be real', comment=self.comment) # FIXME: Add line information

        exp_type = self.type3

        try:
            return self.generate_router(self, exp_type)
        except NoRouteForTypeException:
            raise NotImplementedError(exp_type)

    def human_readable(self) -> HumanReadableRet:
        return (
            '{literal} : {type3}',
            {
                'literal': self.literal,
                'type3': self.type3,
            },
        )

    def __repr__(self) -> str:
        return f'LiteralFitsConstraint({repr(self.type3)}, {repr(self.literal)}, comment={repr(self.comment)})'

class CanBeSubscriptedConstraint(ConstraintBase):
    """
    A value that is subscipted, i.e. a[0] (tuple) or a[b] (static array)
    """
    __slots__ = ('ret_type3', 'type3', 'index_type3', 'index_const', 'generate_router', )

    ret_type3: PlaceholderForType
    type3: PlaceholderForType
    index_type3: PlaceholderForType
    index_const: int | None
    generate_router: TypeApplicationRouter['CanBeSubscriptedConstraint', CheckResult]

    def __init__(
            self,
            context: Context,
            ret_type3: PlaceholderForType,
            type3: PlaceholderForType,
            index_type3: PlaceholderForType,
            index_const: int | None,
            comment: Optional[str] = None,
        ) -> None:
        super().__init__(context=context, comment=comment)

        self.ret_type3 = ret_type3
        self.type3 = type3
        self.index_type3 = index_type3
        self.index_const = index_const


        self.generate_router = TypeApplicationRouter()
        self.generate_router.add_n(context.build.types['bytes'], self.__class__._generate_bytes)
        self.generate_router.add(context.build.static_array, self.__class__._generate_static_array)
        self.generate_router.add(context.build.tuple_, self.__class__._generate_tuple)

    def _generate_bytes(self) -> CheckResult:
        return [
            SameTypeConstraint(self.context, self.context.build.types['u32'], self.index_type3, comment='([]) :: bytes -> u32 -> u8'),
            SameTypeConstraint(self.context, self.context.build.types['u8'], self.ret_type3, comment='([]) :: bytes -> u32 -> u8'),
        ]

    def _generate_static_array(self, sa_args: tuple[Type3, IntType3]) -> CheckResult:
        sa_type, sa_len = sa_args

        if self.index_const is not None and (self.index_const < 0 or sa_len.value <= self.index_const):
            return Error('Tuple index out of range')

        return [
            SameTypeConstraint(self.context, self.context.build.types['u32'], self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
            SameTypeConstraint(self.context, sa_type, self.ret_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
        ]

    def _generate_tuple(self, tp_args: tuple[Type3, ...]) -> CheckResult:
        # We special case tuples to allow for ease of use to the programmer
        # e.g. rather than having to do `fst a` and `snd a` and only have to-sized tuples
        # we use a[0] and a[1] and allow for a[2] and on.

        if self.index_const is None:
            return Error('Must index with integer literal')

        if self.index_const < 0 or len(tp_args) <= self.index_const:
            return Error('Tuple index out of range')

        return [
            SameTypeConstraint(self.context, self.context.build.types['u32'], self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
            SameTypeConstraint(self.context, tp_args[self.index_const], self.ret_type3, comment=f'Tuple subscript index {self.index_const}'),
        ]

    def check(self) -> CheckResult:
        if self.type3.resolve_as is None:
            return RequireTypeSubstitutes()

        exp_type = self.type3.resolve_as

        try:
            return self.generate_router(self, exp_type)
        except NoRouteForTypeException:
            return Error(f'{exp_type.name} cannot be subscripted')

    def human_readable(self) -> HumanReadableRet:
        return (
            '{type3}[{index}]',
            {
                'type3': self.type3,
                'index': self.index_type3 if self.index_const is None else self.index_const,
            },
        )

    def __repr__(self) -> str:
        return f'CanBeSubscriptedConstraint({self.ret_type3!r}, {self.type3!r}, {self.index_type3!r}, {self.index_const!r}, comment={repr(self.comment)})'