package prolog.builtins

import prolog.Answers
import prolog.Substitutions
import prolog.ast.arithmetic.Integer
import prolog.ast.terms.Atom
import prolog.ast.terms.Goal
import prolog.ast.terms.Structure
import prolog.ast.terms.Term
import prolog.flags.AppliedShift
import prolog.logic.applySubstitution
import prolog.logic.unifyLazy

/**
 * These classes provide support for delimited continuations in Prolog.
 * More specifically, they implement the reset/3 and shift/1 operators.
 *
 * See also [SWI-Prolog 4.9 Delimited Continuations](https://www.swi-prolog.org/pldoc/man?section=delcont)
 */

/**
 * Call [Goal]. If Goal calls [Shift], and the arguments of Shift can be unified with Ball, Shift causes Reset to
 * return, unifying Cont with a Goal that represents the continuation after shift.
 */
class Reset(private val goal: Goal, private val ball: Term, private val cont: Term) :
    Structure("reset", goal, ball, cont) {
    override fun satisfy(subs: Substitutions): Answers = sequence {
        goal.satisfy(subs).forEach { goalResult ->
            goalResult.fold(
                // If Goal succeeds, then reset/3 also succeeds and binds Cont and Term1 to 0.
                onSuccess = { goalSubs ->
                    unifyLazy(cont, Integer(0), goalSubs).forEach { contResult ->
                        contResult.map { contSubs ->
                            yield(Result.success(goalSubs + contSubs))
                        }
                    }
                },
                onFailure = { failure ->
                    // If at some point Goal calls shift(Term2), then its further execution is suspended.
                    if (failure is AppliedShift) {
                        require(failure.cont != null) { "Shift must have a continuation" }
                        // Reset/3 succeeds immediately, binding Term1 to Term2 and Cont to the remainder of Goal.
                        val shiftSubs = failure.subs
                        val appliedBall = applySubstitution(failure.ball, shiftSubs)
                        val appliedCont = applySubstitution(failure.cont, shiftSubs)
                        Conjunction(Unify(ball, appliedBall), Unify(appliedCont, cont))
                            .satisfy(shiftSubs)
                            .forEach { conResult ->
                                conResult.map { conSubs ->
                                    yield(Result.success(shiftSubs + conSubs))
                                }
                            }
                    } else {
                        // If Goal fails, then reset also fails.
                        yield(Result.failure(failure))
                    }
                }
            )
        }
    }
}

/**
 * Variables that have been bound during the procedure called by reset/3 stay bound after a shift/1:
 */
class Shift(private val ball: Term) : Structure("shift", ball) {
    override fun satisfy(subs: Substitutions): Answers = sequence {
        val shift = AppliedShift(
            subs = subs,
            ball = ball,
            cont = null
        )
        yield(Result.failure(shift))
    }
}