package prolog

import org.junit.jupiter.api.Assertions.*
import org.junit.jupiter.api.BeforeEach
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

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

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