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

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

from .. import ourlang
from . import typeclasses, types


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[types.PlaceholderForType, types.Type3]

NewConstraintList = List['ConstraintBase']

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

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

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

    __slots__ = ()

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

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

    def __init__(self, comment: Optional[str] = None) -> None:
        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[types.Type3OrPlaceholder]

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

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

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

            if isinstance(typ, (types.PrimitiveType3, types.StructType3, )):
                known_types.append(typ)
                continue

            if isinstance(typ, types.AppliedType3):
                known_types.append(typ)
                continue

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

            raise NotImplementedError(typ)

        if not known_types:
            return RequireTypeSubstitutes()

        new_constraint_list: List[ConstraintBase] = []

        first_type = known_types[0]
        for typ in known_types[1:]:
            if isinstance(first_type, types.AppliedType3) and isinstance(typ, types.AppliedType3):
                if first_type.base is types.tuple and typ.base is types.static_array:
                    # Swap so we can reuse the code below
                    # Hope that it still gives proper type errors
                    first_type, typ = typ, first_type

                if first_type.base is types.static_array and typ.base is types.tuple:
                    assert isinstance(first_type.args[1], types.IntType3)
                    length = first_type.args[1].value

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

                    for typ_arg in typ.args:
                        new_constraint_list.append(SameTypeConstraint(
                            first_type.args[0], typ_arg
                        ))
                    continue

                if first_type.base != typ.base:
                    return Error('Mismatch between applied types base', comment=self.comment)

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

                for first_type_arg, typ_arg in zip(first_type.args, typ.args):
                    new_constraint_list.append(SameTypeConstraint(
                        first_type_arg, typ_arg
                    ))

                continue

            if typ != first_type:
                return Error(f'{typ:s} must be {first_type:s} instead', comment=self.comment)

        if new_constraint_list:
            # If this happens, make CheckResult a class that can have both
            assert not placeholders, 'Cannot (yet) return both new placeholders and new constraints'

            return new_constraint_list

        if not placeholders:
            return None

        for typ in placeholders:
            typ.resolve_as = first_type

        return {
            typ: first_type
            for typ in placeholders
        }

    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 IntegerCompareConstraint(ConstraintBase):
    """
    Verifies that the given IntType3 are in order (<=)
    """
    __slots__ = ('int_type3_list', )

    int_type3_list: List[types.IntType3]

    def __init__(self, *int_type3: types.IntType3, comment: Optional[str] = None) -> None:
        super().__init__(comment=comment)

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

    def check(self) -> CheckResult:
        val_list = [x.value for x in self.int_type3_list]

        prev_val = val_list.pop(0)
        for next_val in val_list:
            if prev_val > next_val:
                return Error(f'{prev_val} must be less or equal than {next_val}')

            prev_val = next_val

        return None

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

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

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

