package prolog

import org.junit.jupiter.api.Assertions.*
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Nested
import org.junit.jupiter.api.Test
import prolog.ast.logic.Fact
import prolog.ast.logic.Rule
import prolog.builtins.Conjunction
import prolog.builtins.Disjunction
import prolog.builtins.Query
import prolog.logic.equivalent
import prolog.ast.terms.Atom
import prolog.ast.terms.Structure
import prolog.ast.terms.Variable
import prolog.ast.Database.Program
import prolog.ast.arithmetic.Integer
import prolog.ast.terms.AnonymousVariable

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

    @Test
    fun successful_single_fact_query() {
        val fact = Fact(Atom("a"))
        Program.load(listOf(fact))

        val result = Program.query(Atom("a"))

        assertTrue(result.any())
    }

    @Test
    fun failing_single_fact_query() {
        val fact = Fact(Atom("a"))
        Program.load(listOf(fact))

        val result = Program.query(Atom("b"))

        assertFalse(result.any())
    }

    @Test
    fun multiple_fact_query() {
        val fact1 = Fact(Atom("a"))
        val fact2 = Fact(Atom("b"))
        Program.load(listOf(fact1, fact2))

        assertTrue(Program.query(Atom("a")).any())
        assertTrue(Program.query(Atom("b")).any())
        assertFalse(Program.query(Atom("c")).any())
    }

    @Test
    fun single_compound_query() {
        val structure = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        Program.load(listOf(Fact(structure)))

        assertTrue(Program.query(Structure(Atom("f"), listOf(Atom("a"), Atom("b")))).any())
    }

    @Test
    fun multiple_compound_query() {
        val structure1 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        val structure2 = Structure(Atom("g"), listOf(Atom("c"), Atom("d")))
        Program.load(listOf(Fact(structure1), Fact(structure2)))

        assertTrue(Program.query(Structure(Atom("f"), listOf(Atom("a"), Atom("b")))).any())
        assertTrue(Program.query(Structure(Atom("g"), listOf(Atom("c"), Atom("d")))).any())
        assertFalse(Program.query(Structure(Atom("h"), listOf(Atom("e")))).any())
    }

    /**
     * John and Jane are Jimmy's parents.
     */
    @Test
    fun parents_prolog_way() {
        val father = Fact(Structure(Atom("father"), listOf(Atom("john"), Atom("jimmy"))))
        val mother = Fact(Structure(Atom("mother"), listOf(Atom("jane"), Atom("jimmy"))))

        val parent1 = Rule(
            Structure(Atom("parent"), listOf(Variable("X"), Variable("Y"))),
            /* :- */
            Structure(Atom("father"), listOf(Variable("X"), Variable("Y")))
        )
        val parent2 = Rule(
            Structure(Atom("parent"), listOf(Variable("X"), Variable("Y"))),
            /* :- */
            Structure(Atom("mother"), listOf(Variable("X"), Variable("Y")))
        )

        Program.load(listOf(father, mother, parent1, parent2))

        assertTrue(Program.query(Structure(Atom("parent"), listOf(Atom("john"), Atom("jimmy")))).any())
        assertTrue(Program.query(Structure(Atom("parent"), listOf(Atom("jane"), Atom("jimmy")))).any())
    }

    /**
     * John and Jane are Jimmy's parents.
     */
    @Test
    fun parents_disjunction_way() {
        val father = Fact(Structure(Atom("father"), listOf(Atom("john"), Atom("jimmy"))))
        val mother = Fact(Structure(Atom("mother"), listOf(Atom("jane"), Atom("jimmy"))))

        val variable1 = Variable("X")
        val variable2 = Variable("Y")

        val parent = Rule(
            Structure(Atom("parent"), listOf(variable1, variable2)), /* :- */
            Disjunction(
                Structure(Atom("father"), listOf(variable1, variable2)),
                /* ; */
                Structure(Atom("mother"), listOf(variable1, variable2))
            )
        )

        Program.load(listOf(father, mother, parent))

        val result1 = Program.query(Structure(Atom("parent"), listOf(Atom("john"), Atom("jimmy")))).toList()
        assertEquals(1, result1.size, "Expected 1 result")
        assertTrue(result1[0].isSuccess, "Expected success")
        assertTrue(result1[0].getOrNull()!!.isEmpty(), "Expected no substitutions")
        val result2 = Program.query(Structure(Atom("parent"), listOf(Atom("jane"), Atom("jimmy")))).toList()
        assertEquals(1, result2.size, "Expected 1 result")
        assertTrue(result2[0].isSuccess, "Expected success")
        assertTrue(result2[0].getOrNull()!!.isEmpty(), "Expected no substitutions")

        val result3 = Program.query(Structure(Atom("parent"), listOf(Atom("john"), Atom("jane"))))
        assertFalse(result3.any())
        val result4 = Program.query(Structure(Atom("father"), listOf(Atom("john"), Atom("jane"))))
        assertFalse(result4.any())
    }

    /**
     * John is Jimmy's father.
     * Jane is Jimmy's mother.
     */
    @Test
    fun father_or_mother() {
        val male = Fact(Structure(Atom("male"), listOf(Atom("john"))))
        val female = Fact(Structure(Atom("female"), listOf(Atom("jane"))))

        val parent1 = Fact(Structure(Atom("parent"), listOf(Atom("john"), Atom("jimmy"))))
        val parent2 = Fact(Structure(Atom("parent"), listOf(Atom("jane"), Atom("jimmy"))))

        val variable1 = Variable("X")
        val variable2 = Variable("Y")

        val isFather = Rule(
            Structure(Atom("isFather"), listOf(variable1, variable2)),
            Conjunction(
                Structure(Atom("parent"), listOf(variable1, variable2)),
                Structure(Atom("male"), listOf(variable1))
            )
        )

        val variable3 = Variable("X")
        val variable4 = Variable("Y")

        val isMother = Rule(
            Structure(Atom("isMother"), listOf(variable3, variable4)),
            Conjunction(
                Structure(Atom("parent"), listOf(variable3, variable4)),
                Structure(Atom("female"), listOf(variable3))
            )
        )

        Program.load(listOf(male, female, parent1, parent2, isFather, isMother))

        val result1 = Program.query(Structure(Atom("isFather"), listOf(Atom("john"), Atom("jimmy"))))
        assertTrue(result1.any())

        val result2 = Program.query(Structure(Atom("isMother"), listOf(Atom("jane"), Atom("jimmy"))))
        assertTrue(result2.any())

        val result3 = Program.query(Structure(Atom("isFather"), listOf(Atom("jane"), Atom("jimmy"))))
        assertFalse(result3.any())

        val result4 = Program.query(Structure(Atom("isMother"), listOf(Atom("john"), Atom("jimmy"))))
        assertFalse(result4.any())

        val result5 = Program.query(Structure(Atom("parent"), listOf(Atom("trudy"), Atom("jimmy"))))
        assertFalse(result5.any())
    }

    @Test
    fun requires_backtracking() {
        val fact1 = Fact(Structure(Atom("a"), listOf(Atom("b"))))
        val fact2 = Fact(Structure(Atom("a"), listOf(Atom("c"))))
        val fact3 = Fact(Structure(Atom("a"), listOf(Atom("d"))))

        Program.load(listOf(fact1, fact2, fact3))

        assertTrue(Program.query(Structure(Atom("a"), listOf(Atom("b")))).any())
        assertTrue(Program.query(Structure(Atom("a"), listOf(Atom("c")))).any())
        assertTrue(Program.query(Structure(Atom("a"), listOf(Atom("d")))).any())
    }

    @Test
    fun multiple_choices() {
        val fact1 = Fact(Structure(Atom("a"), listOf(Atom("b"))))
        val fact2 = Fact(Structure(Atom("a"), listOf(Atom("c"))))
        val fact3 = Fact(Structure(Atom("a"), listOf(Atom("d"))))

        Program.load(listOf(fact1, fact2, fact3))

        val results = Query(Structure(Atom("a"), listOf(Variable("X")))).satisfy(emptyMap())

        val expectedResults = listOf(
            mapOf(Variable("X") to Atom("b")),
            mapOf(Variable("X") to Atom("c")),
            mapOf(Variable("X") to Atom("d"))
        )
        val actualResults = results.toList()

        assertEquals(expectedResults.size, actualResults.size, "Number of results should match")
        for (i in expectedResults.indices) {
            assertEquals(expectedResults[i].size, actualResults[i].getOrNull()!!.size, "Substitution size should match")
            assertTrue(expectedResults[i].all {
                actualResults[i].getOrNull()!![it.key]?.let { it1 ->
                    equivalent(
                        it.value,
                        it1,
                        emptyMap()
                    )
                } ?: false
            }, "Substitution values should match")
        }
    }

    @Test
    fun `likes(alice, pizza)`() {
        val fact = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza"))))
        val goal = Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza")))

        Program.load(listOf(fact))

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

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        val subs = result[0].getOrNull()!!
        assertEquals(0, subs.size, "Expected no substitutions")
    }

    @Test
    fun `likes(Person, pizza)`() {
        val fact = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza"))))
        val goal = Structure(Atom("likes"), listOf(Variable("Person"), Atom("pizza")))

        Program.load(listOf(fact))

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

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        val subs = result[0].getOrNull()!!
        assertEquals(1, subs.size, "Expected 1 substitution")
        assertEquals(Atom("alice"), subs[Variable("Person")], "Expected Person to be alice")
    }

    @Test
    fun `likes_food(alice)`() {
        val fact = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza"))))
        val rule = Rule(
            Structure(Atom("likes_food"), listOf(Variable("Person"))),
            Structure(Atom("likes"), listOf(Variable("Person"), Atom("pizza")))
        )

        val goal = Structure(Atom("likes_food"), listOf(Atom("alice")))

        Program.load(listOf(fact, rule))

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

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        val subs = result[0].getOrNull()!!
        assertEquals(0, subs.size, "Expected no substitutions")
    }

    @Test
    fun `likes_food(Person)`() {
        val fact = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza"))))
        val rule = Rule(
            Structure(Atom("likes_food"), listOf(Variable("Person"))),
            Structure(Atom("likes"), listOf(Variable("Person"), Atom("pizza")))
        )

        val goal = Structure(Atom("likes"), listOf(Variable("X"), Atom("pizza")))

        Program.load(listOf(fact, rule))

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

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        val subs = result[0].getOrNull()!!
        assertEquals(1, subs.size, "Expected 1 substitution")
        assertEquals(Atom("alice"), subs[Variable("X")], "Expected Person to be alice")
    }

    @Test
    fun `requires querying exact`() {
        val fact1 = Fact(Atom("a"))
        val fact2 = Fact(Atom("b"))
        val rule1 = Rule(
            Atom("c"),
            Conjunction(
                Atom("a"),
                Atom("b")
            )
        )

        Program.load(listOf(fact1, fact2, rule1))

        val result = Program.query(Atom("c")).toList()

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

    @Test
    fun `requires querying with variable`() {
        val fact1 = Fact(Atom("a"))
        val fact2 = Fact(Atom("b"))
        val rule1 = Rule(
            Structure(Atom("has fact"), listOf(Variable("X"))),
            Variable("X")
        )

        Program.load(listOf(fact1, fact2, rule1))

        val result = Program.query(Structure(Atom("has fact"), listOf(Atom("a")))).toList()

        assertEquals(1, result.size, "Expected 1 result")
        assertTrue(result[0].isSuccess, "Expected success")
        val subs = result[0].getOrNull()!!
        assertEquals(0, subs.size, "Expected no substitutions")
    }

    @Nested
    class `requires querying with filled variable` {
        @BeforeEach
        fun setup() {
            val fact1 = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pizza"))))
            val fact2 = Fact(Structure(Atom("likes"), listOf(Atom("alice"), Atom("pasta"))))
            val fact3 = Fact(Structure(Atom("likes"), listOf(Atom("bob"), Atom("pasta"))))
            val rule1 = Rule(
                Structure(Atom("likes_italian_food"), listOf(Variable("Person"))),
                Disjunction(
                    Structure(Atom("likes"), listOf(Variable("Person"), Atom("pizza"))),
                    Structure(Atom("likes"), listOf(Variable("Person"), Atom("pasta")))
                )
            )

            Program.reset()
            Program.load(listOf(fact1, fact2, fact3, rule1))
        }

        @Test
        fun `likes_italian_food(alice)`() {
            val result = Program.query(Structure(Atom("likes_italian_food"), listOf(Atom("alice")))).toList()

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

            assertTrue(result[0].isSuccess, "Expected success")
            val subs1 = result[0].getOrNull()!!
            assertEquals(0, subs1.size, "Expected no substitutions")

            assertTrue(result[1].isSuccess, "Expected success")
            val subs2 = result[1].getOrNull()!!
            assertEquals(0, subs2.size, "Expected no substitutions")
        }

        @Test
        fun `likes_italian_food(X)`() {
            val result = Program.query(Structure(Atom("likes_italian_food"), listOf(Variable("X")))).toList()

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

            assertTrue(result[0].isSuccess, "Expected success")
            val subs3 = result[0].getOrNull()!!
            assertEquals(1, subs3.size, "Expected 1 substitution, especially without 'Person'")
            assertEquals(Atom("alice"), subs3[Variable("X")], "Expected alice")

            assertTrue(result[1].isSuccess, "Expected success")
            val subs4 = result[1].getOrNull()!!
            assertEquals(1, subs4.size, "Expected 1 substitution, especially without 'Person'")
            assertEquals(Atom("alice"), subs4[Variable("X")], "Expected alice")

            assertTrue(result[2].isSuccess, "Expected success")
            val subs5 = result[2].getOrNull()!!
            assertEquals(1, subs5.size, "Expected 1 substitution, especially without 'Person'")
            assertEquals(Atom("bob"), subs5[Variable("X")], "Expected bob")
        }

        @Test
        fun `likes_italian_food(Person)`() {
            val result = Program.query(Structure(Atom("likes_italian_food"), listOf(Variable("Person")))).toList()

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

            assertTrue(result[0].isSuccess, "Expected success")
            val subs3 = result[0].getOrNull()!!
            assertEquals(1, subs3.size, "Expected 1 substitution")
            assertEquals(Atom("alice"), subs3[Variable("Person")], "Expected alice")

            assertTrue(result[1].isSuccess, "Expected success")
            val subs4 = result[1].getOrNull()!!
            assertEquals(1, subs4.size, "Expected 1 substitution")
            assertEquals(Atom("alice"), subs4[Variable("Person")], "Expected alice")

            assertTrue(result[2].isSuccess, "Expected success")
            val subs5 = result[2].getOrNull()!!
            assertEquals(1, subs5.size, "Expected 1 substitution")
            assertEquals(Atom("bob"), subs5[Variable("Person")], "Expected bob")
        }
    }

    @Test
    fun `leq Peano`() {
        val fact = Fact(Structure(Atom("leq"), listOf(Integer(0), AnonymousVariable.create())))
        val rule = Rule(
            Structure(
                Atom("leq"),
                listOf(Structure(Atom("s"), listOf(Variable("X"))), Structure(Atom("s"), listOf(Variable("Y"))))
            ),
            Structure(Atom("leq"), listOf(Variable("X"), Variable("Y"))),
        )

        Program.db.load(listOf(fact, rule))

        val result1 = Program.query(Structure(Atom("leq"), listOf(Variable("X"), Integer(0)))).toList()

        assertEquals(1, result1.size, "Expected 1 result")
        assertTrue(result1[0].isSuccess, "Expected success")
        val subs = result1[0].getOrNull()!!
        assertEquals(1, subs.size, "Expected 1 substitution")
        assertEquals(Integer(0), subs[Variable("X")], "Expected X to be 0")

        val result2 =
            Program.query(Structure(Atom("leq"), listOf(Variable("X"), Structure(Atom("s"), listOf(Integer(0))))))
                .toList()

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

        assertTrue(result2[0].isSuccess, "Expected success")
        val subs2a = result2[0].getOrNull()!!
        assertEquals(1, subs2a.size, "Expected 1 substitution")
        assertEquals(Integer(0), subs2a[Variable("X")], "Expected X to be 0")

        assertTrue(result2[1].isSuccess, "Expected success")
        val subs2b = result2[1].getOrNull()!!
        assertEquals(1, subs2b.size, "Expected 1 substitution")
        assertEquals(Structure(Atom("s"), listOf(Integer(0))), subs2b[Variable("X")], "Expected X to be s(0)")

        val result3 = Program.query(
            Structure(
                Atom("leq"),
                listOf(Variable("X"), Structure(Atom("s"), listOf(Structure(Atom("s"), listOf(Integer(0))))))
            )
        ).toList()

        assertEquals(3, result3.size, "Expected 3 results")
        assertTrue(result3[0].isSuccess, "Expected success")
        val subs3a = result3[0].getOrNull()!!
        assertEquals(1, subs3a.size, "Expected 1 substitution")
        assertEquals(Integer(0), subs3a[Variable("X")], "Expected X to be 0")
        assertTrue(result3[1].isSuccess, "Expected success")
        val subs3b = result3[1].getOrNull()!!
        assertEquals(1, subs3b.size, "Expected 1 substitution")
        assertEquals(Structure(Atom("s"), listOf(Integer(0))), subs3b[Variable("X")], "Expected X to be s(0)")
        assertTrue(result3[2].isSuccess, "Expected success")
        val subs3c = result3[2].getOrNull()!!
        assertEquals(1, subs3c.size, "Expected 1 substitution")
        assertEquals(Structure(Atom("s"), listOf(Structure(Atom("s"), listOf(Integer(0))))), subs3c[Variable("X")], "Expected X to be s(s(0))")
    }
}