localize_mf2_treesitter

/*
 * Localize.Mf2.TreeSitter NIF.
 *
 * Thin glue between Elixir and the vendored tree-sitter runtime + the
 * tree-sitter-mf2 grammar. The design is deliberately minimal: TSTree
 * is held as a NIF resource, TSNode is serialised as a binary and
 * carried alongside the resource on the Erlang side so it stays rooted
 * for as long as any node referring to it is alive.
 */

#include <erl_nif.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

#include "tree_sitter/api.h"

/* Declared by the generated grammar parser.c. */
const TSLanguage *tree_sitter_mf2(void);

/* ----- Module-local state --------------------------------------- */

static ErlNifResourceType *TREE_RESOURCE_TYPE = NULL;
static ErlNifResourceType *QUERY_RESOURCE_TYPE = NULL;
static const TSLanguage *MF2_LANGUAGE = NULL;

static ERL_NIF_TERM ATOM_OK;
static ERL_NIF_TERM ATOM_ERROR;
static ERL_NIF_TERM ATOM_NIL;
static ERL_NIF_TERM ATOM_TRUE;
static ERL_NIF_TERM ATOM_FALSE;
static ERL_NIF_TERM ATOM_NODE;
static ERL_NIF_TERM ATOM_PARSE_FAILED;
static ERL_NIF_TERM ATOM_QUERY_ERROR;
static ERL_NIF_TERM ATOM_SYNTAX;
static ERL_NIF_TERM ATOM_NODE_TYPE;
static ERL_NIF_TERM ATOM_FIELD;
static ERL_NIF_TERM ATOM_CAPTURE;
static ERL_NIF_TERM ATOM_STRUCTURE;
static ERL_NIF_TERM ATOM_LANGUAGE;
static ERL_NIF_TERM ATOM_UNKNOWN;
static ERL_NIF_TERM ATOM_PATTERN_INDEX;
static ERL_NIF_TERM ATOM_CAPTURES;

/* ----- Resource wrapper ----------------------------------------- */

typedef struct {
    TSTree *tree;
} TreeResource;

typedef struct {
    TSQuery *query;
} QueryResource;

static void tree_resource_dtor(ErlNifEnv *env, void *obj)
{
    (void)env;
    TreeResource *res = (TreeResource *)obj;
    if (res->tree) {
        ts_tree_delete(res->tree);
        res->tree = NULL;
    }
}

static void query_resource_dtor(ErlNifEnv *env, void *obj)
{
    (void)env;
    QueryResource *res = (QueryResource *)obj;
    if (res->query) {
        ts_query_delete(res->query);
        res->query = NULL;
    }
}

/* ----- Helpers --------------------------------------------------- */

static ERL_NIF_TERM make_atom(ErlNifEnv *env, const char *name)
{
    ERL_NIF_TERM atom;
    if (enif_make_existing_atom(env, name, &atom, ERL_NIF_LATIN1)) {
        return atom;
    }
    return enif_make_atom(env, name);
}

static ERL_NIF_TERM make_binary_from_cstr(ErlNifEnv *env, const char *s, size_t len)
{
    ERL_NIF_TERM bin_term;
    unsigned char *data = enif_make_new_binary(env, len, &bin_term);
    memcpy(data, s, len);
    return bin_term;
}

/*
 * Encode a TSNode as {:node, tree_resource_term, node_bytes_binary}.
 *
 * The tree_resource_term is passed through verbatim so the GC keeps
 * the TSTree alive as long as any encoded TSNode is reachable from
 * Elixir.
 */
static ERL_NIF_TERM encode_node(ErlNifEnv *env, ERL_NIF_TERM tree_term, TSNode node)
{
    ERL_NIF_TERM bin_term;
    unsigned char *data = enif_make_new_binary(env, sizeof(TSNode), &bin_term);
    memcpy(data, &node, sizeof(TSNode));
    return enif_make_tuple3(env, ATOM_NODE, tree_term, bin_term);
}

/*
 * Decode a TSNode tuple and also yield the original tree_term so the
 * caller can pass it on to child-node encodings.
 */
