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

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"))))
        assertTrue(result1.toList().isNotEmpty())
        val result2 = Program.query(Structure(Atom("parent"), listOf(Atom("jane"), Atom("jimmy"))))
        assertTrue(result2.toList().isNotEmpty())

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