package prolog.logic

import org.junit.jupiter.api.Test
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import prolog.ast.terms.Atom
import prolog.ast.terms.Structure
import prolog.ast.terms.Variable

class TermsTests {
    @Test
    fun `rename vars in atom`() {
        val term = Atom("a")
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start, end, "Expected end to still be at start")
        assertTrue(subs.isEmpty(), "Expected no substitutions")
    }

    @Test
    fun `rename vars in var`() {
        val term = Variable("X")
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start + 1, end, "Expected end to be incremented by 1")
        assertEquals(1, subs.size, "Expected one substitution")
        assertTrue(subs.containsKey(term), "Expected subs to contain the original term")
        assertEquals(Variable("X@$start"), subs[term], "Expected subs to contain the new term")
    }

    @Test
    fun `rename vars in compound term without vars`() {
        val term = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start, end, "Expected end to still be at start")
        assertTrue(subs.isEmpty(), "Expected no substitutions")
    }

    @Test
    fun `rename vars in compound term`() {
        val term = Structure(Atom("f"), listOf(Variable("X"), Variable("Y")))
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start + 2, end, "Expected end to be incremented by 2")
        assertEquals(2, subs.size, "Expected two substitutions")
        assertTrue(subs.containsKey(term.arguments[0]), "Expected subs to contain the first original term")
        assertEquals(Variable("X@$start"), subs[term.arguments[0]], "Expected subs to contain the new term")
        assertTrue(subs.containsKey(term.arguments[1]), "Expected subs to contain the second original term")
        assertEquals(Variable("Y@${start + 1}"), subs[term.arguments[1]], "Expected subs to contain the new term")
    }

    @Test
    fun `renaming identical vars should keep the same name`() {
        val term = Structure(Atom("f"), listOf(Variable("X"), Variable("X")))
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start + 1, end, "Expected end to be incremented by 1")
        assertEquals(1, subs.size, "Expected one substitution")
        assertTrue(subs.containsKey(term.arguments[0]), "Expected subs to contain the first original term")
        assertEquals(Variable("X@$start"), subs[term.arguments[0]], "Expected subs to contain the new term")
        assertTrue(subs.containsKey(term.arguments[1]), "Expected subs to contain the second original term")
        assertEquals(Variable("X@$start"), subs[term.arguments[1]], "Expected subs to contain the new term")
    }

    @Test
    fun `renaming identical vars should keep the same name in nested terms`() {
        val term = Structure(Atom("f"), listOf(Variable("X"), Structure(Atom("g"), listOf(Variable("X")))))
        val start = 0

        val (end, subs) = numbervars(term, start)

        assertEquals(start + 1, end, "Expected end to be incremented by 1")
        assertEquals(1, subs.size, "Expected one substitution")
        assertTrue(subs.containsKey(Variable("X")), "Expected subs to contain the variable")
        assertEquals(Variable("X@$start"), subs[term.arguments[0]], "Expected subs to contain the new term")
    }

    @Test
    fun `renaming multiple times`() {
        val variable = Variable("X")
        val term = Structure(Atom("f"), listOf(variable))
        val start = 0

        val (end1, subs1) = numbervars(term, start, emptyMap())

        assertEquals(start + 1, end1, "Expected end to be incremented by 1")
        assertEquals(1, subs1.size, "Expected one substitution")
        assertTrue(subs1.containsKey(variable), "Expected subs to contain the variable")
        assertEquals(Variable("X@$start"), subs1[variable], "Expected subs to contain the new term")
        val variable1 = subs1[variable] as Variable

        val (end2, subs2) = numbervars(term, end1, subs1)

        assertEquals(start + 2, end2, "Expected end to be incremented by 2")
        assertEquals(1, subs2.size, "Expected one substitution")
        assertTrue(subs2.containsKey(variable1), "Expected subs to contain the variable")
        assertEquals(Variable("X@$start@$end1"), subs2[variable1], "Expected subs to contain the new term")
    }
}