static bool decode_node(ErlNifEnv *env, ERL_NIF_TERM term,
                        TSNode *out_node, ERL_NIF_TERM *out_tree_term)
{
    const ERL_NIF_TERM *elems;
    int arity;
    if (!enif_get_tuple(env, term, &arity, &elems) || arity != 3) return false;
    if (!enif_is_identical(elems[0], ATOM_NODE)) return false;

    ErlNifBinary bin;
    if (!enif_inspect_binary(env, elems[2], &bin)) return false;
    if (bin.size != sizeof(TSNode)) return false;

    memcpy(out_node, bin.data, sizeof(TSNode));
    if (out_tree_term) *out_tree_term = elems[1];
    return true;
}

/* ----- NIF implementations -------------------------------------- */

static ERL_NIF_TERM nif_parse(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    ErlNifBinary source;
    if (!enif_inspect_binary(env, argv[0], &source)) {
        return enif_make_badarg(env);
    }

    TSParser *parser = ts_parser_new();
    if (!parser) return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);

    if (!ts_parser_set_language(parser, MF2_LANGUAGE)) {
        ts_parser_delete(parser);
        return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);
    }

    TSTree *tree = ts_parser_parse_string(
        parser, NULL, (const char *)source.data, (uint32_t)source.size);

    ts_parser_delete(parser);

    if (!tree) return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);

    TreeResource *res = enif_alloc_resource(TREE_RESOURCE_TYPE, sizeof(TreeResource));
    res->tree = tree;

    ERL_NIF_TERM tree_term = enif_make_resource(env, res);
    enif_release_resource(res);

    return enif_make_tuple2(env, ATOM_OK, tree_term);
}

static ERL_NIF_TERM nif_tree_root_node(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TreeResource *res;
    if (!enif_get_resource(env, argv[0], TREE_RESOURCE_TYPE, (void **)&res)) {
        return enif_make_badarg(env);
    }
    TSNode root = ts_tree_root_node(res->tree);
    return encode_node(env, argv[0], root);
}

static ERL_NIF_TERM nif_tree_to_sexp(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TreeResource *res;
    if (!enif_get_resource(env, argv[0], TREE_RESOURCE_TYPE, (void **)&res)) {
        return enif_make_badarg(env);
    }
    TSNode root = ts_tree_root_node(res->tree);
    char *s = ts_node_string(root);
    if (!s) return make_binary_from_cstr(env, "", 0);
    ERL_NIF_TERM out = make_binary_from_cstr(env, s, strlen(s));
    free(s);
    return out;
}

static ERL_NIF_TERM nif_node_type(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    const char *type = ts_node_type(node);
    return make_binary_from_cstr(env, type, strlen(type));
}

static ERL_NIF_TERM nif_node_start_byte(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return enif_make_uint(env, ts_node_start_byte(node));
}

static ERL_NIF_TERM nif_node_end_byte(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return enif_make_uint(env, ts_node_end_byte(node));
}

static ERL_NIF_TERM make_point(ErlNifEnv *env, TSPoint p)
{
    return enif_make_tuple2(env, enif_make_uint(env, p.row), enif_make_uint(env, p.column));
}

static ERL_NIF_TERM nif_node_start_point(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return make_point(env, ts_node_start_point(node));
}

static ERL_NIF_TERM nif_node_end_point(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return make_point(env, ts_node_end_point(node));
}

static ERL_NIF_TERM nif_node_is_named(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return ts_node_is_named(node) ? ATOM_TRUE : ATOM_FALSE;
}

static ERL_NIF_TERM nif_node_has_error(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return ts_node_has_error(node) ? ATOM_TRUE : ATOM_FALSE;
}

static ERL_NIF_TERM nif_node_child_count(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return enif_make_uint(env, ts_node_child_count(node));
}

static ERL_NIF_TERM nif_node_named_child_count(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return enif_make_uint(env, ts_node_named_child_count(node));
}

static ERL_NIF_TERM nif_node_child(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    unsigned int index;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[1], &index)) return enif_make_badarg(env);
    if (index >= ts_node_child_count(node)) return ATOM_NIL;
    TSNode child = ts_node_child(node, index);
    if (ts_node_is_null(child)) return ATOM_NIL;
    return encode_node(env, tree_term, child);
}