class CastableConstraint(ConstraintBase):
    """
    A type can be cast to another type
    """
    __slots__ = ('from_type3', 'to_type3', )

    from_type3: types.Type3OrPlaceholder
    to_type3: types.Type3OrPlaceholder

    def __init__(self, from_type3: types.Type3OrPlaceholder, to_type3: types.Type3OrPlaceholder, comment: Optional[str] = None) -> None:
        super().__init__(comment=comment)

        self.from_type3 = from_type3
        self.to_type3 = to_type3

    def check(self) -> CheckResult:
        ftyp = self.from_type3
        if isinstance(ftyp, types.PlaceholderForType) and ftyp.resolve_as is not None:
            ftyp = ftyp.resolve_as

        ttyp = self.to_type3
        if isinstance(ttyp, types.PlaceholderForType) and ttyp.resolve_as is not None:
            ttyp = ttyp.resolve_as

        if isinstance(ftyp, types.PlaceholderForType) or isinstance(ttyp, types.PlaceholderForType):
            return RequireTypeSubstitutes()

        if ftyp is types.u8 and ttyp is types.u32:
            return None

        return Error(f'Cannot cast {ftyp.name} to {ttyp.name}')

    def human_readable(self) -> HumanReadableRet:
        return (
            '{to_type3}({from_type3})',
            {
                'to_type3': self.to_type3,
                'from_type3': self.from_type3,
            },
        )

    def __repr__(self) -> str:
        return f'CastableConstraint({repr(self.from_type3)}, {repr(self.to_type3)}, comment={repr(self.comment)})'

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

    type_class3: Union[str, typeclasses.Type3Class]
    type3: types.Type3OrPlaceholder

    DATA = {
        'u8': {'BitWiseOperation', 'EqualComparison', 'StrictPartialOrder'},
        'u32': {'BitWiseOperation', 'EqualComparison', 'StrictPartialOrder'},
        'u64': {'BitWiseOperation', 'EqualComparison', 'StrictPartialOrder'},
        'i32': {'EqualComparison', 'StrictPartialOrder'},
        'i64': {'EqualComparison', 'StrictPartialOrder'},
        'bytes': {'Foldable', 'Sized'},
        'f32': {'Fractional', 'FloatingPoint'},
        'f64': {'Fractional', 'FloatingPoint'},
    }

    def __init__(self, type_class3: Union[str, typeclasses.Type3Class], type3: types.Type3OrPlaceholder, comment: Optional[str] = None) -> None:
        super().__init__(comment=comment)

        self.type_class3 = type_class3
        self.type3 = type3

    def check(self) -> CheckResult:
        typ = self.type3
        if isinstance(typ, types.PlaceholderForType) and typ.resolve_as is not None:
            typ = typ.resolve_as

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

        if isinstance(self.type_class3, typeclasses.Type3Class):
            if self.type_class3 in typ.classes:
                return None
        else:
            if self.type_class3 in self.__class__.DATA.get(typ.name, set()):
                return None

        return Error(f'{typ.name} does not implement the {self.type_class3} type class')

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

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

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

    type3: types.Type3OrPlaceholder
    literal: Union[ourlang.ConstantPrimitive, ourlang.ConstantBytes, ourlang.ConstantTuple, ourlang.ConstantStruct]

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

        self.type3 = type3
        self.literal = literal

    def check(self) -> CheckResult:
        int_table: Dict[str, Tuple[int, bool]] = {
            'u8': (1, False),
            'u32': (4, False),
            'u64': (8, False),
            'i8': (1, True),
            'i32': (4, True),
            'i64': (8, True),
        }

        float_table: Dict[str, None] = {
            'f32': None,
            'f64': None,
        }

        if isinstance(self.type3, types.PlaceholderForType):
            if self.type3.resolve_as is None:
                return RequireTypeSubstitutes()

            self.type3 = self.type3.resolve_as

        if self.type3.name in int_table:
            bts, sgn = int_table[self.type3.name]

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

                return None

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

        if self.type3.name in float_table:
            _ = float_table[self.type3.name]

            if isinstance(self.literal.value, float):
                # FIXME: Bit check

                return None

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

        if self.type3 is types.bytes:
            if isinstance(self.literal.value, bytes):
                return None

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

        res: NewConstraintList

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

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

                res = []

                res.extend(
                    LiteralFitsConstraint(x, y)
                    for x, y in zip(self.type3.args, self.literal.value)
                )
                res.extend(
                    SameTypeConstraint(x, y.type3)
                    for x, y in zip(self.type3.args, self.literal.value)
                )

                return res

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

                assert 2 == len(self.type3.args)
                assert isinstance(self.type3.args[1], types.IntType3)

                if self.type3.args[1].value != len(self.literal.value):
                    return Error('Member count mismatch', comment=self.comment)

                res = []

                res.extend(
                    LiteralFitsConstraint(self.type3.args[0], y)
                    for y in self.literal.value
                )
                res.extend(
                    SameTypeConstraint(self.type3.args[0], y.type3)
                    for y in self.literal.value
                )

                return res

        if isinstance(self.type3, types.StructType3):
            if not isinstance(self.literal, ourlang.ConstantStruct):
                return Error('Must be struct')

            if self.literal.struct_name != self.type3.name:
                return Error('Struct mismatch')


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

            res = []

            res.extend(
                LiteralFitsConstraint(x, y)
                for x, y in zip(self.type3.members.values(), self.literal.value)
            )
            res.extend(
                SameTypeConstraint(x_t, y.type3, comment=f'{self.literal.struct_name}.{x_n}')
                for (x_n, x_t, ), y in zip(self.type3.members.items(), self.literal.value)
            )

            return res

        raise NotImplementedError(self.type3, self.literal)

    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', 'index_type3', )

    ret_type3: types.Type3OrPlaceholder
    type3: types.Type3OrPlaceholder
    index: ourlang.Expression
    index_type3: types.Type3OrPlaceholder

    def __init__(self, ret_type3: types.Type3OrPlaceholder, type3: types.Type3OrPlaceholder, index: ourlang.Expression, comment: Optional[str] = None) -> None:
        super().__init__(comment=comment)

        self.ret_type3 = ret_type3
        self.type3 = type3
        self.index = index
        self.index_type3 = index.type3

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

            self.type3 = self.type3.resolve_as

        if isinstance(self.type3, types.AppliedType3):
            if self.type3.base == types.static_array:
                result: List[ConstraintBase] = [
                    SameTypeConstraint(types.u32, self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
                    SameTypeConstraint(self.type3.args[0], self.ret_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
                ]

                if isinstance(self.index, ourlang.ConstantPrimitive):
                    assert isinstance(self.index.value, int)
                    assert isinstance(self.type3.args[1], types.IntType3)

                    result.append(
                        IntegerCompareConstraint(
                            types.IntType3(0), types.IntType3(self.index.value), types.IntType3(self.type3.args[1].value - 1),
                            comment='Subscript static array must fit the size of the array'
                        )
                    )

                return result

            if self.type3.base == types.tuple:
                if not isinstance(self.index, ourlang.ConstantPrimitive):
                    return Error('Must index with literal')

                if not isinstance(self.index.value, int):
                    return Error('Must index with integer literal')

                if self.index.value < 0 or len(self.type3.args) <= self.index.value:
                    return Error('Tuple index out of range')

                return [
                    SameTypeConstraint(types.u32, self.index_type3, comment=f'Tuple subscript index {self.index.value}'),
                    SameTypeConstraint(self.type3.args[self.index.value], self.ret_type3, comment=f'Tuple subscript index {self.index.value}'),
                ]

        if self.type3 is types.bytes:
            return [
                SameTypeConstraint(types.u32, self.index_type3, comment='([]) :: bytes -> u32 -> u8'),
                SameTypeConstraint(types.u8, self.ret_type3, comment='([]) :: bytes -> u32 -> u8'),
            ]

        if self.type3.name in types.LOOKUP_TABLE:
            return Error(f'{self.type3.name} cannot be subscripted')

        raise NotImplementedError(self.type3)

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

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