"""Termination metrics (Liquid Haskell-style) for recursive definitions."""

from __future__ import annotations

from aeon.core.liquid import LiquidApp, LiquidLiteralBool, LiquidLiteralInt, LiquidTerm, LiquidVar, liquid_free_vars
from aeon.core.liquid_ops import mk_liquid_and, mk_liquid_or
from aeon.core.substitutions import inline_lets, liquefy, substitution, substitution_in_liquid
from aeon.core.terms import (
    Abstraction,
    Annotation,
    Application,
    If,
    Let,
    Rec,
    RefinementAbstraction,
    RefinementApplication,
    Term,
    TypeAbstraction,
    TypeApplication,
    Var,
)
from aeon.core.types import AbstractionType, RefinedType, Type
from aeon.typechecking.context import TypingContext
from aeon.utils.location import Location
from aeon.utils.name import Name
from aeon.verification.vcs import Conjunction, Constraint, LiquidConstraint

ctrue = LiquidConstraint(LiquidLiteralBool(True))


def peel_abstractions(t: Term) -> tuple[list[Name], Term]:
    names: list[Name] = []
    cur = t
    while isinstance(cur, Abstraction):
        names.append(cur.var_name)
        cur = cur.body
    return names, cur


def peel_type_formal_names(ty: Type) -> list[Name]:
    names: list[Name] = []
    cur: Type = ty
    while isinstance(cur, AbstractionType):
        names.append(cur.var_name)
        cur = cur.type
    return names


def _opened_refinement_liquid(
    var_type: Type, binder_name: Name, term_formals: list[Name], type_formals: list[Name]
) -> LiquidTerm | None:
    match var_type:
        case RefinedType(bname, _bty, ref):
            opened = substitution_in_liquid(ref, LiquidVar(binder_name), bname)
            return align_liquid_to_type_formals(opened, term_formals, type_formals)
        case _:
            return None


def entry_refinement_liquids(
    fn_arrow_type: Type, term_formals: list[Name], type_formals: list[Name]
) -> list[LiquidTerm]:
    """Opened refinement predicates for each curried ``RefinedType`` parameter (entry guards)."""
    out: list[LiquidTerm] = []
    cur: Type = fn_arrow_type
    while isinstance(cur, AbstractionType):
        liq = _opened_refinement_liquid(cur.var_type, cur.var_name, term_formals, type_formals)
        if liq is not None:
            out.append(liq)
        cur = cur.type
    return out


def align_liquid_to_type_formals(liq: LiquidTerm, term_names: list[Name], type_names: list[Name]) -> LiquidTerm:
    """VCs use ``AbstractionType`` binders; liquefy may use distinct ``Name`` ids from lambdas."""
    out = liq
    for t_n, ty_n in zip(term_names, type_names, strict=True):
        if t_n != ty_n:
            out = substitution_in_liquid(out, LiquidVar(ty_n), t_n)
    return out


def peel_application_chain(t: Term) -> tuple[Term, list[Term]]:
    """Left-associate application: ((f x) y) -> f, [x, y]."""
    args: list[Term] = []
    cur = t
    while isinstance(cur, Application):
        args.append(cur.arg)
        cur = cur.fun
    args.reverse()
    return cur, args


def _term_bool_not(cond: Term) -> Term:
    """Surface ``!cond`` as a core application (see parser ``nnot``)."""
    return Application(Var(Name("!", 0), loc=cond.loc), cond, loc=cond.loc)


def _synth_type(ctx: TypingContext, t: Term) -> Type | None:
    """Lazy import: ``typeinfer`` imports this module at load time."""
    try:
        from aeon.typechecking.typeinfer import synth

        return synth(ctx, t)[1]
    except Exception:
        return None


