package prolog.builtins

import org.junit.jupiter.api.Assertions.*
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import prolog.Program
import prolog.ast.logic.Fact
import prolog.ast.logic.Rule
import prolog.ast.terms.Atom
import prolog.ast.terms.CompoundTerm
import prolog.ast.terms.Integer
import prolog.ast.terms.Variable

class ControlOperatorsTests {
    @BeforeEach
    fun setUp() {
        Program.clear()
    }

    @Test
    fun `simple cut program`() {
        // First an example without cut
        Program.load(
            listOf(
                Fact(CompoundTerm(Atom("foo"), listOf(Atom("a")))),
                Fact(CompoundTerm(Atom("foo"), listOf(Atom("b")))),
            )
        )

        val goal = CompoundTerm(Atom("foo"), listOf(Variable("X")))

        var result = Program.query(goal).toList()

        assertEquals(2, result.size, "Expected 2 results")

        // Now with cut

        Program.clear()

        Program.load(
            listOf(
                Rule(CompoundTerm(Atom("foo"), listOf(Atom("a"))), Cut()),
                Fact(CompoundTerm(Atom("foo"), listOf(Atom("b")))),
            )
        )

        result = Program.query(goal).toList()

        assertEquals(1, result.size, "Expected 1 result")
    }

    @Test
    fun `cut with pruning`() {
        // First an example without cut
        Program.load(
            listOf(
                Fact(CompoundTerm(Atom("a"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("a"), listOf(Integer(2)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(2)))),

                Rule(
                    CompoundTerm(Atom("foo"), listOf(Variable("X"), Variable("Y"))),
                    Conjunction(
                        CompoundTerm(Atom("a"), listOf(Variable("X"))),
                        CompoundTerm(Atom("b"), listOf(Variable("Y")))
                    )
                )
            )
        )

        val goal = CompoundTerm(Atom("foo"), listOf(Variable("X"), Variable("Y")))

        /* ?- foo(X, Y).
         * X = Y, Y = 1 ;
         * X = 1,
         * Y = 2 ;
         * X = 2,
         * Y = 1 ;
         * X = Y, Y = 2. */
        var result = Program.query(goal).toList()

        assertEquals(4, result.size, "Expected 4 results")

        // Now with cut in the middle

        Program.clear()

        Program.load(
            listOf(
                Fact(CompoundTerm(Atom("a"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("a"), listOf(Integer(2)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(2)))),

                Rule(
                    CompoundTerm(Atom("foo"), listOf(Variable("X"), Variable("Y"))),
                    Conjunction(
                        CompoundTerm(Atom("a"), listOf(Variable("X"))),
                        Conjunction(
                            Cut(),
                            CompoundTerm(Atom("b"), listOf(Variable("Y")))
                        )
                    )
                )
            )
        )

        /*
        ?- foo(X, Y).
        X = Y, Y = 1 ;
        X = 1,
        Y = 2.
         */
        result = Program.query(goal).toList()

        assertEquals(2, result.size, "Expected 2 results")

        // Now with cut at the end

        Program.clear()

        Program.load(
            listOf(
                Fact(CompoundTerm(Atom("a"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("a"), listOf(Integer(2)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(1)))),
                Fact(CompoundTerm(Atom("b"), listOf(Integer(2)))),

                Rule(
                    CompoundTerm(Atom("foo"), listOf(Variable("X"), Variable("Y"))),
                    Conjunction(
                        CompoundTerm(Atom("a"), listOf(Variable("X"))),
                        Conjunction(
                            CompoundTerm(Atom("b"), listOf(Variable("Y"))),
                            Cut()
                        )
                    )
                )
            )
        )

        /*
         ?- foo(X, Y).
         X = Y, Y = 1 ;
         */
        result = Program.query(goal).toList()

        assertEquals(1, result.size, "Expected 1 result")
    }

    @Test
    fun not_atom() {
        val success = Fact(Atom("a"))

        Program.load(listOf(success))

        val goal = Atom("a")
        val notGoal = Not(goal)

        val result1 = Program.query(goal)
        val result2 = Program.query(notGoal)

        assertTrue(result1.any(), "Expected query to succeed")
        assertFalse(result2.any(), "Expected query to fail")
    }

    @Test
    fun not_compound() {
        val success = Fact(Atom("f"))
        val failure = Fact(Atom("g"))

        Program.load(listOf(success, failure))

        val goal = Atom("f")
        val notGoal = Not(Atom("g"))

        val result1 = Program.query(goal)
        val result2 = Program.query(notGoal)

        assertTrue(result1.any(), "Expected query to succeed")
        assertFalse(result2.any(), "Expected query to fail")
    }

    @Test
    fun fail_should_cause_fails() {
        val success = Fact(Atom("a"))
        val failure = Fact(Atom("b"))

        Program.load(listOf(success, failure))

        val goal = Atom("a")
        val failGoal = Fail

        val result1 = Program.query(goal)
        val result2 = Program.query(failGoal)

        assertTrue(result1.any(), "Expected query to succeed")
        assertFalse(result2.any(), "Expected query to fail")
    }
}