#include <clingo/control/term.hh>

#include <clingo/input/rewrite/analyze.hh>

#include <clingo/util/type_traits.hh>

namespace CppClingo::Control {

namespace {

//! Maps a binary input operator to the corresponding ground one.
//!
//! @note: The input interval operator has no matching ground version.
//! @note: Candidate for core.
inline auto map_binary_op(Input::BinaryOperator op) -> Ground::BinaryOperator {
    switch (op) {
        case Input::BinaryOperator::and_: {
            return Ground::BinaryOperator::and_;
        }
        case Input::BinaryOperator::div: {
            return Ground::BinaryOperator::div;
        }
        case Input::BinaryOperator::dots: {
            break;
        }
        case Input::BinaryOperator::minus: {
            return Ground::BinaryOperator::minus;
        }
        case Input::BinaryOperator::mod: {
            return Ground::BinaryOperator::mod;
        }
        case Input::BinaryOperator::or_: {
            return Ground::BinaryOperator::or_;
        }
        case Input::BinaryOperator::plus: {
            return Ground::BinaryOperator::plus;
        }
        case Input::BinaryOperator::pow: {
            return Ground::BinaryOperator::pow;
        }
        case Input::BinaryOperator::times: {
            return Ground::BinaryOperator::times;
        }
        case Input::BinaryOperator::xor_: {
            return Ground::BinaryOperator::xor_;
        }
    }
    throw std::runtime_error("unsupported binary operator");
}

//! Translates input terms to their ground representation.
//!
//! Input terms must be rewritten before translation.
class BuilderTerm {
  public:
    //! Construct a term builder.
    //!
    //! The reference has_projection is used to track, whether a projection
    //! star occurred in the term. The var_map dictionary is used to map
    //! variables to integers. Each name must have been assigned an integer
    //! beforehand.
    BuilderTerm(bool &has_projection, Util::unordered_map<String, size_t> const &var_map)
        : has_projection_{&has_projection}, var_map_{&var_map} {}

    //! Translate variables.
    auto operator()(Input::TermVariable const &term) const -> Ground::UTerm {
        assert(var_map_->contains(term.name()));
        return std::make_unique<Ground::TermVariable>(var_map_->find(term.name())->second);
    }
    //! Translate symbols.
    auto operator()(Input::TermSymbol const &term) const -> Ground::UTerm {
        return std::make_unique<Ground::TermSymbol>(term.value());
    }
    // Translate format strings.
    auto operator()(Input::TermFormatString const &term) const -> Ground::UTerm {
        auto elems = Ground::FormatFieldVec{};
        elems.reserve(term.elems().size());
        for (auto const &elem : term.elems()) {
            std::visit(
                [&, this]<class T>(T const &elem) {
                    if constexpr (Util::matches<T, Input::FormatFieldExpression>) {
                        elems.emplace_back(std::in_place_type<std::pair<Ground::UTerm, FormatSpec>>,
                                           std::visit(*this, *elem.lhs()), elem.rhs());
                    } else {
                        static_assert(Util::matches<T, Input::FormatFieldLiteral>);
                        elems.emplace_back(std::in_place_type<SharedString>, elem.value());
                    }
                },
                elem);
        }
        return std::make_unique<Ground::TermFormatString>(std::move(elems));
    }
    //! Translate arguments of tuples and functions.
    [[nodiscard]] auto handle_args(Input::ArgumentArray const &args) const -> Ground::UTermVec {
        Ground::UTermVec g_args;
        g_args.reserve(args.size());
        for (auto const &arg : args) {
            g_args.emplace_back(std::visit(
                [this]<class T>(T const &arg) -> Ground::UTerm {
                    if constexpr (Util::matches<T, Input::Projection>) {
                        *has_projection_ = true;
                        return std::make_unique<Ground::TermProjection>();
                    } else {
                        return std::visit(*this, arg);
                    }
                },
                arg));
        }
        return g_args;
    }
    //! Translate tuples.
    //!
    //! Assumes that the arguments consists of a single pool.
    auto operator()(Input::TermTuple const &term) const -> Ground::UTerm {
        assert(term.pool().size() == 1 && std::holds_alternative<Input::ArgumentTuple>(term.pool().front()));
        return std::make_unique<Ground::TermTuple>(
            handle_args(std::get<Input::ArgumentTuple>(term.pool().front()).elems()));
    }
    //! Translate functions.
    //!
    //! Assumes that the arguments consist of a single pool.
    auto operator()(Input::TermFunction const &term) const -> Ground::UTerm {
        assert(!term.external() && term.pool().size() == 1);
        return std::make_unique<Ground::TermFunction>(term.name(), handle_args(term.pool().front().elems()));
    }
    //! Translate absolute value terms.
    //!
    //! Assumes that there is a single argument.
    auto operator()(Input::TermAbs const &term) const -> Ground::UTerm {
        assert(term.pool().size() == 1);
        auto const &rhs = term.pool().front();
        return std::make_unique<Ground::TermUnary>(Ground::UnaryOperator::abs, location(rhs), std::visit(*this, rhs));
    }
    //! Translate unary terms.
    auto operator()(Input::TermUnary const &term) const -> Ground::UTerm {
        Ground::UnaryOperator op =
            term.op() == Input::UnaryOperator::minus ? Ground::UnaryOperator::minus : Ground::UnaryOperator::negate;
        return std::make_unique<Ground::TermUnary>(op, location(*term.rhs()), std::visit(*this, *term.rhs()));
    }
    //! Translate binary terms.
    //!
    //! Intervals must be removed by rewriting beforehand.
    auto operator()(Input::TermBinary const &term) const -> Ground::UTerm {
        assert(term.op() != Input::BinaryOperator::dots);
        if (auto lin = Input::check_linear(term); lin) {
            assert(var_map_->contains(lin->x()));
            return std::make_unique<Ground::TermLinear>(term.loc(), lin->m(), var_map_->find(lin->x())->second,
                                                        lin->n());
        }
        return std::make_unique<Ground::TermBinary>(location(*term.lhs()), std::visit(*this, *term.lhs()),
                                                    map_binary_op(term.op()), location(*term.rhs()),
                                                    std::visit(*this, *term.rhs()));
    }

