from __future__ import annotations

from functools import reduce
from typing import Any, Generator

from aeon.core.types import builtin_core_types
from aeon.core.liquid import LiquidApp
from aeon.core.types import LiquidHornApplication, TypeConstructor
from aeon.core.liquid import LiquidLiteralBool
from aeon.core.liquid import LiquidLiteralUnit
from aeon.core.liquid import LiquidLiteralFloat
from aeon.core.liquid import LiquidLiteralInt
from aeon.core.liquid import LiquidLiteralString
from aeon.core.liquid import LiquidTerm
from aeon.core.liquid import LiquidVar
from aeon.core.liquid_ops import mk_liquid_and
from aeon.core.substitutions import substitution_in_liquid
from aeon.core.types import AbstractionType
from aeon.core.types import ExistentialType
from aeon.core.types import RefinedType
from aeon.core.types import t_bool
from aeon.core.types import Top
from aeon.core.types import Type
from aeon.core.types import TypePolymorphism
from aeon.core.types import RefinementPolymorphism
from aeon.core.types import TypeVar
from aeon.core.liquid_ops import liquid_prelude
from aeon.typechecking.context import TypingContext
from aeon.typechecking.context import UninterpretedBinder
from aeon.typechecking.liquid import (
    LiquidTypeCheckException,
    LiquidTypeCheckingContext,
    check_liquid,
    liquid_admissible,
    lower_abstraction_type,
    lower_context,
    type_infer_liquid,
)
from aeon.typechecking.qualifiers import extract_qualifier_atoms
from aeon.verification.helpers import constraint_builder
from aeon.verification.helpers import end
from aeon.verification.helpers import imp
from aeon.verification.smt import smt_valid
from aeon.verification.vcs import Conjunction
from aeon.verification.vcs import Constraint
from aeon.verification.vcs import Implication
from aeon.verification.vcs import LiquidConstraint
from aeon.verification.vcs import ReflectedFunctionDeclaration
from aeon.verification.vcs import UninterpretedFunctionDeclaration
from aeon.utils.name import Name, fresh_counter

Assignment = dict[Name, list[LiquidTerm]]


def smt_base_type(ty: Type) -> str | None:
    if isinstance(ty, AbstractionType):
        return None
    if isinstance(ty, TypeConstructor):
        return str(ty)
    elif isinstance(ty, RefinedType):
        return smt_base_type(ty.type)
    else:
        return None


def fresh(context: TypingContext, ty: Type) -> Type:
    match ty:
        case RefinedType(name, ity, LiquidHornApplication(_, _)):
            vname = Name("√", fresh_counter.fresh())
            hole_name = Name("hole", fresh_counter.fresh())

            argtypes: list[tuple[LiquidTerm, TypeConstructor | TypeVar]] = []
            for n, t in context.vars() + [(vname, ty.type)]:
                sort = liquid_admissible(t)
                if sort is not None:
                    argtypes.append((LiquidVar(n), sort))
            return RefinedType(
                vname,
                ity,
                LiquidHornApplication(hole_name, argtypes),
            )
        case RefinedType(_, _, _) | Top() | TypeVar():
            return ty
        case AbstractionType(name, aty, rty):
            sp = fresh(context, aty)
            tp = fresh(context.with_var(name, aty), rty)
            return AbstractionType(name, sp, tp, multiplicity=ty.multiplicity)
        case TypePolymorphism(name, kind, body):
            return TypePolymorphism(name, kind, fresh(context, body))
        case RefinementPolymorphism(name, sort, body):
            return RefinementPolymorphism(name, fresh(context, sort), fresh(context, body))
        case TypeConstructor(name, args):
            return TypeConstructor(name, [fresh(context, c) for c in args])
        case ExistentialType(binders, body):
            ext_ctx = context
            new_binders: list[tuple[Name, Type]] = []
            for bn, bt in binders:
                fresh_bt = fresh(ext_ctx, bt)
                new_binders.append((bn, fresh_bt))
                ext_ctx = ext_ctx.with_var(bn, fresh_bt)
            return ExistentialType(tuple(new_binders), fresh(ext_ctx, body))
        case _:
            assert False, f"Type not freshable: {ty}, {type(ty)}"


