#include "ast.hh"
#include "control.hh"
#include "core.hh"
#include "iterable.hh"
#include "symbol.hh"
#include "util.hh"

#include <pybind11/functional.h>
#include <pybind11/native_enum.h>
#include <pybind11/operators.h>

#include <algorithm>
#include <numeric>

// NOLINTBEGIN(readability-convert-member-functions-to-static,performance-enum-size)

namespace PyClingo::AST {

namespace py = pybind11;

using StringArray = std::vector<std::string>;
using StringIterable = Iterable<std::string>;

using SymbolArray = std::vector<Symbol>;
using SymbolIterable = Iterable<Symbol>;

auto to_string_array(std::vector<std::string_view> arr) -> StringArray {
    auto res = StringArray();
    res.insert(res.end(), arr.begin(), arr.end());
    return res;
}

template <class T>
auto c_cast(std::optional<T> const &opt) -> clingo_ast_t*;

template <class... Ts>
auto c_cast(std::variant<Ts...> const &var) -> clingo_ast_t*;

template <class T>
auto c_cast(std::vector<T> const &arr) -> std::vector<clingo_ast_t*>;

auto c_cast(StringArray const &arr) -> std::vector<clingo_string_t> {
    std::vector<clingo_string_t> ret;
    ret.reserve(arr.size());
    for (auto const &str : arr) {
        ret.emplace_back(str.data(), str.size());
    }
    return ret;
}

auto c_cast(SymbolArray const &arr) -> std::vector<clingo_symbol_t> {
    std::vector<clingo_symbol_t> ret;
    ret.reserve(arr.size());
    for (auto const &sym : arr) {
        ret.emplace_back(sym.handle());
    }
    return ret;
}

template <class T>
auto c_cast(Iterable<T> const &it) {
    return c_cast(it.vector());
}

template <class Cons>
void visit_array(clingo_ast_t *ast, clingo_ast_attribute_t attr, py::handle visitor, py::args const &args, py::kwargs const &kwargs, Cons cons);

template <class Array>
auto transform_array(Array arr, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Array, bool>;

template <class Value>
auto transform_value(Value val, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Value, bool>;

template <class Value>
auto transform_opt_value(Value opt, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Value, bool>;

class ASTBase {
public:
    // Note: for pybind
    ASTBase() = default;

    ASTBase(ASTBase const &x) {
        handle_error(clingo_ast_copy(x.ast_, &ast_));
    }

    ASTBase(ASTBase &&x) noexcept {
        std::swap(ast_, x.ast_);
    }

    ~ASTBase() noexcept {
        clingo_ast_free(ast_);
    }

    // NOLINTNEXTLINE(bugprone-unhandled-self-assignment)
    auto operator=(ASTBase const &x) -> ASTBase& {
        if (ast_ != x.ast_) {
            clingo_ast_free(ast_);
            ast_ = nullptr;
            handle_error(clingo_ast_copy(x.ast_, &ast_));
        }
        return *this;
    }

    auto operator=(ASTBase &&x) noexcept -> ASTBase& {
        std::swap(ast_, x.ast_);
        return *this;
    }

    [[nodiscard]] auto hash() const -> size_t {
        return clingo_ast_hash(ast_);
    }

    auto to_string() -> std::string_view {
        auto *bld = string_builder();
        handle_error(clingo_ast_to_string(ast_, bld));
        clingo_string_t str;
        handle_error(clingo_string_builder_string(bld, &str));
        return {str.data, str.size};
    }

    friend auto operator==(ASTBase const &a, ASTBase const &b) -> bool {
        return clingo_ast_equal(a.ast_, b.ast_);
    }

    friend auto operator<=>(ASTBase const &a, ASTBase const &b) -> std::strong_ordering {
        return clingo_ast_compare(a.ast_, b.ast_) <=> 0;
    }

    friend auto c_cast(ASTBase const &x) -> clingo_ast_t *;
    friend auto c_cast(ASTBase const *x) -> clingo_ast_t *;

protected:
    ASTBase(clingo_ast_t *ast) : ast_{ast} { }

    // NOLINTNEXTLINE
    clingo_ast_t *ast_ = nullptr;
};

auto c_cast(ASTBase const &x) -> clingo_ast_t * {
    return x.ast_;
}

auto c_cast(ASTBase const *x) -> clingo_ast_t * {
    return x->ast_;
}

struct CString {
    CString(std::string_view str) : str_{str} {}
    CString(std::string &&str) : str_{std::move(str)} {}
    operator std::string_view() const {
        return std::visit([]<class T>(T const &x){
            if constexpr (std::is_same_v<T, std::string>) {
                return std::string_view{x};
            }
            else {
                return x;
            }
        }, str_);
    }
    std::variant<std::string, std::string_view> str_;
};

template <class T>
using update_result_t = std::conditional_t<std::is_same_v<T, std::string_view>, CString, T>;

template <class T, class F, class M> auto update_value(F *self, M fun, py::kwargs const &kwargs, char const *attr) -> update_result_t<T>;

enum class ProjectionMode {
    Disabled = clingo_projection_mode_disabled,
    Anonymous = clingo_projection_mode_anonymous,
    Pure = clingo_projection_mode_pure,
};

{%  for type_name, type in types.items() -%}
{%    set c_type_name = type_name | camel -%}
{%    if type.type == 'optional' %}
using {{ c_type_name }} = std::optional<{{ type.value_type | camel }}>;
{%    elif type.type == 'array' %}
{%      for name in type["value_type"] | forward(type_name, types) -%}
class {{ name | camel }};

{%      endfor -%}
using {{ c_type_name }} = std::vector<{{ type.value_type | camel }}>;
using {{ type.value_type | camel }}Iterable = Iterable<{{ type.value_type | camel }}>;

auto construct_{{ type_name }}(clingo_ast_t **ast, size_t size) -> {{ c_type_name }};
{%    elif type.type == 'union' %}
{%     for name in type["types"] | forward(type_name, types) -%}
class {{ name | camel }};

{%     endfor -%}
using {{ c_type_name }} = std::variant<{{ type.types | map("camel") | join(", ") }}>;

auto construct_{{ type_name }}(clingo_ast_t *ast) -> {{ c_type_name }};
{%    elif type.type == 'enum' %}
enum class {{ c_type_name }} {
{%-     for name, attr in type["values"].items() %}
    {{ name | camel }} = {{ attr.value }},
{%-     endfor %}
};
{%    elif type.type == 'record' %}
{% include "_record_declare_cpp.j2" %}
{%    endif %}
{%- endfor %}
{%- for type_name, type in types.items() -%}
{%    set c_type_name = type_name | camel -%}
{%    if type.type == 'array' %}
{%      if types[type.value_type].type == "union" -%}
{%        set cons = "construct_" + type.value_type -%}
{%      else -%}
{%        set cons = type.value_type | camel + "::acquire" -%}
{%      endif -%}
auto construct_{{ type_name }}(clingo_ast_t **ast, size_t size) -> {{ c_type_name }} {
    {{ c_type_name }} ret;
    try {
        ret.reserve(size);
        std::for_each_n(ast, size, [&ret](auto &arg){
            auto tmp = arg;
            arg = nullptr;
            ret.emplace_back({{ cons }}(tmp));
        });
        clingo_ast_array_free(ast, size);
    }
    catch (...) {
        clingo_ast_array_free(ast, size);
        throw;
    }
    return ret;
}
{%    elif type.type == 'union' %}
auto construct_{{ type_name }}(clingo_ast_t *ast) -> {{ c_type_name }} {
    clingo_ast_type_t type = 0;
    if (!clingo_ast_get_type(ast, &type)) {
        clingo_ast_free(ast);
        raise_error();
    }
    switch (type) {
{%-     for sub_type in type.types | flatten_types(types) %}
        case clingo_ast_type_{{sub_type}}: {
            return {{ sub_type | camel }}::acquire(ast);
        }
{%-     endfor %}
        default: {
            clingo_ast_free(ast);
            throw std::runtime_error("unexpected ast type");
        }
    }
}
{%    elif type.type == 'record' %}
{%- include "_record_define_cpp.j2" %}
{%    endif %}
{%- endfor %}
template <class T>
auto c_cast(std::optional<T> const &opt) -> clingo_ast_t* {
    if (opt) {
        return c_cast(*opt);
    }
    return nullptr;
}

template <class... Ts>
auto c_cast(std::variant<Ts...> const &var) -> clingo_ast_t* {
    return std::visit([](auto const &x) { return c_cast(x); }, var);
}

template <class T>
auto c_cast(std::vector<T> const &arr) -> std::vector<clingo_ast_t*> {
    std::vector<clingo_ast_t*> ret;
    ret.reserve(arr.size());
    for (auto const &x : arr)  {
        ret.emplace_back(c_cast(x));
    }
    return ret;
}

template <class Cons>
void visit_array(clingo_ast_t *ast, clingo_ast_attribute_t attr, py::handle visitor, py::args const &args, py::kwargs const &kwargs, Cons cons) {
    struct Array {
        ~Array() { clingo_ast_array_free(begin, size); }
        clingo_ast_t **begin = nullptr;
        size_t size = 0;
    };
    auto array = Array{};
    handle_error(clingo_ast_attribute_get_ast_array(ast, attr, &array.begin, &array.size));
    std::for_each_n(array.begin, array.size, [&](auto *&cld) {
        auto *cpy = cld;
        cld = nullptr;
        visitor(std::invoke(cons, cpy), *args, **kwargs);
    });
}

template <class Array>
auto transform_array(Array arr, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Array, bool> {
    bool changed = false;
    for (auto &elem : arr) {
        if (auto res = transform(elem, *args, **kwargs); !res.is_none()) {
            elem = py::cast<typename Array::value_type>(std::move(res));
            changed = true;
        }
    }
    return {std::move(arr), changed};
}

template <class Value>
auto transform_value(Value val, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Value, bool> {
    auto res = transform(val, *args, **kwargs);
    return {res.is_none() ? std::move(val) : py::cast<Value>(res), !res.is_none()};
}

template <class Value>
auto transform_opt_value(Value opt, py::handle transform, py::args const &args, py::kwargs const &kwargs) -> std::pair<Value, bool> {
    if (opt) {
        auto res = transform(*opt, *args, **kwargs);
        return {res.is_none() ? std::move(opt) : py::cast<Value>(res), !res.is_none()};
    }
    return {std::nullopt, false};
}

namespace {

struct CString {
    CString(std::string &&str) : str_{std::move(str)} { }
    operator std::string_view() const {
        return std::string_view{str_};
    }
    std::string str_;
};

void add_union(py::module_ &mod, char const *name, std::initializer_list<char const *> names) {
    mod.add_object(name, std::accumulate(std::next(names.begin()), names.end(), mod.attr(*names.begin()),
                                         [&](py::object const &a, const char *t) { return a | mod.attr(t); }));
};

}

template <class T, class F, class M>
auto update_value(F *self, M fun, py::kwargs const &kwargs, char const *attr) -> update_result_t<T> {
    if constexpr (std::is_same_v<T, std::string_view>) {
        if (kwargs.contains(attr)) {
            return py::cast<std::string>(kwargs[attr]);
        }
    } else {
        if (kwargs.contains(attr)) {
            return py::cast<T>(kwargs[attr]);
        }
    }
    if constexpr (std::is_same_v<T, StringArray>) {
        auto arr = (self->*fun)();
        return StringArray{arr.begin(), arr.end()};
    }
    else {
        return (self->*fun)();
    }
}

auto parse_term(Library &lib, std::string_view string) -> Term {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_term, string.data(), string.size(), &ast));
    return construct_term(ast);
}

auto parse_theory_term(Library &lib, std::string_view string) -> TheoryTerm {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_theory_term, string.data(), string.size(), &ast));
    return construct_theory_term(ast);
}

auto parse_literal(Library &lib, std::string_view string) -> Literal {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_literal, string.data(), string.size(), &ast));
    return construct_literal(ast);
}