def collect_recursive_calls_with_paths(
    fn: Name,
    arity: int,
    t: Term,
    typing_ctx: TypingContext | None,
    inner_expect_ty: Type | None,
    term_formals: list[Name],
    type_formals: list[Name],
) -> list[tuple[list[Term], Location | None, tuple[Term, ...], tuple[LiquidTerm, ...]]]:
    """Self-calls with ``arity`` args, ``If`` path guards, and nested binder refinements.

    The walk traverses the *original* body — ``let`` bindings are not inlined
    away beforehand. Each ``let`` contributes:

    * a substitution (binder → inlined value) applied to recorded call
      arguments and path guards, so obligations stay phrased over the
      function's formals (what the former global ``inline_lets`` achieved);
    * a path-scoped hypothesis: the binder's synthesised refinement, opened at
      the liquefied value (issue #378). This is what lets an instantiated
      axiom (``b := fdiv_bounds n 10``) or a refined native's return type
      guard the termination obligation of a recursive call under it. Facts are
      attached at the let's position in the walk, never hoisted, so a fact
      established under one branch cannot leak into another.
    """
    found: list[tuple[list[Term], Location | None, tuple[Term, ...], tuple[LiquidTerm, ...]]] = []
    base_vars = set(term_formals) | {fn}
    if typing_ctx is not None:
        base_vars |= {n for n, _ in typing_ctx.vars()}

    def apply_env(term: Term, env: tuple[tuple[Name, Term], ...]) -> Term:
        for n, v in reversed(env):
            term = substitution(term, v, n)
        return inline_lets(term)

    def let_fact(name: Name, val: Term, ty_v: Type | None, env: tuple[tuple[Name, Term], ...]) -> LiquidTerm | None:
        """The let binder's refinement, opened at the binder's (inlined) value."""
        match ty_v:
            case RefinedType(bname, _, ref):
                if any(v == bname for v in liquid_free_vars(ref)):
                    lv = liquefy(apply_env(val, env))
                    if lv is None:
                        return None
                    ref = substitution_in_liquid(ref, lv, bname)
                if any(v not in base_vars for v in liquid_free_vars(ref)):
                    return None
                return align_liquid_to_type_formals(ref, term_formals, type_formals)
            case _:
                return None

    def walk(
        tt: Term,
        path: tuple[Term, ...],
        nested_refs: tuple[LiquidTerm, ...],
        ctx: TypingContext | None,
        expect_ty: Type | None,
        env: tuple[tuple[Name, Term], ...],
    ) -> None:
        match tt:
            case Application(_, _, _):
                head, args = peel_application_chain(tt)
                if isinstance(head, Var) and head.name == fn and len(args) == arity:
                    found.append(([apply_env(a, env) for a in args], tt.loc, path, nested_refs))
            case _:
                pass
        match tt:
            case Application(Abstraction(bname, abody, _), arg, _):
                # Beta redex: bind like a let so recorded terms stay closed.
                walk(arg, path, nested_refs, ctx, None, env)
                walk(abody, path, nested_refs, ctx, expect_ty, env + ((bname, apply_env(arg, env)),))
            case Application(fun, arg, _):
                f_ty = _synth_type(ctx, fun) if ctx is not None else None
                walk(fun, path, nested_refs, ctx, f_ty, env)
                arg_ty: Type | None = None
                match f_ty:
                    case AbstractionType(_, vt, _):
                        arg_ty = vt
                walk(arg, path, nested_refs, ctx, arg_ty, env)
            case Abstraction(name, body, _):
                if isinstance(expect_ty, AbstractionType):
                    vty = expect_ty.var_type
                    new_ctx = ctx.with_var(name, vty) if ctx is not None else None
                    ref_l = _opened_refinement_liquid(vty, name, term_formals, type_formals)
                    nrefs = nested_refs + (ref_l,) if ref_l is not None else nested_refs
                    walk(body, path, nrefs, new_ctx, expect_ty.type, env)
                else:
                    walk(body, path, nested_refs, ctx, None, env)
            case Let(name, val, body, _):
                walk(val, path, nested_refs, ctx, None, env)
                ty_v = _synth_type(ctx, val) if ctx is not None else None
                fact = let_fact(name, val, ty_v, env)
                nrefs = nested_refs + (fact,) if fact is not None else nested_refs
                new_ctx = ctx.with_var(name, ty_v) if ctx is not None and ty_v is not None else ctx
                walk(body, path, nrefs, new_ctx, expect_ty, env + ((name, apply_env(val, env)),))
            case Rec(name, _, val, body, _, _):
                walk(val, path, nested_refs, ctx, None, env)
                walk(body, path, nested_refs, ctx, expect_ty, env + ((name, apply_env(val, env)),))
            case Annotation(expr, _, _):
                walk(expr, path, nested_refs, ctx, expect_ty, env)
            case If(cond, then, otherwise, _):
                cond_s = apply_env(cond, env)
                walk(cond, path, nested_refs, ctx, None, env)
                walk(then, path + (cond_s,), nested_refs, ctx, expect_ty, env)
                walk(otherwise, path + (_term_bool_not(cond_s),), nested_refs, ctx, expect_ty, env)
            case TypeApplication(body, _, _):
                walk(body, path, nested_refs, ctx, expect_ty, env)
            case TypeAbstraction(_, _, body, _):
                walk(body, path, nested_refs, ctx, expect_ty, env)
            case RefinementApplication(body, refinement, _):
                walk(body, path, nested_refs, ctx, expect_ty, env)
                walk(refinement, path, nested_refs, ctx, None, env)
            case RefinementAbstraction(_, _, body, _):
                walk(body, path, nested_refs, ctx, expect_ty, env)
            case _:
                pass

    walk(t, (), (), typing_ctx, inner_expect_ty, ())
    return found