def obtain_holes(t: LiquidTerm) -> list[LiquidHornApplication]:
    if isinstance(t, LiquidHornApplication):
        return [t]
    elif (
        isinstance(t, LiquidLiteralBool)
        or isinstance(t, LiquidLiteralInt)
        or isinstance(t, LiquidLiteralFloat)
        or isinstance(t, LiquidLiteralString)
        or isinstance(t, LiquidLiteralUnit)
    ):
        return []
    elif isinstance(t, LiquidVar):
        return []
    elif isinstance(t, LiquidApp):
        holes: list[LiquidHornApplication] = []
        for h in t.args:
            holes = holes + obtain_holes(h)
        return holes
    else:
        assert False, f"Unkown term type: {t} ({type(t)})"


def obtain_holes_constraint(c: Constraint) -> list[LiquidHornApplication]:
    match c:
        case Conjunction(c1, c2):
            return obtain_holes_constraint(c1) + obtain_holes_constraint(c2)
        case Implication(_, _, pre, post):
            return obtain_holes(pre) + obtain_holes_constraint(post)
        case LiquidConstraint(e):
            return obtain_holes(e)
        case UninterpretedFunctionDeclaration(_, _, post):
            return obtain_holes_constraint(post)
        case ReflectedFunctionDeclaration(_, _, _, _, post):
            return obtain_holes_constraint(post)
        case _:
            assert False, c


def contains_horn(t: LiquidTerm) -> bool:
    if (
        isinstance(t, LiquidLiteralInt)
        or isinstance(t, LiquidLiteralBool)
        or isinstance(t, LiquidLiteralString)
        or isinstance(t, LiquidLiteralFloat)
        or isinstance(t, LiquidLiteralUnit)
    ):
        return False
    elif isinstance(t, LiquidVar):
        return False
    elif isinstance(t, LiquidHornApplication):
        return True
    elif isinstance(t, LiquidApp):
        return all([contains_horn(arg) for arg in t.args])
    else:
        assert False


def contains_horn_constraint(c: Constraint) -> bool:
    if isinstance(c, LiquidConstraint):
        return contains_horn(c.expr)
    elif isinstance(c, Conjunction):
        return contains_horn_constraint(c.c1) or contains_horn_constraint(c.c2)
    elif isinstance(c, Implication):
        return contains_horn(c.pred) or contains_horn_constraint(c.seq)
    elif isinstance(c, UninterpretedFunctionDeclaration):
        return contains_horn_constraint(c.seq)
    elif isinstance(c, ReflectedFunctionDeclaration):
        return contains_horn(c.body) or contains_horn_constraint(c.seq)
    else:
        assert False


def wellformed_horn(predicate: LiquidTerm) -> bool:
    if not contains_horn(predicate):
        return True
    elif (
        isinstance(predicate, LiquidApp)
        and predicate.fun == "&&"
        and not contains_horn(predicate.args[0])
        and isinstance(predicate.args[1], LiquidHornApplication)
    ):
        return True
    elif isinstance(predicate, LiquidHornApplication):
        return True
    else:
        return False


def mk_arg(i: int) -> Name:
    return Name(f"arg_{i}", 0)


def _liquid_var_names_in_term(t: LiquidTerm) -> set[Name]:
    match t:
        case LiquidVar(n):
            return {n}
        case LiquidApp(_, args):
            return {v for a in args for v in _liquid_var_names_in_term(a)}
        case LiquidHornApplication():
            return set()
        case (
            LiquidLiteralBool()
            | LiquidLiteralInt()
            | LiquidLiteralFloat()
            | LiquidLiteralString()
            | LiquidLiteralUnit()
        ):
            return set()
        case _:
            return set()


def liquid_typechecking_ctx_for_hole(
    typing_ctx: TypingContext | None,
    hole: LiquidHornApplication,
) -> LiquidTypeCheckingContext:
    vars_h = {mk_arg(i): ty for i, (_, ty) in enumerate(hole.argtypes)}
    if typing_ctx is None:
        return LiquidTypeCheckingContext(
            known_types=[TypeConstructor(Name(bn.name.name, 0)) for bn in builtin_core_types],
            functions=dict(liquid_prelude),
            variables=vars_h,
        )
    lc = lower_context(typing_ctx)
    fn = dict(liquid_prelude)
    for k, v in lc.functions.items():
        if k not in fn:
            fn[k] = v
    return LiquidTypeCheckingContext(lc.known_types, vars_h, fn)


