package prolog.logic

import prolog.Answer
import prolog.Answers
import prolog.Substitutions
import prolog.ast.arithmetic.Expression
import prolog.ast.arithmetic.Float
import prolog.ast.arithmetic.Integer
import prolog.ast.arithmetic.Number
import prolog.ast.lists.List.Cons
import prolog.ast.lists.List.Empty
import prolog.ast.logic.Fact
import prolog.ast.terms.*
import prolog.ast.lists.List as PList

// Apply substitutions to a term
fun applySubstitution(term: Term, subs: Substitutions): Term = when {
    term is Fact -> term.applySubstitution(subs)

    variable(term, emptyMap()) -> {
        val variable = term as Variable
        subs[variable]?.let { applySubstitution(term = it, subs = subs) } ?: term
    }

    atomic(term, subs) -> term
    compound(term, subs) -> term.applySubstitution(subs)

    else -> term
}

fun applySubstitution(expr: Expression, subs: Substitutions): Expression {
    return applySubstitution(expr as Term, subs) as Expression
}

// Check if a variable occurs in a term
fun occurs(variable: Variable, term: Term, subs: Substitutions = emptyMap()): Boolean = when {
    variable(term, subs) -> term == variable
    atomic(term, subs) -> false
    compound(term, subs) -> {
        val structure = term as Structure
        structure.arguments.any { occurs(variable, it, subs) }
    }

    else -> false
}

// Unify two terms with backtracking and lazy evaluation
fun unifyLazy(term1: Term, term2: Term, subs: Substitutions): Answers = sequence {
    val t1 = applySubstitution(term1, subs)
    val t2 = applySubstitution(term2, subs)

    when {
        equivalent(t1, t2, subs) -> {
            yield(Result.success(emptyMap()))
        }

        variable(t1, subs) -> {
            val variable = t1 as Variable
            if (!occurs(variable, t2, subs)) {
                yield(Result.success(mapOf(term1 to t2)))
            }
        }

        variable(t2, subs) -> {
            val variable = t2 as Variable
            if (!occurs(variable, t1, subs)) {
                yield(Result.success(mapOf(term2 to t1)))
            }
        }

        compound(t1, subs) && compound(t2, subs) -> {
            val structure1 = t1 as Structure
            val structure2 = t2 as Structure
            if (structure1.functor == structure2.functor) {
                // Unify each argument at the same time, and yield the result
                val args1 = structure1.arguments
                val args2 = structure2.arguments
                if (args1.size == args2.size) {
                    yieldAll(unifyArgs(args1, args2, subs))
                }
            }
        }

        t1 is PList && t2 is PList -> {
            if (equivalent(t1, t2, subs)) {
                yield(Result.success(emptyMap()))
            } else if (t1.size == t2.size) {
                val head1 = t1.head
                val head2 = t2.head

                unifyLazy(head1, head2, subs).forEach { headSubs ->
                    headSubs.map { headResult ->
                        if (t1 is Empty && t2 is Empty) {
                            yield(Result.success(headResult))
                        } else if (t1 is Cons && t2 is Cons) {
                            val tail1 = t1.tail
                            val tail2 = t2.tail
                            unifyLazy(tail1, tail2, headResult).forEach { tailSubs ->
                                tailSubs.map { tailResult ->
                                    yield(Result.success(headResult + tailResult))
                                }
                            }
                        }
                    }
                }
            }
        }

        else -> {}
    }
}

private fun unifyArgs(
    args1: List<Argument>,
    args2: List<Argument>,
    subs: Substitutions
): Answers = sequence {
    // Using the current subs, unify the first argument of each list
    val arg1 = applySubstitution(args1[0], subs)
    val arg2 = applySubstitution(args2[0], subs)
    unifyLazy(arg1, arg2, subs).forEach { headResult ->
        // Use the resulting substitutions to unify the rest of the arguments
        headResult.map { headSubs ->
            val rest1 = args1.drop(1).map { applySubstitution(it, headSubs) }
            val rest2 = args2.drop(1).map { applySubstitution(it, headSubs) }
            if (rest1.isEmpty() && rest2.isEmpty()) {
                yield(Result.success(headSubs))
            } else if (rest1.isNotEmpty() && rest2.isNotEmpty()) {
                unifyArgs(rest1, rest2, subs).forEach { tailResult ->
                    tailResult.map { tailSubs ->
                        yield(Result.success(headSubs + tailSubs))
                    }
                }
            }
        }
    }
}