static ERL_NIF_TERM nif_node_named_child(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    unsigned int index;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[1], &index)) return enif_make_badarg(env);
    if (index >= ts_node_named_child_count(node)) return ATOM_NIL;
    TSNode child = ts_node_named_child(node, index);
    if (ts_node_is_null(child)) return ATOM_NIL;
    return encode_node(env, tree_term, child);
}

/*
 * Return the parent / sibling / descendant or `:nil` if the target
 * TSNode is null. Each wraps a single libtree-sitter call and carries
 * the tree-term through so the GC keeps the owning tree alive.
 */

static ERL_NIF_TERM nif_node_parent(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    TSNode parent = ts_node_parent(node);
    if (ts_node_is_null(parent)) return ATOM_NIL;
    return encode_node(env, tree_term, parent);
}

static ERL_NIF_TERM nif_node_next_sibling(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    TSNode sib = ts_node_next_sibling(node);
    if (ts_node_is_null(sib)) return ATOM_NIL;
    return encode_node(env, tree_term, sib);
}

static ERL_NIF_TERM nif_node_prev_sibling(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    TSNode sib = ts_node_prev_sibling(node);
    if (ts_node_is_null(sib)) return ATOM_NIL;
    return encode_node(env, tree_term, sib);
}

static ERL_NIF_TERM nif_node_next_named_sibling(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    TSNode sib = ts_node_next_named_sibling(node);
    if (ts_node_is_null(sib)) return ATOM_NIL;
    return encode_node(env, tree_term, sib);
}

static ERL_NIF_TERM nif_node_prev_named_sibling(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    TSNode sib = ts_node_prev_named_sibling(node);
    if (ts_node_is_null(sib)) return ATOM_NIL;
    return encode_node(env, tree_term, sib);
}

static ERL_NIF_TERM nif_node_descendant_for_byte_range(
    ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    unsigned int start_byte, end_byte;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[1], &start_byte)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[2], &end_byte)) return enif_make_badarg(env);
    TSNode found = ts_node_descendant_for_byte_range(node, start_byte, end_byte);
    if (ts_node_is_null(found)) return ATOM_NIL;
    return encode_node(env, tree_term, found);
}

static ERL_NIF_TERM nif_node_named_descendant_for_byte_range(
    ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    ERL_NIF_TERM tree_term;
    unsigned int start_byte, end_byte;
    if (!decode_node(env, argv[0], &node, &tree_term)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[1], &start_byte)) return enif_make_badarg(env);
    if (!enif_get_uint(env, argv[2], &end_byte)) return enif_make_badarg(env);
    TSNode found = ts_node_named_descendant_for_byte_range(node, start_byte, end_byte);
    if (ts_node_is_null(found)) return ATOM_NIL;
    return encode_node(env, tree_term, found);
}

static ERL_NIF_TERM nif_node_is_error(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return ts_node_is_error(node) ? ATOM_TRUE : ATOM_FALSE;
}

static ERL_NIF_TERM nif_node_is_missing(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    return ts_node_is_missing(node) ? ATOM_TRUE : ATOM_FALSE;
}

static ERL_NIF_TERM nif_node_to_sexp(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TSNode node;
    if (!decode_node(env, argv[0], &node, NULL)) return enif_make_badarg(env);
    char *s = ts_node_string(node);
    if (!s) return make_binary_from_cstr(env, "", 0);
    ERL_NIF_TERM out = make_binary_from_cstr(env, s, strlen(s));
    free(s);
    return out;
}

/* ----- Incremental parse ---------------------------------------- */

/*
 * Decode an edit tuple of the form
 *   {start_byte, old_end_byte, new_end_byte,
 *    {start_row, start_col}, {old_end_row, old_end_col},
 *    {new_end_row, new_end_col}}.
 */