def adapt_qualifier_to_hole(hole: LiquidHornApplication, q: LiquidTerm) -> LiquidTerm | None:
    t: LiquidTerm = q
    for i, (slot, _) in enumerate(hole.argtypes):
        match slot:
            case LiquidVar(nm):
                t = substitution_in_liquid(t, LiquidVar(mk_arg(i)), nm)
            case _:
                return None
    for n in _liquid_var_names_in_term(t):
        if not n.name.startswith("arg_"):
            return None
    return t


def build_qualifier_candidates(
    typing_ctx: TypingContext | None,
    hole: LiquidHornApplication,
    atoms: frozenset[LiquidTerm],
) -> list[LiquidTerm]:
    if not atoms:
        return []
    liq_ctx = liquid_typechecking_ctx_for_hole(typing_ctx, hole)
    seen: set[LiquidTerm] = set()
    out: list[LiquidTerm] = []
    for q in atoms:
        adapted = adapt_qualifier_to_hole(hole, q)
        if adapted is None:
            continue
        if not check_liquid(liq_ctx, adapted, t_bool):
            continue
        if adapted in seen:
            continue
        seen.add(adapted)
        out.append(adapted)
    return out


def get_possible_args(
    vars: list[tuple[LiquidTerm, TypeConstructor | TypeVar]],
    arity: int,
    extra_atoms: list[LiquidTerm] | None = None,
):
    extra = extra_atoms or []
    if arity == 0:
        yield []
    else:
        for base in get_possible_args(vars, arity - 1, extra):
            for i, (_, _) in enumerate(vars):
                yield [LiquidVar(mk_arg(i))] + base
                yield [LiquidLiteralBool(True)] + base
                yield [LiquidLiteralBool(False)] + base
                yield [LiquidLiteralInt(0)] + base
                yield [LiquidLiteralInt(1)] + base
            for atom in extra:
                yield [atom] + base


def context_measures(typing_ctx: TypingContext | None) -> dict[Name, list[TypeConstructor | TypeVar]]:
    """Custom measures available for Horn predicate generation.

    These are the user-declared uninterpreted functions associated with
    constructors (``def feats : (ds:Dataset) -> Int = uninterpreted``, the
    auto-generated inductive projections ``Pair_mk_fst``, …). They are
    lowered to the liquid signature shape so they can be applied to a hole's
    argument slots. Only ``UninterpretedBinder`` entries are taken — plain
    first-order program functions are deliberately excluded to keep the
    candidate space bounded.
    """
    if typing_ctx is None:
        return {}
    measures: dict[Name, list[TypeConstructor | TypeVar]] = {}
    for entry in typing_ctx.entries:
        if isinstance(entry, UninterpretedBinder):
            lowered = lower_abstraction_type(entry.type)
            if len(lowered) >= 2:
                measures[entry.name] = lowered
    return measures


def build_measure_atoms(
    ctx: LiquidTypeCheckingContext,
    hole: LiquidHornApplication,
    measures: dict[Name, list[TypeConstructor | TypeVar]],
) -> list[LiquidTerm]:
    """Applications of custom measures to the hole's variable slots.

    A measure such as ``feats : Dataset -> Int`` is not itself a boolean
    predicate, so it can only appear in a refinement nested inside a prelude
    operator (e.g. ``feats(arg_0) <= 0``). These applications become extra
    atoms that ``get_possible_args`` can place in operator argument
    positions. Only single applications over the hole variables are produced;
    deeper nesting is intentionally out of scope.
    """
    atoms: list[LiquidTerm] = []
    seen: set[LiquidTerm] = set()
    for fname, ftype in measures.items():
        arity = len(ftype) - 1
        if arity == 0:
            continue
        for args in get_possible_args(hole.argtypes, arity):
            if not any(isinstance(a, LiquidVar) for a in args):
                continue
            app = LiquidApp(fname, list(args))
            if app in seen:
                continue
            try:
                type_infer_liquid(ctx, app)
            except LiquidTypeCheckException:
                continue
            seen.add(app)
            atoms.append(app)
    return atoms