auto parse_head_literal(Library &lib, std::string_view string) -> HeadLiteral {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_head_literal, string.data(), string.size(), &ast));
    return construct_head_literal(ast);
}

auto parse_body_literal(Library &lib, std::string_view string) -> BodyLiteral {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_body_literal, string.data(), string.size(), &ast));
    return construct_body_literal(ast);
}

auto parse_statement(Library &lib, std::string_view string) -> Statement {
    clingo_ast_t *ast = nullptr;
    handle_error(clingo_ast_parse_expression(lib, clingo_ast_parse_type_statement, string.data(), string.size(), &ast));
    return construct_statement(ast);
}

void parse_string(Library const &lib, std::string_view program, std::function<void(Statement)> callback,
                  TypeHint<"clingo.control.Control|None"> const &ctl) {
    auto *ptr = !ctl.is_none() ? ctl.cast<Control *>()->c_ptr() : nullptr;
    handle_error(clingo_ast_parse_string(
        lib, program.data(), program.size(), ptr,
        [](clingo_ast_t *ast, void *data) {
            CLINGO_TRY {
                auto &callback = *static_cast<std::function<void(Statement)> *>(data);
                clingo_ast_t *copy = nullptr;
                handle_error(clingo_ast_copy(ast, &copy));
                callback(construct_statement(copy));
            }
            CLINGO_CATCH;
        },
        &callback));
}