fun unify(term1: Term, term2: Term): Answer {
    val substitutions = unifyLazy(term1, term2, emptyMap()).iterator()
    return if (substitutions.hasNext()) {
        substitutions.next()
    } else {
        Result.failure(NoSuchElementException())
    }
}

/**
 * True if Term1 is equivalent to Term2. A variable is only identical to a sharing variable.
 */
fun equivalent(term1: Term, term2: Term, subs: Substitutions): Boolean {
    return when {
        term1 is Atom && term2 is Atom -> compare(term1, term2, subs) == 0
        term1 is Structure && term2 is Structure -> {
            term1.functor == term2.functor && term1.arguments.zip(term2.arguments).all { (arg1, arg2) ->
                equivalent(arg1, arg2, subs)
            }
        }

        term1 is Integer && term2 is Integer -> compare(term1, term2, subs) == 0
        term1 is Number && term2 is Number -> compare(term1, term2, subs) == 0
        term1 is Variable && term2 is Variable -> term1 == term2
        term1 is Variable -> term1 in subs && equivalent(subs[term1]!!, term2, subs)
        term2 is Variable -> term2 in subs && equivalent(subs[term2]!!, term1, subs)
        term1 is PList && term2 is PList -> {
            if (term1.isEmpty() && term2.isEmpty()) {
                true
            } else if (term1.isEmpty() || term2.isEmpty()) {
                false
            } else {
                require(term1 is Cons && term2 is Cons)
                val head1 = term1.head
                val tail1 = term1.tail
                val head2 = term2.head
                val tail2 = term2.tail
                equivalent(head1, head2, subs) && equivalent(tail1, tail2, subs)
            }
        }

        else -> false
    }
}

/**
 * See also [SWI Prolog Standard Order of Terms](https://www.swi-prolog.org/pldoc/man?section=standardorder)
 */
fun compare(term1: Term, term2: Term, subs: Substitutions): Int {
    val t1 = applySubstitution(term1, subs)
    val t2 = applySubstitution(term2, subs)

    return when (t1) {
        is Variable -> {
            when (t2) {
                is Variable -> t1.name.compareTo(t2.name)
                is Number -> -1
                is Atom -> -1
                is Structure -> -1
                else -> throw IllegalArgumentException("Cannot compare $t1 with $t2")
            }
        }

        is Number -> {
            when (t2) {
                is Variable -> 1
                is Integer -> (t1.value as Int).compareTo(t2.value)
                is Float -> (t1.value as kotlin.Float).compareTo(t2.value)
                is Atom -> -1
                is Structure -> -1
                else -> throw IllegalArgumentException("Cannot compare $t1 with $t2")
            }
        }

        is Atom -> {
            when (t2) {
                is Variable -> 1
                is Number -> 1
                is Atom -> t1.name.compareTo(t2.name)
                is Structure -> -1
                else -> throw IllegalArgumentException("Cannot compare $t1 with $t2")
            }
        }

        is Structure -> {
            when (t2) {
                is Variable -> 1
                is Number -> 1
                is Atom -> 1
                is Structure -> {
                    val arityComparison = t1.arguments.size.compareTo(t2.arguments.size)
                    if (arityComparison != 0) return arityComparison
                    val nameComparison = t1.name.compareTo(t2.name)
                    if (nameComparison != 0) return nameComparison
                    t1.arguments.zip(t2.arguments).forEach { (arg1, arg2) ->
                        val argsComparison = equivalent(arg1, arg2, emptyMap())
                        if (!argsComparison) return compare(arg1, arg2, subs)
                    }
                    return 0
                }

                else -> throw IllegalArgumentException("Cannot compare $t1 with $t2")
            }
        }

        else -> throw IllegalArgumentException("Cannot compare $t1 with $t2")
    }
}