static bool decode_edit(ErlNifEnv *env, ERL_NIF_TERM term, TSInputEdit *out)
{
    const ERL_NIF_TERM *elems;
    int arity;
    if (!enif_get_tuple(env, term, &arity, &elems) || arity != 6) return false;

    unsigned int start_byte, old_end_byte, new_end_byte;
    if (!enif_get_uint(env, elems[0], &start_byte)) return false;
    if (!enif_get_uint(env, elems[1], &old_end_byte)) return false;
    if (!enif_get_uint(env, elems[2], &new_end_byte)) return false;

    const ERL_NIF_TERM *pt;
    int pt_arity;
    unsigned int r, c;

    if (!enif_get_tuple(env, elems[3], &pt_arity, &pt) || pt_arity != 2) return false;
    if (!enif_get_uint(env, pt[0], &r) || !enif_get_uint(env, pt[1], &c)) return false;
    out->start_point.row = r; out->start_point.column = c;

    if (!enif_get_tuple(env, elems[4], &pt_arity, &pt) || pt_arity != 2) return false;
    if (!enif_get_uint(env, pt[0], &r) || !enif_get_uint(env, pt[1], &c)) return false;
    out->old_end_point.row = r; out->old_end_point.column = c;

    if (!enif_get_tuple(env, elems[5], &pt_arity, &pt) || pt_arity != 2) return false;
    if (!enif_get_uint(env, pt[0], &r) || !enif_get_uint(env, pt[1], &c)) return false;
    out->new_end_point.row = r; out->new_end_point.column = c;

    out->start_byte = start_byte;
    out->old_end_byte = old_end_byte;
    out->new_end_byte = new_end_byte;
    return true;
}

static ERL_NIF_TERM nif_parse_incremental(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TreeResource *old_res;
    if (!enif_get_resource(env, argv[0], TREE_RESOURCE_TYPE, (void **)&old_res)) {
        return enif_make_badarg(env);
    }

    ErlNifBinary source;
    if (!enif_inspect_binary(env, argv[2], &source)) return enif_make_badarg(env);

    /*
     * Work on a copy of the old tree so the caller can still use it
     * against the new tree (e.g. for `ts_tree_get_changed_ranges`).
     * `ts_tree_edit` mutates in place, so we must not mutate the
     * caller's resource directly.
     */
    TSTree *edited = ts_tree_copy(old_res->tree);
    if (!edited) return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);

    ERL_NIF_TERM edits_list = argv[1];
    ERL_NIF_TERM head, tail = edits_list;
    while (enif_get_list_cell(env, tail, &head, &tail)) {
        TSInputEdit edit;
        if (!decode_edit(env, head, &edit)) {
            ts_tree_delete(edited);
            return enif_make_badarg(env);
        }
        ts_tree_edit(edited, &edit);
    }

    TSParser *parser = ts_parser_new();
    if (!parser) {
        ts_tree_delete(edited);
        return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);
    }
    if (!ts_parser_set_language(parser, MF2_LANGUAGE)) {
        ts_parser_delete(parser);
        ts_tree_delete(edited);
        return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);
    }

    TSTree *new_tree = ts_parser_parse_string(
        parser, edited,
        (const char *)source.data, (uint32_t)source.size);

    ts_parser_delete(parser);
    ts_tree_delete(edited);

    if (!new_tree) return enif_make_tuple2(env, ATOM_ERROR, ATOM_PARSE_FAILED);

    TreeResource *res = enif_alloc_resource(TREE_RESOURCE_TYPE, sizeof(TreeResource));
    res->tree = new_tree;
    ERL_NIF_TERM tree_term = enif_make_resource(env, res);
    enif_release_resource(res);
    return enif_make_tuple2(env, ATOM_OK, tree_term);
}

/*
 * Return the byte/point ranges that differ between two trees parsed
 * from the same source (old tree edited to match the new text, then
 * reparsed). List of `{start_byte, end_byte, start_point, end_point}`.
 */