inline void parse_files(Library const &lib, std::span<std::string_view const> files,
                        std::function<void(Statement)> callback, TypeHint<"clingo.control.Control|None"> const &ctl) {
    auto *ptr = !ctl.is_none() ? ctl.cast<Control *>()->c_ptr() : nullptr;
    auto cfiles = transform_vec(files, [](auto const &x) { return clingo_string_t{x.data(), x.size()}; });
    handle_error(clingo_ast_parse_files(
        lib, cfiles.data(), cfiles.size(), ptr,
        [](clingo_ast_t *ast, void *data) {
            CLINGO_TRY {
                auto &callback = *static_cast<std::function<void(Statement)> *>(data);
                clingo_ast_t *copy = nullptr;
                handle_error(clingo_ast_copy(ast, &copy));
                callback(construct_statement(copy));
            }
            CLINGO_CATCH;
        },
        &callback));
}

class RewriteContext {
  public:
    RewriteContext(Library &lib) { handle_error(clingo_ast_rewrite_context_create(lib, &ctx_)); }
    ~RewriteContext() { clingo_ast_rewrite_context_free(ctx_); }
    void set_project_mode(ProjectionMode value) {
        clingo_ast_rewrite_context_set_project_mode(ctx_, static_cast<clingo_projection_mode_t>(value));
    }
    auto get_project_mode() -> ProjectionMode {
        return static_cast<ProjectionMode>(clingo_ast_rewrite_context_get_project_mode(ctx_));
    }
    void set_project_anoymous(bool value) { clingo_ast_rewrite_context_set_project_anonymous(ctx_, value); }
    auto get_project_anonymous() -> bool { return clingo_ast_rewrite_context_get_project_mode(ctx_) != 0; }
    void add_param(std::string const &name) {
        handle_error(clingo_ast_rewrite_context_add_param(ctx_, name.data(), name.size()));
    }
    void clear_params() {
        clingo_ast_rewrite_context_clear_params(ctx_);
    }
    void add_theory(StatementTheory const &stm) {
        handle_error(clingo_ast_rewrite_context_add_theory(ctx_, c_cast(stm)));
    }
    friend auto c_cast(RewriteContext const &x) -> clingo_ast_rewrite_context_t * { return x.ctx_; }

