package prolog.builtins

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

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

    @Test
    fun `rule with cut as body`() {
        Program.load(
            listOf(
                Rule(Atom("foo"), Cut()),
                Fact(Atom("foo"))
            )
        )

        val goal = Atom("foo")

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

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        assertTrue(result[0].getOrNull()!!.isEmpty(), "Expected empty substitutions")
    }

    // See also: https://stackoverflow.com/a/23292126

    @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.reset()

        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.reset()

        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.reset()

        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_equal cut test`() {
        Program.load(
            listOf(
                Rule(
                    CompoundTerm(Atom("not_equal"), listOf(Variable("X"), Variable("Y"))),
                    Conjunction(
                        Unify(Variable("X"), Variable("Y")),
                        Conjunction(Cut(), Fail)
                    )
                ),
                Fact(CompoundTerm(Atom("not_equal"), listOf(Variable("_1"), Variable("_2"))))
            )
        )

        var goal = CompoundTerm(Atom("not_equal"), listOf(Integer(1), Integer(1)))
        var result = Program.query(goal).toList()

        assertTrue(result.none(), "Expected no results")

        goal = CompoundTerm(Atom("not_equal"), listOf(Integer(1), Integer(2)))
        result = Program.query(goal).toList()

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

    @Test
    fun `cut with disjunction`() {
        Program.load(
            listOf(
                Rule(
                    CompoundTerm(Atom("choice"), listOf(Variable("X"))),
                    Conjunction(
                        Unify(Variable("X"), Integer(1)),
                        Disjunction(
                            Cut(),
                            Unify(Variable("X"), Integer(2))
                        )
                    )
                )
            )
        )

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

        val 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")
    }
}