static ERL_NIF_TERM nif_tree_get_changed_ranges(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    TreeResource *old_res, *new_res;
    if (!enif_get_resource(env, argv[0], TREE_RESOURCE_TYPE, (void **)&old_res)) {
        return enif_make_badarg(env);
    }
    if (!enif_get_resource(env, argv[1], TREE_RESOURCE_TYPE, (void **)&new_res)) {
        return enif_make_badarg(env);
    }

    uint32_t count = 0;
    TSRange *ranges = ts_tree_get_changed_ranges(old_res->tree, new_res->tree, &count);

    ERL_NIF_TERM list = enif_make_list(env, 0);
    for (int32_t i = (int32_t)count - 1; i >= 0; i--) {
        const TSRange *r = &ranges[i];
        ERL_NIF_TERM tup = enif_make_tuple4(
            env,
            enif_make_uint(env, r->start_byte),
            enif_make_uint(env, r->end_byte),
            make_point(env, r->start_point),
            make_point(env, r->end_point));
        list = enif_make_list_cell(env, tup, list);
    }
    if (ranges) free(ranges);
    return list;
}

/* ----- Query NIFs ----------------------------------------------- */

static ERL_NIF_TERM query_error_atom(TSQueryError err)
{
    switch (err) {
        case TSQueryErrorSyntax:    return ATOM_SYNTAX;
        case TSQueryErrorNodeType:  return ATOM_NODE_TYPE;
        case TSQueryErrorField:     return ATOM_FIELD;
        case TSQueryErrorCapture:   return ATOM_CAPTURE;
        case TSQueryErrorStructure: return ATOM_STRUCTURE;
        case TSQueryErrorLanguage:  return ATOM_LANGUAGE;
        default:                    return ATOM_UNKNOWN;
    }
}

static ERL_NIF_TERM nif_query_new(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    ErlNifBinary source;
    if (!enif_inspect_binary(env, argv[0], &source)) return enif_make_badarg(env);

    uint32_t err_offset = 0;
    TSQueryError err_type = TSQueryErrorNone;

    TSQuery *query = ts_query_new(
        MF2_LANGUAGE,
        (const char *)source.data,
        (uint32_t)source.size,
        &err_offset,
        &err_type);

    if (!query) {
        ERL_NIF_TERM reason = enif_make_tuple3(
            env,
            ATOM_QUERY_ERROR,
            enif_make_uint(env, err_offset),
            query_error_atom(err_type));
        return enif_make_tuple2(env, ATOM_ERROR, reason);
    }

    QueryResource *res = enif_alloc_resource(QUERY_RESOURCE_TYPE, sizeof(QueryResource));
    res->query = query;
    ERL_NIF_TERM term = enif_make_resource(env, res);
    enif_release_resource(res);
    return enif_make_tuple2(env, ATOM_OK, term);
}

static ERL_NIF_TERM nif_query_pattern_count(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    QueryResource *res;
    if (!enif_get_resource(env, argv[0], QUERY_RESOURCE_TYPE, (void **)&res)) {
        return enif_make_badarg(env);
    }
    return enif_make_uint(env, ts_query_pattern_count(res->query));
}

static ERL_NIF_TERM nif_query_capture_count(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    QueryResource *res;
    if (!enif_get_resource(env, argv[0], QUERY_RESOURCE_TYPE, (void **)&res)) {
        return enif_make_badarg(env);
    }
    return enif_make_uint(env, ts_query_capture_count(res->query));
}

static ERL_NIF_TERM nif_query_capture_names(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    QueryResource *res;
    if (!enif_get_resource(env, argv[0], QUERY_RESOURCE_TYPE, (void **)&res)) {
        return enif_make_badarg(env);
    }

    uint32_t count = ts_query_capture_count(res->query);
    ERL_NIF_TERM list = enif_make_list(env, 0);
    for (int32_t i = (int32_t)count - 1; i >= 0; i--) {
        uint32_t len = 0;
        const char *name = ts_query_capture_name_for_id(res->query, (uint32_t)i, &len);
        ERL_NIF_TERM bin = make_binary_from_cstr(env, name ? name : "", len);
        list = enif_make_list_cell(env, bin, list);
    }
    return list;
}

/*
 * Execute a query on a node and return all matches as a list of
 * `%{pattern_index: i, captures: [{name, node_tuple}]}` maps.
 *
 * For each match, captures preserve the order the tree-sitter cursor
 * yields them. For typical highlight queries the list is small.
 */
