from __future__ import annotations

from loguru import logger

from aeon.core.liquid import LiquidLiteralBool
from aeon.core.liquid import LiquidVar
from aeon.core.substitutions import substitution_in_liquid
from aeon.core.substitutions import substitution_in_type
from aeon.core.terms import Var
from aeon.core.types import AbstractionType, TypeConstructor, TypeVar
from aeon.core.types import RefinedType
from aeon.core.types import Top
from aeon.core.types import Type
from aeon.core.types import TypePolymorphism
from aeon.core.types import t_unit
from aeon.typechecking.context import TypingContext
from aeon.utils.location import Location
from aeon.verification.vcs import Conjunction, UninterpretedFunctionDeclaration
from aeon.verification.vcs import Constraint
from aeon.verification.vcs import Implication
from aeon.verification.vcs import LiquidConstraint
from aeon.utils.name import Name, fresh_counter

ctrue = LiquidConstraint(LiquidLiteralBool(True))
cfalse = LiquidConstraint(LiquidLiteralBool(False))


def ensure_refined(t: Type) -> Type:
    """Ensures that the Base Types and TypeVars are refined. All other types remain the same."""
    match t:
        case RefinedType(_, _, _):
            return t
        case TypeVar(Name(name)):
            return RefinedType(Name(f"singleton_{name}", fresh_counter.fresh()), t, LiquidLiteralBool(True))
        case TypeConstructor(Name(name), _):
            return RefinedType(Name(f"singleton_{name}", fresh_counter.fresh()), t, LiquidLiteralBool(True))
        case _:
            return t


def is_first_order_function(at: AbstractionType):
    v: Type = at
    while isinstance(v, AbstractionType):
        match v.var_type:
            case AbstractionType(_, _, _):
                return False
            case Top() | RefinedType(_, _, _) | TypeVar(_) | TypeConstructor(_, _):
                pass
            case _:
                assert False
        v = v.type
    return True


def lower_constraint_type(ttype: Type) -> Type:
    match ttype:
        case TypeVar(name):
            return TypeConstructor(name)
        case Top():
            return t_unit
        case AbstractionType(name, b, r):
            return AbstractionType(name, lower_constraint_type(b), lower_constraint_type(r))
        case RefinedType(_, t, _):
            return lower_constraint_type(t)
        case TypeConstructor(name, args):
            if args:
                # Polymorphic types are represented by top
                return TypeConstructor(Name("Top", 0))
            else:
                return TypeConstructor(name)

        case _:
            assert False, f"Unsupport type in constraint {ttype} ({type(ttype)})"


def implication_constraint(name: Name, ty: Type, c: Constraint, loc: Location | None = None) -> Constraint:
    match ty:
        case Top() | TypeVar(_) | TypeConstructor(_, _):
            basety = lower_constraint_type(ty)
            assert isinstance(basety, TypeConstructor)
            return Implication(name, basety, LiquidLiteralBool(True), c, loc=loc)
        case RefinedType(tname, ttype, tref):
            ref_subs = substitution_in_liquid(tref, LiquidVar(name), tname)
            ltype = lower_constraint_type(ttype)
            assert isinstance(ltype, TypeConstructor) or isinstance(ltype, TypeVar)
            return Implication(name, ltype, ref_subs, c, loc=loc)
        case AbstractionType(_, _, _):
            # TODO Poly Refl: instead of true, reflect the implementation of the function?
            if is_first_order_function(ty):
                absty = lower_constraint_type(ty)
                assert isinstance(absty, AbstractionType)
                return UninterpretedFunctionDeclaration(name, absty, c)
            else:
                return c
        case TypePolymorphism(_, _, _):
            return c
        case _:
            assert False


def sub(ctx: TypingContext, t1: Type, t2: Type, loc: Location | None = None) -> Constraint:
    if t2 == Top():
        return ctrue
    match (ensure_refined(t1), ensure_refined(t2)):
        case RefinedType(n1, ty1, r1), RefinedType(n2, ty2, r2):
            if ty2 == Top():
                return ctrue
            if ty1 != ty2:
                return cfalse

            new_name: Name = Name(n1.name + n2.name, fresh_counter.fresh())

            # Performs substition on t2 to have the same name of t1
            r2_ = substitution_in_liquid(r2, LiquidVar(new_name), n2)
            r1_ = substitution_in_liquid(r1, LiquidVar(new_name), n1)
            lowert = lower_constraint_type(ty1)
            assert isinstance(lowert, TypeConstructor)
            rconstraint = Implication(new_name, lowert, r1_, LiquidConstraint(r2_, loc=loc), loc=loc)

            return rconstraint
        case TypePolymorphism(_, _, _), _:
            return ctrue
        case AbstractionType(a1, t1, rt1), AbstractionType(a2, t2, rt2):
            new_name_a: Name = Name(a1.name + a2.name, fresh_counter.fresh())

            c0 = sub(ctx, t2, t1, loc)
            rt1_ = substitution_in_type(rt1, Var(new_name_a), a1)
            rt2_ = substitution_in_type(rt2, Var(new_name_a), a2)
            c1 = sub(ctx, rt1_, rt2_, loc)
            return Conjunction(c0, implication_constraint(new_name_a, t2, c1, loc))
        case TypeConstructor(name1, args1), TypeConstructor(name2, args2):
            if name1 != name2:
                return cfalse
            if len(args1) != len(args2):
                return cfalse
            for it1, it2 in zip(args1, args2):
                if it1 != it2:  # TODO polytypes: subtyping here?
                    return cfalse
            return ctrue
        case _:
            logger.error(f"Failed subtyping by exhaustion: {t1} <: {t2}")
            return cfalse
