package interpreter

import io.Logger
import prolog.ast.arithmetic.Expression
import prolog.ast.arithmetic.Integer
import prolog.ast.lists.List as PList
import prolog.ast.logic.Clause
import prolog.ast.logic.Fact
import prolog.ast.logic.LogicOperand
import prolog.ast.logic.Rule
import prolog.ast.terms.*
import prolog.builtins.*

/**
 * Preprocessor for Prolog
 *
 * This class preprocesses Prolog code and applies various transformations such as recognizing builtins.
 */
open class Preprocessor {
    /**
     * Preprocesses the input Prolog code.
     *
     * @param input The already parsed Prolog code as a list of clauses.
     * @return The preprocessed Prolog code as a list of clauses.
     */
    fun preprocess(input: List<Clause>): List<Clause> {
        return input.map { preprocess(it) }
    }

    fun preprocess(input: Query): Query {
        return Query(preprocess(input.query) as Goal)
    }

    private fun preprocess(clause: Clause): Clause {
        return when (clause) {
            is Fact -> {
                Fact(preprocess(clause.head) as Head)
            }

            is Rule -> {
                Rule(
                    preprocess(clause.head) as Head,
                    preprocess(clause.body as Term) as Body
                )
            }

            else -> clause
        }
    }

    protected open fun preprocess(term: Term, nested: Boolean = false): Term {
        // TODO Remove hardcoding by storing the functors as constants in operators?

        val prepped = when (term) {
            Variable("_") -> AnonymousVariable.create()
            is Atom, is Structure -> {
                // Preprocess the arguments first to recognize builtins
                val args = if (term is Structure) {
                    term.arguments.map { preprocess(it, nested = true) }
                } else emptyList()

                when (term.functor) {
                    // Analysis
                    Functor.of("functor/3") -> FunctorOp(args[0], args[1], args[2])
                    Functor.of("arg/3") -> Arg(args[0], args[1], args[2])
                    Functor.of("clause/2") -> ClauseOp(args[0] as Head, args[1] as Body)
                    Functor.of("atomic/1") -> AtomicOp(args[0])
                    Functor.of("compound/1") -> CompoundOp(args[0])
                    Functor.of("=../2") -> Univ(args[0], args[1])

                    // Arithmetic
                    Functor.of("inf/0") -> Integer(Int.MAX_VALUE)
                    Functor.of("=\\=/2") -> if (args.all { it is Expression }) {
                        EvaluatesToDifferent(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("=:=/2") -> if (args.all { it is Expression }) {
                        EvaluatesTo(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("is/2") -> if (args.all { it is Expression }) {
                        Is(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("-/1") -> if (args.all { it is Expression }) Negate(args[0] as Expression) else term
                    Functor.of("+/1") -> if (args.all { it is Expression }) Positive(args[0] as Expression) else term
                    Functor.of("+/2") -> if (args.all { it is Expression }) {
                        Add(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("-/2") -> if (args.all { it is Expression }) {
                        Subtract(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("*/2") -> if (args.all { it is Expression }) {
                        Multiply(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("//2") -> if (args.all { it is Expression }) {
                        Divide(args[0] as Expression, args[1] as Expression)
                    } else term

                    Functor.of("between/3") -> if (args.all { it is Expression }) {
                        Between(args[0] as Expression, args[1] as Expression, args[2] as Expression)
                    } else term

                    Functor.of("succ/2") -> if (args.all { it is Expression }) {
                        Successor(args[0] as Expression, args[1] as Expression)
                    } else term

                    // Control
                    Functor.of("fail/0") -> Fail
                    Functor.of("false/0") -> False
                    Functor.of("true/0") -> True
                    Functor.of("!/0") -> Cut()
                    Functor.of(",/2") -> Conjunction(args[0] as LogicOperand, args[1] as LogicOperand)
                    Functor.of(";/2") -> Disjunction(args[0] as LogicOperand, args[1] as LogicOperand)
                    Functor.of("\\+/1") -> Not(args[0] as Goal)

                    // Database
                    Functor.of("dynamic/1") -> Dynamic(Functor.of((args[0] as Atom).name))
                    Functor.of("retract/1") -> Retract(args[0])
                    Functor.of("retractall/1") -> RetractAll(args[0])
                    Functor.of("assert/1") -> {
                        if (args[0] is Rule) {
                            Assert(args[0] as Rule)
                        } else {
                            Assert(Fact(args[0] as Head))
                        }
                    }

                    Functor.of("asserta/1") -> {
                        if (args[0] is Rule) {
                            AssertA(args[0] as Rule)
                        } else {
                            AssertA(Fact(args[0] as Head))
                        }
                    }

                    Functor.of("assertz/1") -> {
                        if (args[0] is Rule) {
                            AssertZ(args[0] as Rule)
                        } else {
                            AssertZ(Fact(args[0] as Head))
                        }
                    }

                    // IO
                    Functor.of("write/1") -> Write(args[0])
                    Functor.of("nl/0") -> Nl
                    Functor.of("read/1") -> Read(args[0])

                    // Lists
                    Functor.of("member/2") -> Member(args[0], args[1] as PList)

                    // Meta
                    Functor.of("call/1") -> Call(args[0] as Goal)
                    Functor.of("ignore/1") -> Ignore(args[0] as Goal)

                    // Other
                    Functor.of("initialization/1") -> Initialization(args[0] as Goal)
                    Functor.of("forall/2") -> ForAll(args[0] as LogicOperand, args[1] as Goal)

                    // Unification
                    Functor.of("=/2") -> Unify(args[0], args[1])
                    Functor.of("\\=/2") -> NotUnify(args[0], args[1])
                    Functor.of("==/2") -> Equivalent(args[0], args[1])
                    Functor.of("\\==/2") -> NotEquivalent(args[0], args[1])

                    Functor.of(":-/2") -> Rule(args[0] as Head, args[1] as Body)

                    else -> {
                        if (term is Structure) term.arguments = args
                        term
                    }
                }
            }

            else -> term
        }

        Logger.debug(
            "Preprocessed term $term into $prepped (kind ${prepped::class.simpleName})",
            !nested && (prepped != term || prepped::class != term::class)
        )

        return prepped
    }
}