def _inner_ctx_and_expect_ty(nrctx: TypingContext, var_type: Type, formals: list[Name]) -> tuple[TypingContext, Type]:
    ctx = nrctx
    cur_ty = var_type
    for nm in formals:
        assert isinstance(cur_ty, AbstractionType)
        ctx = ctx.with_var(nm, cur_ty.var_type)
        cur_ty = cur_ty.type
    return ctx, cur_ty


def _full_path_guard_liquid(
    entry_refs: list[LiquidTerm],
    nested_refs: tuple[LiquidTerm, ...],
    path: tuple[Term, ...],
    formals: list[Name],
    type_formals: list[Name],
) -> LiquidTerm | None:
    parts: list[LiquidTerm] = list(entry_refs) + list(nested_refs)
    if path:
        pl = _liquefy_path_conjuncts(path, formals, type_formals)
        if pl is None:
            return None
        if pl != LiquidLiteralBool(True):
            parts.append(pl)
    if not parts:
        return LiquidLiteralBool(True)
    acc = parts[0]
    for p in parts[1:]:
        acc = mk_liquid_and(acc, p)
    return acc


def _liquefy_path_conjuncts(path: tuple[Term, ...], formals: list[Name], type_formals: list[Name]) -> LiquidTerm | None:
    if not path:
        return LiquidLiteralBool(True)
    parts: list[LiquidTerm] = []
    for g in path:
        liq = liquefy(inline_lets(g))
        if liq is None:
            return None
        parts.append(align_liquid_to_type_formals(liq, formals, type_formals))
    acc = parts[0]
    for p in parts[1:]:
        acc = mk_liquid_and(acc, p)
    return acc


def _metric_call_non_negative(call_ms: list[LiquidTerm]) -> LiquidTerm:
    acc = LiquidApp(Name(">=", 0), [call_ms[0], LiquidLiteralInt(0)])
    for m in call_ms[1:]:
        acc = mk_liquid_and(acc, LiquidApp(Name(">=", 0), [m, LiquidLiteralInt(0)]))
    return acc