  private:
    clingo_ast_rewrite_context_t *ctx_ = nullptr;
};

auto rewrite_statement(RewriteContext &ctx, Statement &stm) -> std::vector<Statement> {
    auto *c_ctx = c_cast(ctx);
    struct Array {
        ~Array() { clingo_ast_array_free(result, result_size); }
        clingo_ast_t **result = nullptr;
        size_t result_size = 0;
    };
    auto arr = Array{};
    handle_error(clingo_ast_rewrite(c_ctx, c_cast(stm), &arr.result, &arr.result_size));
    std::vector<Statement> res;
    res.reserve(arr.result_size);
    std::for_each_n(arr.result, arr.result_size, [&res](clingo_ast *&ast) {
        auto *cpy = ast;
        ast = nullptr;
        res.emplace_back(construct_statement(cpy));
    });
    return res;
}

Program::Program(Library &lib) {
    clingo_program_t *prg = nullptr;
    handle_error(clingo_program_new(lib, &prg));
    prg_.reset(prg);
}

void Program::free(clingo_program_t *prg) noexcept {
    clingo_program_free(prg);
}

void add(Program &prg, Statement &stm) {
    handle_error(clingo_program_add(prg, c_cast(stm)));
}

template <typename Variant> struct pointer_variant;
template <typename... Ts> struct pointer_variant<std::variant<Ts...>> {
    using type = std::variant<std::add_pointer_t<Ts>...>;
};
template <typename Variant> using pointer_variant_t = typename pointer_variant<Variant>::type;