static ERL_NIF_TERM nif_query_matches(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    QueryResource *qres;
    if (!enif_get_resource(env, argv[0], QUERY_RESOURCE_TYPE, (void **)&qres)) {
        return enif_make_badarg(env);
    }
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[1], &node, &tree_term)) return enif_make_badarg(env);

    TSQueryCursor *cursor = ts_query_cursor_new();
    if (!cursor) return enif_make_list(env, 0);

    ts_query_cursor_exec(cursor, qres->query, node);

    /*
     * Collect matches in forward order. We accumulate reversed and
     * reverse at the end so that the returned list is in match-order.
     */
    ERL_NIF_TERM matches_reversed = enif_make_list(env, 0);
    TSQueryMatch match;
    while (ts_query_cursor_next_match(cursor, &match)) {
        ERL_NIF_TERM captures_reversed = enif_make_list(env, 0);
        for (uint16_t i = 0; i < match.capture_count; i++) {
            const TSQueryCapture *cap = &match.captures[i];
            uint32_t name_len = 0;
            const char *name = ts_query_capture_name_for_id(
                qres->query, cap->index, &name_len);
            ERL_NIF_TERM name_bin = make_binary_from_cstr(env, name ? name : "", name_len);
            ERL_NIF_TERM node_term = encode_node(env, tree_term, cap->node);
            ERL_NIF_TERM pair = enif_make_tuple2(env, name_bin, node_term);
            captures_reversed = enif_make_list_cell(env, pair, captures_reversed);
        }

        ERL_NIF_TERM captures;
        enif_make_reverse_list(env, captures_reversed, &captures);

        ERL_NIF_TERM match_map;
        enif_make_map_from_arrays(
            env,
            (ERL_NIF_TERM[]){ATOM_PATTERN_INDEX, ATOM_CAPTURES},
            (ERL_NIF_TERM[]){enif_make_uint(env, match.pattern_index), captures},
            2,
            &match_map);

        matches_reversed = enif_make_list_cell(env, match_map, matches_reversed);
    }

    ts_query_cursor_delete(cursor);

    ERL_NIF_TERM matches;
    enif_make_reverse_list(env, matches_reversed, &matches);
    return matches;
}

/*
 * Flat capture list — `[{name, node_tuple}]` in the order the cursor
 * yields captures. Useful for highlight queries where pattern
 * provenance is irrelevant.
 */
static ERL_NIF_TERM nif_query_captures(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    (void)argc;
    QueryResource *qres;
    if (!enif_get_resource(env, argv[0], QUERY_RESOURCE_TYPE, (void **)&qres)) {
        return enif_make_badarg(env);
    }
    TSNode node;
    ERL_NIF_TERM tree_term;
    if (!decode_node(env, argv[1], &node, &tree_term)) return enif_make_badarg(env);

    TSQueryCursor *cursor = ts_query_cursor_new();
    if (!cursor) return enif_make_list(env, 0);

    ts_query_cursor_exec(cursor, qres->query, node);

    ERL_NIF_TERM list_reversed = enif_make_list(env, 0);
    TSQueryMatch match;
    uint32_t capture_index;
    while (ts_query_cursor_next_capture(cursor, &match, &capture_index)) {
        const TSQueryCapture *cap = &match.captures[capture_index];
        uint32_t name_len = 0;
        const char *name = ts_query_capture_name_for_id(
            qres->query, cap->index, &name_len);
        ERL_NIF_TERM name_bin = make_binary_from_cstr(env, name ? name : "", name_len);
        ERL_NIF_TERM node_term = encode_node(env, tree_term, cap->node);
        ERL_NIF_TERM pair = enif_make_tuple2(env, name_bin, node_term);
        list_reversed = enif_make_list_cell(env, pair, list_reversed);
    }

    ts_query_cursor_delete(cursor);

    ERL_NIF_TERM list;
    enif_make_reverse_list(env, list_reversed, &list);
    return list;
}

/* ----- Load / register ------------------------------------------ */