def build_possible_assignment(
    hole: LiquidHornApplication,
    typing_ctx: TypingContext | None = None,
) -> Generator[LiquidApp]:
    ctx = liquid_typechecking_ctx_for_hole(typing_ctx, hole)
    measures = context_measures(typing_ctx)
    measure_atoms = build_measure_atoms(ctx, hole, measures)
    measure_atom_set = set(measure_atoms)

    # Boolean-returning measures may serve as predicates directly; the prelude
    # supplies the relational operators that consume the non-boolean ones.
    gen_functions: dict[Name, list[TypeConstructor | TypeVar]] = dict(liquid_prelude)
    for fname, ftype in measures.items():
        gen_functions.setdefault(fname, ftype)

    for fname in gen_functions:
        ftype = gen_functions[fname]
        arity = len(ftype) - 1
        for args in get_possible_args(hole.argtypes, arity, measure_atoms):
            # At least one hole variable (directly or via a measure application) must be used.
            if not any(isinstance(a, LiquidVar) or a in measure_atom_set for a in args):
                continue
            arg_list = list(args)
            app = LiquidApp(fname, arg_list)
            if check_liquid(ctx, app, t_bool):
                yield app


def build_initial_assignment(
    c: Constraint,
    typing_ctx: TypingContext | None = None,
    qualifier_atoms: frozenset[LiquidTerm] | None = None,
) -> Assignment:
    atoms = frozenset() if qualifier_atoms is None else qualifier_atoms
    holes = obtain_holes_constraint(c)
    assign: dict[Name, list[LiquidTerm]] = {}
    for h in holes:
        if h.name not in assign:
            prelude = list(build_possible_assignment(h, typing_ctx))
            extra = build_qualifier_candidates(typing_ctx, h, atoms)
            prelude_set = set(prelude)
            assign[h.name] = prelude + [c for c in extra if c not in prelude_set]
    return assign


def merge_assignments(xs: list[LiquidTerm]) -> LiquidTerm:
    b = LiquidLiteralBool(True)
    for c in xs:
        b = mk_liquid_and(b, c)
    return b


def split(c: Constraint) -> list[Constraint]:
    match c:
        case LiquidConstraint(_):
            return [c]
        case Conjunction(c1, c2):
            return split(c1) + split(c2)
        case Implication(name, base, pre, post):
            return [Implication(name, base, pre, cp) for cp in split(post)]
        case UninterpretedFunctionDeclaration(name, type, seq):
            return [UninterpretedFunctionDeclaration(name, type, c) for c in split(seq)]
        case ReflectedFunctionDeclaration(name, type, params, body, seq):
            return [ReflectedFunctionDeclaration(name, type, params, body, c) for c in split(seq)]
        case _:
            assert False


def build_forall_implication(
    vs: list[tuple[Name, Type]],
    p: LiquidTerm,
    c: Constraint,
) -> Constraint:
    if not vs:
        return c
    else:
        for name, _ in vs:
            assert isinstance(name, Name)
    lastEl = vs[-1]
    assert isinstance(lastEl[1], TypeConstructor)
    cf = Implication(lastEl[0], lastEl[1], p, c)
    for n, t in vs[-2::-1]:
        assert isinstance(t, TypeConstructor)
        assert isinstance(n, Name)
        cf = Implication(n, t, LiquidLiteralBool(True), cf)
    return cf


def simpl(vs: list[tuple[Name, Type]], p: LiquidTerm, c: Constraint) -> Constraint:
    if isinstance(c, Implication):
        return simpl(vs + [(c.name, c.base)], mk_liquid_and(p, c.pred), c.seq)
    else:
        return build_forall_implication(vs, p, c)


def flat(c: Constraint) -> list[Constraint]:
    return [simpl([], LiquidLiteralBool(True), cp) for cp in split(c)]


def has_k_head(c: Constraint) -> bool:
    match c:
        case Conjunction(_, _):
            assert False
        case Implication(_, _, _, post):
            return has_k_head(post)
        case LiquidConstraint(e):
            return isinstance(e, LiquidHornApplication)
        case UninterpretedFunctionDeclaration(_, _, post):
            return has_k_head(post)
        case ReflectedFunctionDeclaration(_, _, _, _, post):
            return has_k_head(post)
        case _:
            assert False, f"Unkown constraint type: {c} ({type(c)})"


def apply_constraint(assign: Assignment, c: Constraint) -> Constraint:
    match c:
        case LiquidConstraint(e):
            return LiquidConstraint(apply_liquid(assign, e))
        case Conjunction(c1, c2):
            return Conjunction(
                apply_constraint(assign, c1),
                apply_constraint(assign, c2),
            )
        case Implication(name, base, pre, post):
            return Implication(
                name,
                base,
                apply_liquid(assign, pre),
                apply_constraint(assign, post),
            )
        case UninterpretedFunctionDeclaration(name, base, post):
            return UninterpretedFunctionDeclaration(name, base, apply_constraint(assign, post))
        case ReflectedFunctionDeclaration(name, base, params, body, post):
            return ReflectedFunctionDeclaration(
                name, base, params, apply_liquid(assign, body), apply_constraint(assign, post)
            )
        case _:
            assert False


