package prolog

import prolog.builtins.atomic
import prolog.builtins.compound
import prolog.builtins.variable
import prolog.components.Goal
import prolog.components.terms.Term
import prolog.components.terms.Variable
import prolog.components.terms.Structure
import java.util.*

typealias Substitution = Map<Variable, Term>

// Apply substitutions to a term
fun applySubstitution(term: Term, substitution: Substitution): Term = when {
    variable(term) -> (term as Variable).alias().map { applySubstitution(it, substitution) }.orElse(term)
    atomic(term)   -> term
    compound(term) -> {
        val structure = term as Structure
        Structure(structure.name, structure.arguments.map { applySubstitution(it, substitution) })
    }
    else -> term
}

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

// Generate possible substitutions
fun generateSubstitutions(term1: Term, term2: Term, substitution: Substitution): Sequence<Substitution> = sequence {
    val t1 = applySubstitution(term1, substitution)
    val t2 = applySubstitution(term2, substitution)

    when {
        t1 == t2      -> yield(substitution)
        variable(t1) -> {
            val variable = t1 as Variable
            if (!occurs(variable, t2)) {
                variable.bind(t2)
                yield(substitution + (variable to t2))
            }
        }
        variable(t2) -> {
            val variable = t2 as Variable
            if (!occurs(variable, t1)) {
                variable.bind(t1)
                yield(substitution + (variable to t1))
            }
        }
        compound(t1) && compound(t2) -> {
            val structure1 = t1 as Structure
            val structure2 = t2 as Structure
            if (structure1.functor == structure2.functor) {
                val newSubstitution = structure1.arguments.zip(structure2.arguments).fold(substitution) { acc, (arg1, arg2) ->
                    generateSubstitutions(arg1, arg2, acc).firstOrNull() ?: return@sequence
                }
                yield(newSubstitution)
            }
        }
        else -> {}
    }
}

// Unify two terms with backtracking and lazy evaluation
fun unifyLazy(term1: Term, term2: Term, substitution: Substitution = emptyMap()): Sequence<Substitution> {
    return generateSubstitutions(term1, term2, substitution)
}

fun unify(term1: Term, term2: Term): Optional<Substitution> {
    val substitutions = unifyLazy(term1, term2).toList()
    return if (substitutions.isEmpty()) {
        Optional.empty()
    } else {
        Optional.of(substitutions.first())
    }
}