static int on_load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info)
{
    (void)priv_data;
    (void)load_info;

    ErlNifResourceType *rt = enif_open_resource_type(
        env, NULL, "Localize.Mf2.TreeSitter.Tree",
        tree_resource_dtor, ERL_NIF_RT_CREATE, NULL);
    if (!rt) return -1;
    TREE_RESOURCE_TYPE = rt;

    ErlNifResourceType *qrt = enif_open_resource_type(
        env, NULL, "Localize.Mf2.TreeSitter.Query",
        query_resource_dtor, ERL_NIF_RT_CREATE, NULL);
    if (!qrt) return -1;
    QUERY_RESOURCE_TYPE = qrt;

    MF2_LANGUAGE = tree_sitter_mf2();
    if (!MF2_LANGUAGE) return -1;

    ATOM_OK = make_atom(env, "ok");
    ATOM_ERROR = make_atom(env, "error");
    ATOM_NIL = make_atom(env, "nil");
    ATOM_TRUE = make_atom(env, "true");
    ATOM_FALSE = make_atom(env, "false");
    ATOM_NODE = make_atom(env, "node");
    ATOM_PARSE_FAILED = make_atom(env, "parse_failed");
    ATOM_QUERY_ERROR = make_atom(env, "query_error");
    ATOM_SYNTAX = make_atom(env, "syntax");
    ATOM_NODE_TYPE = make_atom(env, "node_type");
    ATOM_FIELD = make_atom(env, "field");
    ATOM_CAPTURE = make_atom(env, "capture");
    ATOM_STRUCTURE = make_atom(env, "structure");
    ATOM_LANGUAGE = make_atom(env, "language");
    ATOM_UNKNOWN = make_atom(env, "unknown");
    ATOM_PATTERN_INDEX = make_atom(env, "pattern_index");
    ATOM_CAPTURES = make_atom(env, "captures");

    return 0;
}

static ErlNifFunc nif_funcs[] = {
    {"parse",                   1, nif_parse,                 0},
    {"tree_root_node",          1, nif_tree_root_node,        0},
    {"tree_to_sexp",            1, nif_tree_to_sexp,          0},
    {"node_type",               1, nif_node_type,             0},
    {"node_start_byte",         1, nif_node_start_byte,       0},
    {"node_end_byte",           1, nif_node_end_byte,         0},
    {"node_start_point",        1, nif_node_start_point,      0},
    {"node_end_point",          1, nif_node_end_point,        0},
    {"node_is_named",           1, nif_node_is_named,         0},
    {"node_has_error",          1, nif_node_has_error,        0},
    {"node_child_count",        1, nif_node_child_count,      0},
    {"node_named_child_count",  1, nif_node_named_child_count,0},
    {"node_child",              2, nif_node_child,            0},
    {"node_named_child",        2, nif_node_named_child,      0},
    {"node_to_sexp",            1, nif_node_to_sexp,          0},
    {"node_parent",             1, nif_node_parent,           0},
    {"node_next_sibling",       1, nif_node_next_sibling,     0},
    {"node_prev_sibling",       1, nif_node_prev_sibling,     0},
    {"node_next_named_sibling", 1, nif_node_next_named_sibling,0},
    {"node_prev_named_sibling", 1, nif_node_prev_named_sibling,0},
    {"node_descendant_for_byte_range",       3, nif_node_descendant_for_byte_range,       0},
    {"node_named_descendant_for_byte_range", 3, nif_node_named_descendant_for_byte_range, 0},
    {"node_is_error",           1, nif_node_is_error,         0},
    {"node_is_missing",         1, nif_node_is_missing,       0},
    {"query_new",               1, nif_query_new,             0},
    {"query_pattern_count",     1, nif_query_pattern_count,   0},
    {"query_capture_count",     1, nif_query_capture_count,   0},
    {"query_capture_names",     1, nif_query_capture_names,   0},
    {"query_matches",           2, nif_query_matches,         0},
    {"query_captures",          2, nif_query_captures,        0},
    {"parse_incremental",       3, nif_parse_incremental,     0},
    {"tree_get_changed_ranges", 2, nif_tree_get_changed_ranges,0},
};

ERL_NIF_INIT(Elixir.Localize.Mf2.TreeSitter.Nif, nif_funcs, on_load, NULL, NULL, NULL)