  private:
    bool *has_projection_;
    Util::unordered_map<String, size_t> const *var_map_;
};

//! Translates input theory terms to their ground representation.
//!
//! Input terms must be rewritten before translation.
class BuilderTheoryTerm {
  public:
    //! Construct a theory term builder.
    //!
    //! The var_map dictionary is used to map variables to integers. Each name
    //! must have been assigned an integer beforehand.
    BuilderTheoryTerm(Util::unordered_map<String, size_t> const &var_map) : var_map_{&var_map} {}

    //! Translate unparsed terms.
    auto operator()([[maybe_unused]] Input::TheoryTermUnparsed const &term) const -> Ground::UTheoryTerm {
        throw std::logic_error("unexpected unparsed theory term");
    }
    //! Translate theory variables.
    auto operator()(Input::TheoryTermVariable const &term) const -> Ground::UTheoryTerm {
        assert(var_map_->contains(term.name()));
        return std::make_unique<Ground::TheoryTermVariable>(var_map_->find(term.name())->second);
    }
    //! Translate theory symbols.
    auto operator()(Input::TheoryTermSymbol const &term) const -> Ground::UTheoryTerm {
        return std::make_unique<Ground::TheoryTermSymbol>(term.value());
    }
    //! Translate arguments of tuples and functions.
    [[nodiscard]] auto handle_args(Input::TheoryTermArray const &args) const -> Ground::UTheoryTermVec {
        Ground::UTheoryTermVec g_args;
        g_args.reserve(args.size());
        for (auto const &arg : args) {
            g_args.emplace_back(std::visit(*this, arg));
        }
        return g_args;
    }
    //! Translate tuples.
    auto operator()(Input::TheoryTermTuple const &term) const -> Ground::UTheoryTerm {
        return std::make_unique<Ground::TheoryTermTuple>(term.type(), handle_args(term.elems()));
    }
    //! Translate functions.
    auto operator()(Input::TheoryTermFunction const &term) const -> Ground::UTheoryTerm {
        return std::make_unique<Ground::TheoryTermFunction>(term.name(), handle_args(term.args()));
    }

  private:
    Util::unordered_map<String, size_t> const *var_map_;
};

} // namespace

//! Translates input theory terms to their ground representation.
auto build_term(Util::unordered_map<String, size_t> const &var_map, Input::Term const &term, bool &has_projection)
    -> Ground::UTerm {
    return std::visit(BuilderTerm{has_projection, var_map}, term);
}

//! Translates input theory terms to their ground representation.
auto build_theory_term(Util::unordered_map<String, size_t> const &var_map, Input::TheoryTerm const &term)
    -> Ground::UTheoryTerm {
    return std::visit(BuilderTheoryTerm{var_map}, term);
}

} // namespace CppClingo::Control