auto convert_stm(py::handle hnd) -> clingo_ast_t * {
    return c_cast(hnd.cast<pointer_variant_t<Statement>>());
}

auto convert_stm(clingo_ast_t *ast) -> py::object {
    clingo_ast_t *copy = nullptr;
    handle_error(clingo_ast_copy(ast, &copy));
    return py::cast(construct_statement(copy));
}

void register_ast(pybind11::module &m) {
    auto ast = m.def_submodule("ast", R"doc(This module provides functions to work with Abstract Syntax Trees of logic programs.

# Examples

The following example shows how to parse individual statements and add them to a control object.

```python
>>> from clingo.core import Library
>>> from clingo.control import Control
>>> from clingo import ast
>>>
>>> lib = Library()
>>> ctl = Control(lib, ["--mode=ground"])
>>> ctl.parse_string("a(1).")
>>> prg = ast.Program(lib)
>>> prg.add(ast.parse_statement(lib, "b(X+1) :- a(X)."))
>>> prg.add(ast.parse_statement(lib, "c(X+1) :- b(X)."))
>>> ctl.join(prg)
>>> ctl.ground()
>>> ctl.buffer
"""
a(1).
b(2).
c(3).
#show.
#show a/1.
#show b/1.
#show c/1.
"""
```

The next example hows how to visit and AST and collect variables.

```python
from functools import singledispatch

from clingo.core import Library
from clingo import ast

@singledispatch
def collect(stm, vars):
    """
    Collect all variables occurring in statements.
    """
    return stm.visit(collect, vars)

@collect.register
def _(var: ast.TermVariable | ast.TheoryTermVariable, vars):
    vars.add(var.name)

lib = Library()
vars = set()

ast.parse_string(
    lib,
    """\
    a(X) :- b(X,Y).
    b(Z).
    """,
    lambda stm: collect(stm, vars),
)

print(vars)
```

The last example shows how to use a transformer to modify an AST.

```python
from functools import singledispatch

from clingo.core import Library
from clingo import ast

@singledispatch
def rename(stm, lib):
    """
    Replaces all occurrences of variable `X` with `Y` in the given statement
    returning a new statement. If no replacement was made, `None` is returned.
    """
    return stm.transform(lib, rename, lib)

@rename.register
def _(var: ast.TermVariable, lib):
    return var.update(lib, name="Y") if var.name == "X" else None

lib = Library()
stms = []

ast.parse_string(
    lib,
    """\
    b(1,1).
    b(2,3).
    a(X) :- b(X,Y).
    """,
    stms.append,
)

for stm in stms:
    print(rename(stm, lib) or stm)
```
)doc");

    ast.def("_type_info_yaml", &clingo_ast_type_info_yaml, R"doc(Return a yaml description of the AST.

This can be used to auto-generate most of the binding.)doc");

    py::native_enum<ProjectionMode>(ast, "ProjectionMode", "enum.IntEnum", R"doc(Available projection modes.)doc")
        .value("Disabled", ProjectionMode::Disabled, R"doc(Do not project.)doc")
        .value("Anonymous", ProjectionMode::Anonymous, R"doc(Only project anonymous variables.)doc")
        .value("Pure", ProjectionMode::Pure, R"doc(Project pure variables.)doc")
        .finalize();
{%  for type_name, type in types.items() -%}
{%    set c_type_name = type_name | camel -%}
{%    if type.type == 'enum' %}
    auto py_{{ type_name }} = py::native_enum<{{ c_type_name }}>(ast, "{{ c_type_name }}", "enum.IntEnum", R"doc({{ type.doc }})doc");
{%    elif type.type == 'record' %}
    auto py_{{ type_name }} = py::class_<{{ c_type_name }}>(ast, "{{ c_type_name }}", R"doc({{ type.doc | doc }})doc");
{%    endif %}
{%- endfor %}
{%  for type_name, type in types.items() -%}
{%    set c_type_name = type_name | camel -%}
{%    if type.type == 'enum' %}
    py_{{ type_name }}
{%-     for value_name, value_type in type["values"].items() %}
        .value("{{ value_name | camel }}", {{ c_type_name }}::{{ value_name | camel }}, R"doc({{ value_type.doc }})doc")
{%-       if loop.last %}
        .finalize();{% endif -%}
{%-     endfor %}
{%    elif type.type == 'record' %}
    {{ type_name | comp }}(py_{{ type_name }})
        .def(py::init(&{{ c_type_name }}::construct), py::arg("lib")
        {%- for argument_name, argument in type.arguments.items() -%}
            , py::arg("{{ argument_name }}")
            {%- if "default" in argument %}
              {%- if argument.type == "bool" %} = {{ argument.default | string | lower }}
              {%- elif types[argument.type].type == "array" and argument.default == "empty" %} = {{ argument.type | camel }}{}
              {%- elif types[argument.type].type == "optional" and argument.default == "empty" %} = std::nullopt
              {%- endif %}
            {%- endif %}
        {%- endfor %}, R"doc(Construct a {{ c_type_name }} object.

Args:
    lib: The library object for storing symbols.
{%-     for argument_name, argument in type.arguments.items() %}
    {{ argument_name }}: {{ argument.doc | param_doc }}
{%-     endfor %})doc")
        .def("__str__", &{{ c_type_name }}::to_string)
{%-     for argument_name, argument in type.arguments.items() %}
        .def_property_readonly("{{argument_name}}", &{{ c_type_name }}::{{argument_name}}, R"doc({{ argument.doc }})doc")
{%-     endfor %}
        .def("visit", &{{ c_type_name }}::visit, py::arg("visitor"), R"doc(Visit the children of the expression.

Args:
    visitor: The visitor accepting the sub expressions.
)doc")
        .def("transform", &{{ c_type_name }}::transform, py::arg("lib"), py::arg("transformer"), R"doc(Transform the expression.

Additional arguments are passed to the transformer.

Args:
    lib: The library object for storing symbols.
    transformer: The transformer accepting the sub expressions.
Returns:
    The transformed object or None.
)doc")
        .def("update", &{{ c_type_name }}::update, py::arg("lib"), R"doc(Update the expression.

Accepts keyword arguments with attributes to update.

Args:
    lib: The library object for storing symbols.
Returns:
    The updated object.
)doc");
{%    endif %}
{%- endfor %}
    ast.def("parse_string", &parse_string, py::arg("lib"), py::arg("program"), py::arg("callback"),
            py::arg("control") = std::nullopt,
            R"doc(Parse the program in the given string.

Args:
    lib: A library object to store symbols.
    program: The program to parse.
    callback: Function to report statements.
    control: Optional Control object to handle ASPIF.
)doc");
    ast.def("parse_files", &parse_files, py::arg("lib"), py::arg("files"), py::arg("callback"),
            py::arg("control") = std::nullopt, R"(
Parse the program in the given files.

The parser follows clingo's handling of files on the command line. Filename
"-" is treated as "STDIN" and if an empty list is given, then the parser will
read from "STDIN".

Args:
    lib: A library object to store symbols.
    files: The files to parse.
    callback: Function to report statements.
    control: Optional Control object to handle ASPIF.
)");
    ast.def("parse_term", &parse_term, py::arg("lib"), py::arg("string"), R"doc(Parse a term.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed Term object.)doc");
    ast.def("parse_theory_term", &parse_theory_term, py::arg("lib"), py::arg("string"), R"doc(Parse a theory term.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed TheoryTerm object.)doc");
    ast.def("parse_literal", &parse_literal, py::arg("lib"), py::arg("string"), R"doc(Parse a literal.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed Literal object.)doc");
    ast.def("parse_head_literal", &parse_head_literal, py::arg("lib"), py::arg("string"), R"doc(Parse a head literal.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed HeadLiteral object.)doc");
    ast.def("parse_body_literal", &parse_body_literal, py::arg("lib"), py::arg("string"), R"doc(Parse a body literal.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed BodyLiteral object.)doc");
    ast.def("parse_statement", &parse_statement, py::arg("lib"), py::arg("string"), R"doc(Parse a statement.

Args:
    lib: The library object for storing symbols.
    string: The string to parse.

Returns:
    The parsed Statement object.)doc");

    py::class_<RewriteContext>(ast, "RewriteContext", R"doc(Context to rewrite statements.)doc")
        .def(py::init<Library &>(), py::arg("lib"),
             R"doc(Create a context to rewrite statements.

Args:
    lib: A library object to store symbols.
)doc")
        .def("add_param", &RewriteContext::add_param, py::arg("name"), R"doc(
Add a parameter.

Parameters are protected from simplification.

Args:
    name: The name of the parameter.
)doc")
        .def("clear_params", &RewriteContext::clear_params, R"doc(
Remove previously added params
)doc")
        .def("add_theory", &RewriteContext::add_theory, py::arg("theory"), R"doc(
Add a theory definition statement.

The theory definition is used to rewrite theory atoms in statements.

Args:
    theory: The theory statement to add.
)doc")
        .def_property("project_mode", &RewriteContext::get_project_mode, &RewriteContext::set_project_mode,
                      R"doc(The active projection mode.)doc")
        .def_property("project_anonymous", &RewriteContext::get_project_anonymous,
                      &RewriteContext::set_project_anoymous,
                      R"doc(Whether to project anonymous variables in negative literals.)doc");

    ast.def("rewrite_statement", &rewrite_statement, py::arg("ctx"), py::arg("statement"),
            R"doc(Simplify the given statement.

Args:
    ctx: The rewrite context.
    statement: The statement to rewrite.

Returns:
    A list of rewritten statements.)doc");

    py::class_<Program>(ast, "Program", R"doc(A non-ground program.)doc")
        .def(py::init<Library &>(), py::arg("lib"), R"doc(Create an empty non-ground program.

Args:
    lib: A library object to store symbols.
)doc")
        .def("add", &add, py::arg("statement"), R"doc(Add a statement to a program.

Args:
    statement: The statement to add.)doc");
{%  for type_name, type in types.items() -%}
{%    if type.type == 'union' %}
        add_union(ast, "{{ type_name | camel }}", {
{%-     for name in type["types"] -%}
          {%- if not loop.first %}, {% endif %}"{{ name | camel }}"
{%-     endfor -%}
});
{%-    endif -%}
{%- endfor %}
}

}

// NOLINTEND(readability-convert-member-functions-to-static,performance-enum-size)