def fill_horn_arguments(h: LiquidHornApplication, candidate: LiquidTerm) -> LiquidTerm:
    for i, (n, _) in enumerate(h.argtypes):
        assert isinstance(n, LiquidTerm)
        candidate = substitution_in_liquid(candidate, n, mk_arg(i))
    return candidate


def apply_liquid(assign: Assignment, c: LiquidTerm) -> LiquidTerm:
    if isinstance(c, LiquidHornApplication):
        if c.name in assign:
            ne = assign[c.name]
            return fill_horn_arguments(c, merge_assignments(ne))
        else:
            return c
    elif isinstance(c, LiquidApp):
        return LiquidApp(c.fun, [apply_liquid(assign, ci) for ci in c.args])
    else:
        return c


def apply(assign: Assignment, c: Any):
    if isinstance(c, Constraint):
        return apply_constraint(assign, c)
    elif isinstance(c, LiquidTerm):
        return apply_liquid(assign, c)
    assert False


def extract_components_of_imp(
    c: Constraint,
) -> tuple[list[tuple[Name, TypeConstructor | TypeVar | AbstractionType | Top]], tuple[LiquidTerm, LiquidTerm]]:
    match c:
        case UninterpretedFunctionDeclaration(name, base, post):
            (vs1, (p, h)) = extract_components_of_imp(post)
            vsh: list[tuple[Name, TypeConstructor | TypeVar | AbstractionType | Top]] = [(name, base)]
            return (vsh + vs1, (p, h))
        case ReflectedFunctionDeclaration(name, base, _, _, post):
            (vs1, (p, h)) = extract_components_of_imp(post)
            vsh = [(name, base)]
            return (vsh + vs1, (p, h))
        case Implication(name, base, pre, seq):
            (vs1, (p, h)) = extract_components_of_imp(seq)
            vs: list[tuple[Name, TypeConstructor | TypeVar | AbstractionType | Top]] = [(name, base)]
            return (vs + vs1, (mk_liquid_and(pre, p), h))
        case LiquidConstraint(e):
            return ([], (LiquidLiteralBool(True), e))
        case _:
            assert False, f"Unkown context: {c} ({type(c)})"


def weaken(assign, c: Constraint) -> Assignment:
    (vs, (p, h)) = extract_components_of_imp(c)
    assert isinstance(h, LiquidHornApplication)
    assert h.name in assign
    current_rep = assign[h.name]

    def keep(q: LiquidTerm) -> bool:
        assert isinstance(h, LiquidHornApplication)
        qp = fill_horn_arguments(h, q)
        nc = constraint_builder(vs, imp(apply(assign, p), end(qp)))
        return smt_valid(nc)

    qsp = [q for q in current_rep if keep(q)]
    return {k: assign[k] if k != h.name else qsp for k in assign}


def fixpoint(cs: list[Constraint], assign) -> Assignment:
    ncs = [c for c in cs if not smt_valid(apply(assign, c))]
    if not ncs:
        return assign
    else:
        weakened_assignment = weaken(assign, ncs[0])
        return fixpoint(cs, weakened_assignment)


def solve(
    c: Constraint,
    typing_ctx: TypingContext | None = None,
    qualifier_atoms: frozenset[LiquidTerm] | None = None,
) -> bool:
    # Performance improvement
    if not contains_horn_constraint(c):
        return smt_valid(c)
    if qualifier_atoms is None:
        atoms = extract_qualifier_atoms(typing_ctx) if typing_ctx is not None else frozenset()
    else:
        atoms = qualifier_atoms
    cs = flat(c)
    csk = [c for c in cs if has_k_head(c)]
    csp = [c for c in cs if not has_k_head(c)]
    assignment0: Assignment = build_initial_assignment(c, typing_ctx, atoms)
    subst = fixpoint(csk, assignment0)

    def merge(acc: Constraint, pi: Constraint) -> Constraint:
        return Conjunction(acc, pi)

    seed: Constraint = LiquidConstraint(LiquidLiteralBool(True))
    merged_csps = reduce(merge, csp, seed)
    return smt_valid(apply(subst, merged_csps))
