package prolog.builtins

import prolog.Answers
import prolog.Substitutions
import prolog.ast.Database.Program
import prolog.ast.logic.LogicOperand
import prolog.ast.logic.LogicOperator
import prolog.ast.terms.Atom
import prolog.ast.terms.Body
import prolog.ast.terms.Goal
import prolog.flags.AppliedCut
import prolog.flags.AppliedShift
import prolog.logic.applySubstitution
import prolog.logic.numbervars

/**
 * Always fail.
 */
object Fail : Atom("fail"), Body {
    override fun satisfy(subs: Substitutions): Answers = emptySequence()
    override fun applySubstitution(subs: Substitutions): Fail = this
}

/**
 * Same as fail, but the name has a more declarative connotation.
 */
typealias False = Fail

/**
 * Always succeed.
 */
object True : Atom("true"), Body {
    override fun satisfy(subs: Substitutions): Answers = sequenceOf(Result.success(emptyMap()))
    override fun applySubstitution(subs: Substitutions): True = this
}

// TODO Repeat/0

class Cut() : Atom("!") {
    override fun satisfy(subs: Substitutions): Answers {
        return sequenceOf(Result.failure(AppliedCut(emptyMap())))
    }

    override fun applySubstitution(subs: Substitutions): Cut = this
}

/**
 * Conjunction (and). True if both Goal1 and Goal2 are true.
 */
open class Conjunction(val left: LogicOperand, private val right: LogicOperand) :
    LogicOperator(Atom(","), left, right) {
    override fun satisfy(subs: Substitutions): Answers = sequence {
        fun satisfyRight(leftSubs: Substitutions): Answers = sequence {
            right.satisfy(subs + leftSubs).forEach { right ->
                right.fold(
                    // If the right part succeeds, yield the result with the left substitutions
                    onSuccess = { rightSubs ->
                        yield(Result.success(leftSubs + rightSubs))
                    },
                    onFailure = { exception ->
                        // If the right part fails, check if it's a cut
                        if (exception is AppliedCut) {
                            // If it's a cut, yield the result with the left substitutions
                            val newSubs = if (exception.subs != null) leftSubs + exception.subs else null
                            yield(Result.failure(AppliedCut(newSubs)))
                            return@sequence
                        }

                        // If it's not a cut, yield the failure
                        yield(Result.failure(exception))
                    }
                )
            }
        }

        fun findNextCutSolution(appliedCut: AppliedCut): Answers = sequence {
            val leftSubs = appliedCut.subs
            right.satisfy(subs + (appliedCut.subs!!)).firstOrNull()?.map { rightSubs ->
                // If the right part succeeds, yield the result with the left substitutions
                yield(Result.success(leftSubs + rightSubs))
                return@sequence
            } ?: yield(Result.failure(AppliedCut()))
        }

        // Satisfy the left part first, which either succeeds or fails
        left.satisfy(subs).forEach { left ->
            left.fold(
                // If it succeeds, satisfy the right part with the updated substitutions and return all results
                onSuccess = { leftSubs -> yieldAll(satisfyRight(leftSubs)) },
                // If it fails, check these conditions:
                onFailure = { exception ->
                    when (exception) {
                        // 1. If the left part is a cut, satisfy the right part ONCE, and stop searching for more solutions
                        is AppliedCut -> yieldAll(findNextCutSolution(exception))

                        // 2. If the left part is a shift, we need to check if the right part can be satisfied
                        is AppliedShift -> {
                            if (exception.cont == null) {
                                // Pass the right of our disjunction as the continuation
                                val newShift = AppliedShift(exception.subs, exception.ball, right as Goal)
                                yield(Result.failure(newShift))
                            } else {
                                // Satisfy the right part with the updated substitutions
                                yieldAll(satisfyRight(exception.subs))
                            }
                        }

                        // 3. Any other failure should be returned as is
                        else -> yield(Result.failure(exception))
                    }
                }
            )
        }
    }

    override fun applySubstitution(subs: Substitutions): Conjunction = Conjunction(
        applySubstitution(left, subs) as LogicOperand,
        applySubstitution(right, subs) as LogicOperand
    )
}

/**
 * Disjunction (or). True if either Goal1 or Goal2 succeeds.
 */
open class Disjunction(private val left: LogicOperand, private val right: LogicOperand) :
    LogicOperator(Atom(";"), left, right) {
    override fun satisfy(subs: Substitutions): Answers = sequence {
        left.satisfy(subs).forEach { left ->
            left.fold(
                onSuccess = { leftSubs ->
                    yield(Result.success(leftSubs))
                },
                onFailure = { failure ->
                    if (failure is AppliedCut) {
                        val leftSubs = failure.subs
                        yield(Result.failure(AppliedCut(leftSubs)))
                        return@sequence
                    }
                }
            )
        }
        yieldAll(right.satisfy(subs))
    }

    override fun applySubstitution(subs: Substitutions): Disjunction = Disjunction(
        applySubstitution(left, subs) as LogicOperand,
        applySubstitution(right, subs) as LogicOperand
    )
}

@Deprecated("Use Disjunction instead")
class Bar(leftOperand: LogicOperand, rightOperand: LogicOperand) : Disjunction(leftOperand, rightOperand)

// TODO ->

// TODO *->

/**
 * True if 'Goal' cannot be proven.
 */
class Not(private val goal: Goal) : LogicOperator(Atom("\\+"), rightOperand = goal) {
    override fun satisfy(subs: Substitutions): Answers {
        // If the goal can be proven, return an empty sequence
        val goalResults = goal.satisfy(subs).iterator()
        if (goalResults.hasNext()) {
            return emptySequence()
        }
        // If the goal cannot be proven, return a sequence with an empty map
        return sequenceOf(Result.success(emptyMap()))
    }

    override fun applySubstitution(subs: Substitutions): Not = Not(applySubstitution(goal, subs) as Goal)
}