def _termination_obligation(
    path: tuple[Term, ...],
    lex: LiquidTerm,
    call_ms: list[LiquidTerm],
    formals: list[Name],
    type_formals: list[Name],
    entry_refs: list[LiquidTerm],
    nested_refs: tuple[LiquidTerm, ...],
) -> Constraint:
    r"""``(entry /\ nested /\ branch) => (lex and metric(call) >= 0)``.

    Non-negativity is **always** required: a well-founded measure must strictly
    decrease *and* be bounded below. An unbounded integer metric descends
    forever (``f (n - 1)`` never terminates) yet would satisfy a decrease-only
    obligation, making the inductive hypothesis used for the recursive call
    unsound. See ``tests/recursion_soundness_test.py``.
    """
    oblig = mk_liquid_and(lex, _metric_call_non_negative(call_ms))
    full = _full_path_guard_liquid(entry_refs, nested_refs, path, formals, type_formals)
    if full is None:
        return LiquidConstraint(LiquidLiteralBool(False))
    if full == LiquidLiteralBool(True):
        return LiquidConstraint(oblig)
    return LiquidConstraint(LiquidApp(Name("-->", 0), [full, oblig]))


def substitute_formals(template: Term, formals: list[Name], args: list[Term]) -> Term:
    out = template
    for formal, arg in zip(formals, args, strict=True):
        out = substitution(out, arg, formal)
    return out


def _liquefy_metric_at(
    metric: Term,
    formals: list[Name],
    type_formals: list[Name],
    call_args: list[Term] | None,
) -> LiquidTerm | None:
    t = inline_lets(metric)
    if call_args is not None:
        t = substitute_formals(t, formals, [inline_lets(a) for a in call_args])
        t = inline_lets(t)
    liq = liquefy(t)
    if liq is None:
        return None
    return align_liquid_to_type_formals(liq, formals, type_formals)


def _liquefy_callee_metric_at(
    callee_metric: Term,
    callee_formals: list[Name],
    call_args: list[Term],
    caller_formals: list[Name],
    caller_type_formals: list[Name],
) -> LiquidTerm | None:
    """The callee's metric evaluated at a call's arguments, expressed over the
    *caller's* type formals.

    ``call_args`` are written in terms of the caller's value parameters, so after
    substituting them for the callee's formals the residual free variables are the
    caller's value formals — align those to the caller's type formals (the names
    the surrounding implication binds in the SMT context). For a self-call,
    ``callee_*`` and ``caller_*`` coincide and this matches the single-function
    behaviour."""
    t = inline_lets(callee_metric)
    t = substitute_formals(t, callee_formals, [inline_lets(a) for a in call_args])
    t = inline_lets(t)
    liq = liquefy(t)
    if liq is None:
        return None
    return align_liquid_to_type_formals(liq, caller_formals, caller_type_formals)


def _fold_or(terms: list[LiquidTerm]) -> LiquidTerm:
    acc = terms[0]
    for t in terms[1:]:
        acc = mk_liquid_or(acc, t)
    return acc


def _lexicographic_less(call_ms: list[LiquidTerm], entry_ms: list[LiquidTerm]) -> LiquidTerm:
    """Strict lexicographic decrease (Liquid Haskell ``/ [m0, m1, …]``)."""
    assert len(call_ms) == len(entry_ms) and len(call_ms) >= 1
    k = len(call_ms)
    disjuncts: list[LiquidTerm] = []
    for j in range(k):
        step_lt = LiquidApp(Name("<", 0), [call_ms[j], entry_ms[j]])
        if j == 0:
            disjuncts.append(step_lt)
        else:
            prefix = LiquidApp(Name("==", 0), [call_ms[0], entry_ms[0]])
            for i in range(1, j):
                prefix = mk_liquid_and(prefix, LiquidApp(Name("==", 0), [call_ms[i], entry_ms[i]]))
            disjuncts.append(mk_liquid_and(prefix, step_lt))
    return _fold_or(disjuncts)


def termination_metric_constraints(rec: Rec, typing_ctx: TypingContext | None = None) -> Constraint:
    r"""Lexicographic ``m(call) <* m(entry)`` (Liquid Haskell ``/ [m0, m1, …]``).

    With ``typing_ctx`` (the context used to check ``rec.var_value``), also accumulate
    refinements from **nested** abstractions—typically pattern binders after ``match`` lowers
    to an eliminator—so constructor argument refinements guard termination obligations.

    For a member of a ``mutual`` group, calls to *siblings* are recursive too: the
    obligation for a call ``g a`` is ``g_metric(a) <* f_metric(entry)`` (and
    ``g_metric(a) >= 0``), using the *callee's* metric at the call arguments and
    the *caller's* metric at entry. This makes the whole group descend a single
    well-founded lexicographic measure, so no infinite call chain exists.
    """
    if not rec.decreasing_by:
        return ctrue
    caller_metrics = list(rec.decreasing_by)
    formals, inner = peel_abstractions(rec.var_value)
    type_formals = peel_type_formal_names(rec.var_type)
    if len(type_formals) != len(formals):
        return LiquidConstraint(LiquidLiteralBool(False))
    entry_refs = entry_refinement_liquids(rec.var_type, formals, type_formals)
    # Let-bindings are NOT inlined away here: the walk substitutes them into
    # recorded call arguments and path guards itself, and additionally turns
    # each binder's refinement into a path-scoped hypothesis (issue #378).
    inner_ctx: TypingContext | None = None
    inner_expect: Type | None = None
    if typing_ctx is not None:
        inner_ctx, inner_expect = _inner_ctx_and_expect_ty(typing_ctx, rec.var_type, formals)

    # Caller's entry measure (evaluated at its own parameters).
    caller_entry_ms: list[LiquidTerm] = []
    for metric in caller_metrics:
        m_entry = _liquefy_metric_at(metric, formals, type_formals, None)
        if m_entry is None:
            return LiquidConstraint(LiquidLiteralBool(False))
        caller_entry_ms.append(m_entry)

    # Targets whose calls must decrease the caller's measure: the function itself
    # plus each mutual-group sibling, each with its own metric/formals.
    targets: list[tuple[Name, list[Term], list[Name], list[Name]]] = [
        (rec.var_name, caller_metrics, formals, type_formals)
    ]
    for comp in rec.companions:
        comp_type_formals = peel_type_formal_names(comp.type)
        targets.append((comp.name, list(comp.decreasing_by), list(comp.formals), comp_type_formals))

    parts: list[Constraint] = []
    for tname, tmetrics, tformals, ttype_formals in targets:
        tarity = len(tformals)
        # A sibling without a metric (or with a differing measure length) cannot
        # be shown to descend the shared lexicographic order.
        if not tmetrics or len(tmetrics) != len(caller_metrics) or len(ttype_formals) != tarity:
            calls = collect_recursive_calls_with_paths(
                tname, tarity, inner, inner_ctx, inner_expect, formals, type_formals
            )
            if calls:
                parts.append(LiquidConstraint(LiquidLiteralBool(False)))
            continue
        calls = collect_recursive_calls_with_paths(tname, tarity, inner, inner_ctx, inner_expect, formals, type_formals)
        for call_args, _loc, path, nested_refs in calls:
            if len(call_args) != tarity:
                parts.append(LiquidConstraint(LiquidLiteralBool(False)))
                continue
            call_ms: list[LiquidTerm] = []
            bad = False
            for metric in tmetrics:
                m_call = _liquefy_callee_metric_at(metric, tformals, call_args, formals, type_formals)
                if m_call is None:
                    bad = True
                    break
                call_ms.append(m_call)
            if bad:
                parts.append(LiquidConstraint(LiquidLiteralBool(False)))
                continue
            lex = _lexicographic_less(call_ms, caller_entry_ms)
            parts.append(_termination_obligation(path, lex, call_ms, formals, type_formals, entry_refs, nested_refs))

    if not parts:
        return ctrue
    acc: Constraint = parts[0]
    for p in parts[1:]:
        acc = Conjunction(acc, p)
    return acc
