const std = @import("std");
const cio = @import("cio.zig");
const ContentCache = @import("hot_cache.zig").ContentCache;

// ── Inverted word index ─────────────────────────────────────
// Maps word → list of (path, line) hits. O(1) word lookup.

pub const WordHit = struct {
    doc_id: u32,
    line_num: u32,
};

pub const WordIndex = struct {
    /// word → hits
    index: std.StringHashMap(std.ArrayList(WordHit)),
    /// path → set of words contributed (for efficient re-index cleanup).
    /// WordIndex owns these path keys.
    file_words: std.StringHashMap([]const []const u8),
    allocator: std.mem.Allocator,
    skip_file_words: bool = false,
    enabled: bool = true,
    path_to_id: std.StringHashMap(u32),
    id_to_path: std.ArrayList([]const u8),
    /// doc_id → number of tokens indexed for that doc (BM25 length normalization).
    doc_lengths: std.AutoHashMap(u32, u32),
    /// Sum of all values in doc_lengths.
    total_tokens: u64 = 0,

    pub fn hitPath(self: *const WordIndex, hit: WordHit) []const u8 {
        if (hit.doc_id < self.id_to_path.items.len) return self.id_to_path.items[hit.doc_id];
        return "";
    }

    fn getOrCreateDocId(self: *WordIndex, path: []const u8) !u32 {
        if (self.path_to_id.get(path)) |id| return id;
        const id: u32 = @intCast(self.id_to_path.items.len);
        try self.id_to_path.append(self.allocator, path);
        try self.path_to_id.put(path, id);
        return id;
    }

    pub fn init(allocator: std.mem.Allocator) WordIndex {
        return .{
            .index = std.StringHashMap(std.ArrayList(WordHit)).init(allocator),
            .file_words = std.StringHashMap([]const []const u8).init(allocator),
            .allocator = allocator,
            .path_to_id = std.StringHashMap(u32).init(allocator),
            .id_to_path = .empty,
            .doc_lengths = std.AutoHashMap(u32, u32).init(allocator),
            .total_tokens = 0,
        };
    }

    pub fn deinit(self: *WordIndex) void {
        // Free hit lists and duped word keys
        var iter = self.index.iterator();
        while (iter.next()) |entry| {
            self.allocator.free(entry.key_ptr.*);
            entry.value_ptr.deinit(self.allocator);
        }
        self.index.deinit();

        // Free per-file word sets
        var fw_iter = self.file_words.iterator();
        while (fw_iter.next()) |entry| {
            self.allocator.free(entry.key_ptr.*);
            self.allocator.free(entry.value_ptr.*);
        }
        self.file_words.deinit();

        if (self.skip_file_words) {
            for (self.id_to_path.items) |path| {
                if (path.len > 0) self.allocator.free(path);
            }
        }

        self.path_to_id.deinit();
        self.id_to_path.deinit(self.allocator);
        self.doc_lengths.deinit();
    }

    /// Remove all index entries for a file (call before re-indexing).
    pub fn removeFile(self: *WordIndex, path: []const u8) void {
        const removed = self.file_words.fetchRemove(path) orelse return;
        const stable_path = removed.key;
        const words_slice = removed.value;

        const doc_id = self.path_to_id.get(stable_path) orelse {
            self.allocator.free(words_slice);
            self.allocator.free(stable_path);
            return;
        };
        _ = self.path_to_id.remove(stable_path);
        if (doc_id < self.id_to_path.items.len) {
            self.id_to_path.items[doc_id] = "";
        }
        if (self.doc_lengths.fetchRemove(doc_id)) |kv| {
            self.total_tokens -= kv.value;
        }
        defer {
            self.allocator.free(words_slice);
            self.allocator.free(stable_path);
        }

        // For each word this file contributed, remove hits with this doc_id.
        // Prune empty buckets so churn does not leak key/list entries.
        for (words_slice) |word| {
            const word_ptr = &word;
            if (self.index.getEntry(word_ptr.*)) |entry| {
                const hits = entry.value_ptr;
                var i: usize = 0;
                while (i < hits.items.len) {
                    if (hits.items[i].doc_id == doc_id) {
                        _ = hits.swapRemove(i);
                    } else {
                        i += 1;
                    }
                }
                if (hits.items.len == 0) {
                    const owned_word = entry.key_ptr.*;
                    hits.deinit(self.allocator);
                    _ = self.index.remove(word_ptr.*);
                    self.allocator.free(owned_word);
                }
            }
        }
    }
    fn indexOneToken(
        self: *WordIndex,
        token: []const u8,
        doc_id: u32,
        line_num: u32,
        words_set: *std.StringHashMap(void),
    ) !void {
        const gop = try self.index.getOrPut(token);
        if (!gop.found_existing) {
            const duped = try self.allocator.dupe(u8, token);
            gop.key_ptr.* = duped;
            gop.value_ptr.* = .empty;
        }
        if (gop.value_ptr.items.len > 0) {
            const last = gop.value_ptr.items[gop.value_ptr.items.len - 1];
            if (last.doc_id == doc_id and last.line_num == line_num) {
                const wgop = try words_set.getOrPut(gop.key_ptr.*);
                if (!wgop.found_existing) wgop.key_ptr.* = gop.key_ptr.*;
                return;
            }
        }
        try gop.value_ptr.append(self.allocator, .{
            .doc_id = doc_id,
            .line_num = line_num,
        });
        const wgop = try words_set.getOrPut(gop.key_ptr.*);
        if (!wgop.found_existing) wgop.key_ptr.* = gop.key_ptr.*;
    }

    pub fn indexFile(self: *WordIndex, path: []const u8, content: []const u8) !void {
        if (!self.enabled) return;
        // Clean up old entries first
        self.removeFile(path);

        // If the path is already tracked (e.g. skip_file_words=true and removeFile
        // early-exited), reuse the existing stable copy rather than leaking a new dup.
        const stable_path = if (self.path_to_id.contains(path))
            path
        else
            try self.allocator.dupe(u8, path);
        const owned_path = stable_path.ptr != path.ptr;
        errdefer if (owned_path) self.allocator.free(stable_path);

        const doc_id = try self.getOrCreateDocId(stable_path);

        // Use page_allocator-backed arena for words_set — pages are returned
        // to the OS immediately when the arena is deinitialized, instead of
        // being retained by the GPA (which caused ~1GB of page retention
        // across 14K files of alloc/free churn).
        var words_arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
        defer words_arena.deinit();
        var words_set = std.StringHashMap(void).init(words_arena.allocator());
        var line_num: u32 = 0;
        var lines = std.mem.splitScalar(u8, content, '\n');
        var doc_token_count: u32 = 0;

        while (lines.next()) |line| {
            line_num += 1;
            var tok = WordTokenizer{ .buf = line };
            while (tok.next()) |word| {
                if (word.len < 2) continue;
                doc_token_count +|= 1;

                const aa = words_arena.allocator();

                // Fast path: lowercase into a stack buffer (avoids one alloc per word).
                // Most identifiers fit in 256 bytes; rare longer ones fall through
                // to the arena path.
                var stack_buf: [256]u8 = undefined;
                const lower_word: []const u8 = if (word.len <= stack_buf.len) blk: {
                    for (word, 0..) |c, j| stack_buf[j] = normalizeChar(c);
                    // Check if the lowercase form differs from the original;
                    // if identical (already lowercase) and word is the actual
                    // source token, we can use `word` directly — but we'd need
                    // the string to outlive the stack frame when stored, so
                    // dup only at the insert-key boundary below.
                    break :blk stack_buf[0..word.len];
                } else blk: {
                    const buf = aa.alloc(u8, word.len) catch continue;
                    for (word, 0..) |c, j| buf[j] = normalizeChar(c);
                    break :blk buf;
                };

                // Index the lowercase form.
                try indexOneToken(self, lower_word, doc_id, line_num, &words_set);

                // Sub-tokens from identifier splitting (camelCase, snake_case, etc.).
                // Skip the alloc'd ArrayList when the word is too short or all-lower
                // (no split possible).
                var needs_split: bool = false;
                if (word.len >= 4) {
                    for (word) |c| {
                        if (c == '_' or (c >= 'A' and c <= 'Z')) {
                            needs_split = true;
                            break;
                        }
                    }
                }
                if (needs_split) {
                    var sub_toks: std.ArrayList([]const u8) = .empty;
                    defer sub_toks.deinit(aa);
                    splitIdentifier(word, &sub_toks, aa) catch continue;
                    for (sub_toks.items) |sub| {
                        try indexOneToken(self, sub, doc_id, line_num, &words_set);
                    }
                }
            }
        }

        if (!self.skip_file_words) {
            // Compact the HashMap to a simple slice — saves ~70KB per file
            const compact = try self.allocator.alloc([]const u8, words_set.count());
            var ki: usize = 0;
            var wk_iter = words_set.keyIterator();
            while (wk_iter.next()) |k| : (ki += 1) {
                compact[ki] = k.*;
            }
            try self.file_words.put(stable_path, compact);
        }
        words_set.deinit();

        if (self.doc_lengths.get(doc_id)) |old_len| {
            self.total_tokens -%= old_len;
        }
        try self.doc_lengths.put(doc_id, doc_token_count);
        self.total_tokens += doc_token_count;
    }

    /// Look up all hits for a word. O(1) lookup + O(hits) iteration.
    /// Query is normalized to lowercase (all index keys are stored lowercase).
    pub fn search(self: *WordIndex, word: []const u8) []const WordHit {
        var buf: [512]u8 = undefined;
        const lower = if (word.len <= buf.len) blk: {
            for (word, 0..) |c, i| buf[i] = normalizeChar(c);
            break :blk buf[0..word.len];
        } else word;
        if (self.index.get(lower)) |hits| return hits.items;
        return &.{};
    }

    /// Look up hits, returning results allocated by the caller.
    /// Deduplicates by (path, line_num).
    pub fn searchDeduped(self: *WordIndex, word: []const u8, allocator: std.mem.Allocator) ![]const WordHit {
        const hits = self.search(word);
        if (hits.len == 0) return try allocator.alloc(WordHit, 0);
        if (hits.len == 1) {
            var out = try allocator.alloc(WordHit, 1);
            out[0] = hits[0];
            return out;
        }

        const DedupKey = struct { doc_id: u32, line_num: u32 };
        var seen = std.AutoHashMap(DedupKey, void).init(allocator);
        defer seen.deinit();
        try seen.ensureTotalCapacity(@intCast(hits.len));

        var result: std.ArrayList(WordHit) = .empty;
        errdefer result.deinit(allocator);
        try result.ensureTotalCapacity(allocator, hits.len);

        for (hits) |hit| {
            const key = DedupKey{ .doc_id = hit.doc_id, .line_num = hit.line_num };
            const gop = try seen.getOrPut(key);
            if (!gop.found_existing) {
                result.appendAssumeCapacity(hit);
            }
        }
        return result.toOwnedSlice(allocator);
    }

    /// Collect all hits for index keys that begin with `prefix_raw` (normalized internally to
    /// lowercase). Only keys strictly longer than the normalized prefix are considered —
    /// exact-match keys are already handled by Tier 0 (`search`).
    /// Results are deduplicated by (doc_id, line_num), capped at `max_results`, and returned as an owned slice.
    pub fn searchPrefix(self: *WordIndex, prefix_raw: []const u8, allocator: std.mem.Allocator, max_results: usize) ![]const WordHit {
        if (prefix_raw.len == 0 or max_results == 0) return try allocator.alloc(WordHit, 0);
        var buf: [512]u8 = undefined;
        const prefix = if (prefix_raw.len <= buf.len) blk: {
            for (prefix_raw, 0..) |c, i| buf[i] = normalizeChar(c);
            break :blk buf[0..prefix_raw.len];
        } else prefix_raw;

        var result: std.ArrayList(WordHit) = .empty;
        errdefer result.deinit(allocator);
        try result.ensureTotalCapacity(allocator, max_results);
        const DedupKey = struct { doc_id: u32, line_num: u32 };
        var seen = std.AutoHashMap(DedupKey, void).init(allocator);
        defer seen.deinit();
        try seen.ensureTotalCapacity(@intCast(max_results));

        var key_iter = self.index.keyIterator();
        outer: while (key_iter.next()) |k| {
            if (k.len <= prefix.len) continue; // strictly longer: exact match is Tier 0
            if (!std.mem.startsWith(u8, k.*, prefix)) continue;
            const hits = self.index.get(k.*) orelse continue;
            for (hits.items) |hit| {
                const dk = DedupKey{ .doc_id = hit.doc_id, .line_num = hit.line_num };
                const gop = try seen.getOrPut(dk);
                if (!gop.found_existing) {
                    result.appendAssumeCapacity(hit);
                    if (result.items.len >= max_results) break :outer;
                }
            }
        }
        return result.toOwnedSlice(allocator);
    }

    pub fn fileCount(self: *WordIndex) u32 {
        return @intCast(self.file_words.count());
    }

    /// BM25 helper: number of docs the ranker can see (source of truth regardless of skip_file_words).
    pub fn rankedDocCount(self: *const WordIndex) u32 {
        return @intCast(self.doc_lengths.count());
    }

    /// BM25 helper: number of indexed tokens in a doc, or 0 if unknown.
    pub fn docLength(self: *const WordIndex, doc_id: u32) u32 {
        return self.doc_lengths.get(doc_id) orelse 0;
    }

    /// BM25 helper: average doc length over docs that have a recorded length.
    /// Returns 1.0 when no docs are tracked, so callers can divide safely.
    pub fn avgDocLength(self: *const WordIndex) f32 {
        const n = self.doc_lengths.count();
        if (n == 0) return 1.0;
        return @as(f32, @floatFromInt(self.total_tokens)) / @as(f32, @floatFromInt(n));
    }

    /// Shrink all hit lists and per-file word sets to release excess capacity.
    pub fn shrinkAllocations(self: *WordIndex) void {
        var iter = self.index.iterator();
        while (iter.next()) |entry| {
            if (entry.value_ptr.capacity > entry.value_ptr.items.len) {
                entry.value_ptr.shrinkAndFree(self.allocator, entry.value_ptr.items.len);
            }
        }
    }

    pub const DiskHeader = struct {
        file_count: u32,
        git_head: ?[40]u8,
    };

    const DISK_MAGIC = [4]u8{ 'C', 'D', 'B', 'W' };
    const DISK_FORMAT_VERSION: u16 = 3;

    pub fn writeToDisk(self: *WordIndex, io: std.Io, dir_path: []const u8, git_head: ?[40]u8) !void {
        var file_table: std.ArrayList([]const u8) = .empty;
        defer file_table.deinit(self.allocator);
        var disk_path_to_id = std.StringHashMap(u32).init(self.allocator);
        defer disk_path_to_id.deinit();

        if (self.file_words.count() > 0) {
            var file_iter = self.file_words.iterator();
            while (file_iter.next()) |entry| {
                const path = entry.key_ptr.*;
                if (path.len > std.math.maxInt(u16)) return error.NameTooLong;
                const id: u32 = @intCast(file_table.items.len);
                try file_table.append(self.allocator, path);
                try disk_path_to_id.put(path, id);
            }
        } else {
            for (self.id_to_path.items) |path| {
                if (path.len == 0) continue;
                if (path.len > std.math.maxInt(u16)) return error.NameTooLong;
                const id: u32 = @intCast(file_table.items.len);
                try file_table.append(self.allocator, path);
                try disk_path_to_id.put(path, id);
            }
        }

        var words_sorted: std.ArrayList([]const u8) = .empty;
        defer words_sorted.deinit(self.allocator);
        try words_sorted.ensureTotalCapacity(self.allocator, self.index.count());
        var word_iter = self.index.keyIterator();
        while (word_iter.next()) |word_ptr| {
            if (word_ptr.*.len > std.math.maxInt(u16)) return error.NameTooLong;
            words_sorted.appendAssumeCapacity(word_ptr.*);
        }
        std.mem.sort([]const u8, words_sorted.items, {}, struct {
            fn lt(_: void, a: []const u8, b: []const u8) bool {
                return std.mem.order(u8, a, b) == .lt;
            }
        }.lt);

        const rand_suffix = @as(u64, blk: {
            var ts: std.c.timespec = undefined;
            _ = std.c.clock_gettime(std.c.CLOCK.REALTIME, &ts);
            break :blk @as(u64, @intCast(ts.nsec)) ^ (@as(u64, @intCast(ts.sec)) << 1);
        });
        const tmp_path = try std.fmt.allocPrint(self.allocator, "{s}/word.index.{x}.tmp", .{ dir_path, rand_suffix });
        defer self.allocator.free(tmp_path);
        const final_path = try std.fmt.allocPrint(self.allocator, "{s}/word.index", .{dir_path});
        defer self.allocator.free(final_path);

        const file = try std.Io.Dir.cwd().createFile(io, tmp_path, .{});
        defer file.close(io);

        var file_buf: [256 * 1024]u8 = undefined;
        var writer = file.writer(io, &file_buf);

        var header_buf: [51]u8 = [_]u8{0} ** 51;
        @memcpy(header_buf[0..4], &DISK_MAGIC);
        std.mem.writeInt(u16, header_buf[4..6], DISK_FORMAT_VERSION, .little);
        std.mem.writeInt(u32, header_buf[6..10], @intCast(file_table.items.len), .little);
        if (git_head) |head| {
            header_buf[10] = 40;
            @memcpy(header_buf[11..51], &head);
        }
        try writer.interface.writeAll(&header_buf);

        for (file_table.items) |path| {
            var pl_buf: [2]u8 = undefined;
            std.mem.writeInt(u16, &pl_buf, @intCast(path.len), .little);
            try writer.interface.writeAll(&pl_buf);
            try writer.interface.writeAll(path);
        }

        var wc_buf: [4]u8 = undefined;
        std.mem.writeInt(u32, &wc_buf, @intCast(words_sorted.items.len), .little);
        try writer.interface.writeAll(&wc_buf);

        for (words_sorted.items) |word| {
            const hits = self.index.get(word) orelse return error.InvalidData;
            var wl_buf: [2]u8 = undefined;
            std.mem.writeInt(u16, &wl_buf, @intCast(word.len), .little);
            try writer.interface.writeAll(&wl_buf);
            try writer.interface.writeAll(word);

            var hc_buf: [4]u8 = undefined;
            std.mem.writeInt(u32, &hc_buf, @intCast(hits.items.len), .little);
            try writer.interface.writeAll(&hc_buf);
            for (hits.items) |hit| {
                const hit_path = self.id_to_path.items[hit.doc_id];
                const file_id = disk_path_to_id.get(hit_path) orelse return error.InvalidData;
                var hit_buf: [8]u8 = undefined;
                std.mem.writeInt(u32, hit_buf[0..4], file_id, .little);
                std.mem.writeInt(u32, hit_buf[4..8], hit.line_num, .little);
                try writer.interface.writeAll(&hit_buf);
            }
        }

        // v3 trailer: per-doc length table for BM25.
        // file_id (u32 disk-id) → length (u32). Total tokens follows as u64.
        var dl_count_buf: [4]u8 = undefined;
        std.mem.writeInt(u32, &dl_count_buf, @intCast(file_table.items.len), .little);
        try writer.interface.writeAll(&dl_count_buf);
        for (file_table.items) |path| {
            const in_mem_id = self.path_to_id.get(path) orelse {
                var z: [4]u8 = .{ 0, 0, 0, 0 };
                try writer.interface.writeAll(&z);
                continue;
            };
            const len = self.doc_lengths.get(in_mem_id) orelse 0;
            var lb: [4]u8 = undefined;
            std.mem.writeInt(u32, &lb, len, .little);
            try writer.interface.writeAll(&lb);
        }
        var tt_buf: [8]u8 = undefined;
        std.mem.writeInt(u64, &tt_buf, self.total_tokens, .little);
        try writer.interface.writeAll(&tt_buf);

        try writer.interface.flush();

        try std.Io.Dir.cwd().rename(tmp_path, std.Io.Dir.cwd(), final_path, io);
    }

    pub fn readFromDisk(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) ?WordIndex {
        return readFromDiskInner(io, dir_path, allocator) catch null;
    }

    fn readFromDiskInner(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?WordIndex {
        const index_path = try std.fmt.allocPrint(allocator, "{s}/word.index", .{dir_path});
        defer allocator.free(index_path);

        const data = std.Io.Dir.cwd().readFileAlloc(io, index_path, allocator, .limited(512 * 1024 * 1024)) catch return null;
        defer allocator.free(data);

        if (data.len < 51) return null;
        if (!std.mem.eql(u8, data[0..4], &DISK_MAGIC)) return null;
        const version = std.mem.readInt(u16, data[4..6], .little);
        if (version != DISK_FORMAT_VERSION) return null;
        const file_count = std.mem.readInt(u32, data[6..10], .little);

        var file_paths = try allocator.alloc([]u8, file_count);
        defer allocator.free(file_paths);
        var used_paths = try allocator.alloc(bool, file_count);
        defer allocator.free(used_paths);
        @memset(used_paths, false);
        var parsed_files: u32 = 0;
        defer {
            for (0..parsed_files) |i| {
                if (!used_paths[i]) allocator.free(file_paths[i]);
            }
        }

        var pos: usize = 51;
        for (0..file_count) |i| {
            if (pos + 2 > data.len) return null;
            const path_len = std.mem.readInt(u16, data[pos..][0..2], .little);
            pos += 2;
            if (path_len == 0 or pos + path_len > data.len) return null;
            file_paths[i] = try allocator.dupe(u8, data[pos .. pos + path_len]);
            pos += path_len;
            parsed_files += 1;
        }

        if (pos + 4 > data.len) return null;
        const word_count = std.mem.readInt(u32, data[pos..][0..4], .little);
        pos += 4;

        var result = WordIndex.init(allocator);
        errdefer result.deinit();

        // Temporary HashMap for accumulating file_words during load
        // (compacted to slices below)
        var tmp_file_words = std.StringHashMap(std.StringHashMap(void)).init(allocator);
        defer {
            var tfw_iter = tmp_file_words.iterator();
            while (tfw_iter.next()) |entry| entry.value_ptr.deinit();
            tmp_file_words.deinit();
        }

        for (0..word_count) |_| {
            if (pos + 2 > data.len) return null;
            const word_len = std.mem.readInt(u16, data[pos..][0..2], .little);
            pos += 2;
            if (word_len == 0 or pos + word_len > data.len) return null;
            const word = try allocator.dupe(u8, data[pos .. pos + word_len]);
            pos += word_len;
            errdefer allocator.free(word);

            if (pos + 4 > data.len) return null;
            const hit_count = std.mem.readInt(u32, data[pos..][0..4], .little);
            pos += 4;

            var hits: std.ArrayList(WordHit) = .empty;
            errdefer hits.deinit(allocator);
            try hits.ensureTotalCapacity(allocator, hit_count);

            var last_file_id: ?u32 = null;
            for (0..hit_count) |_| {
                if (pos + 8 > data.len) return null;
                const file_id = std.mem.readInt(u32, data[pos..][0..4], .little);
                pos += 4;
                if (file_id >= file_count) return null;
                const line_num = std.mem.readInt(u32, data[pos..][0..4], .little);
                pos += 4;

                hits.appendAssumeCapacity(.{
                    .doc_id = file_id,
                    .line_num = line_num,
                });

                if (last_file_id == null or last_file_id.? != file_id) {
                    const fw_gop = try tmp_file_words.getOrPut(file_paths[file_id]);
                    if (!fw_gop.found_existing) {
                        fw_gop.key_ptr.* = file_paths[file_id];
                        fw_gop.value_ptr.* = std.StringHashMap(void).init(allocator);
                        used_paths[file_id] = true;
                    }
                    const wgop = try fw_gop.value_ptr.getOrPut(word);
                    if (!wgop.found_existing) wgop.key_ptr.* = word;
                    last_file_id = file_id;
                }
            }

            const gop = try result.index.getOrPut(word);
            if (gop.found_existing) return error.InvalidData;
            gop.key_ptr.* = word;
            gop.value_ptr.* = hits;
        }

        // v3 trailer: per-doc length table.
        if (pos + 4 > data.len) return null;
        const dl_count = std.mem.readInt(u32, data[pos..][0..4], .little);
        pos += 4;
        if (dl_count != file_count) return null;
        if (pos + dl_count * 4 + 8 > data.len) return null;
        var dl_values = try allocator.alloc(u32, dl_count);
        defer allocator.free(dl_values);
        for (0..dl_count) |i| {
            dl_values[i] = std.mem.readInt(u32, data[pos..][0..4], .little);
            pos += 4;
        }
        const total_tokens_loaded = std.mem.readInt(u64, data[pos..][0..8], .little);
        pos += 8;

        if (pos != data.len) return null;

        // Populate path_to_id and id_to_path from file_paths
        try result.id_to_path.ensureTotalCapacity(allocator, file_count);
        for (0..file_count) |i| {
            result.id_to_path.appendAssumeCapacity(file_paths[i]);
            try result.path_to_id.put(file_paths[i], @intCast(i));
            if (dl_values[i] > 0) {
                try result.doc_lengths.put(@intCast(i), dl_values[i]);
            }
        }
        result.total_tokens = total_tokens_loaded;

        // Compact tmp_file_words HashMaps into slices for result.file_words
        var tfw_iter = tmp_file_words.iterator();
        while (tfw_iter.next()) |entry| {
            const compact = try allocator.alloc([]const u8, entry.value_ptr.count());
            var ki2: usize = 0;
            var wk_iter2 = entry.value_ptr.keyIterator();
            while (wk_iter2.next()) |k| : (ki2 += 1) {
                compact[ki2] = k.*;
            }
            try result.file_words.put(entry.key_ptr.*, compact);
        }

        return result;
    }

    pub fn readDiskHeader(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?DiskHeader {
        const index_path = try std.fmt.allocPrint(allocator, "{s}/word.index", .{dir_path});
        defer allocator.free(index_path);

        const file = std.Io.Dir.cwd().openFile(io, index_path, .{}) catch return null;
        defer file.close(io);

        var buf: [51]u8 = undefined;
        const n = file.readPositionalAll(io, &buf, 0) catch return null;
        if (n < 51) return null;
        if (!std.mem.eql(u8, buf[0..4], &DISK_MAGIC)) return null;
        const version = std.mem.readInt(u16, buf[4..6], .little);
        if (version < 1 or version > DISK_FORMAT_VERSION) return null;

        const file_count = std.mem.readInt(u32, buf[6..10], .little);
        var git_head: ?[40]u8 = null;
        if (buf[10] == 40) {
            var head: [40]u8 = undefined;
            @memcpy(&head, buf[11..51]);
            git_head = head;
        }
        return DiskHeader{
            .file_count = file_count,
            .git_head = git_head,
        };
    }

    pub fn readGitHead(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?[40]u8 {
        const header = try readDiskHeader(io, dir_path, allocator) orelse return null;
        return header.git_head;
    }
};

// ── Trigram index ───────────────────────────────────────────
// Maps 3-byte sequences → set of file paths.
// Enables fast substring search: extract trigrams from query,
// intersect candidate file sets, then verify with actual match.

pub const Trigram = u24;

pub fn packTrigram(a: u8, b: u8, c: u8) Trigram {
    return @as(Trigram, a) << 16 | @as(Trigram, b) << 8 | @as(Trigram, c);
}

pub const PostingMask = struct {
    next_mask: u8 = 0, // bloom filter of chars following this trigram
    loc_mask: u8 = 0, // bit mask of (position % 8) where trigram appears
};

pub const DocPosting = struct {
    doc_id: u32,
    next_mask: u8 = 0,
    loc_mask: u8 = 0,
};

pub const PostingList = struct {
    items: std.ArrayList(DocPosting) = .empty,
    path_to_id: ?*const std.StringHashMap(u32) = null,

    pub fn deinit(self: *PostingList, allocator: std.mem.Allocator) void {
        self.items.deinit(allocator);
    }

    pub fn count(self: *const PostingList) usize {
        return self.items.items.len;
    }

    pub fn get(self: *const PostingList, path: []const u8) ?PostingMask {
        const p2id = self.path_to_id orelse return null;
        const doc_id = p2id.get(path) orelse return null;
        return self.getByDocId(doc_id);
    }

    pub fn contains(self: *const PostingList, path: []const u8) bool {
        return self.get(path) != null;
    }

    pub fn getByDocId(self: *const PostingList, doc_id: u32) ?PostingMask {
        // Binary search on sorted doc_id array
        const items = self.items.items;
        var lo: usize = 0;
        var hi: usize = items.len;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            if (items[mid].doc_id == doc_id) return PostingMask{ .next_mask = items[mid].next_mask, .loc_mask = items[mid].loc_mask };
            if (items[mid].doc_id < doc_id) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return null;
    }

    pub fn containsDocId(self: *const PostingList, doc_id: u32) bool {
        const items = self.items.items;
        var lo: usize = 0;
        var hi: usize = items.len;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            if (items[mid].doc_id == doc_id) return true;
            if (items[mid].doc_id < doc_id) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return false;
    }
    pub fn getOrAddPosting(self: *PostingList, allocator: std.mem.Allocator, doc_id: u32) !*DocPosting {
        // Binary search for existing
        const items = self.items.items;
        var lo: usize = 0;
        var hi: usize = items.len;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            if (items[mid].doc_id == doc_id) return &self.items.items[mid];
            if (items[mid].doc_id < doc_id) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        // Insert at sorted position
        try self.items.insert(allocator, lo, .{ .doc_id = doc_id });
        return &self.items.items[lo];
    }

    pub fn removeDocId(self: *PostingList, doc_id: u32) void {
        var lo: usize = 0;
        var hi: usize = self.items.items.len;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            if (self.items.items[mid].doc_id < doc_id) {
                lo = mid + 1;
            } else if (self.items.items[mid].doc_id > doc_id) {
                hi = mid;
            } else {
                _ = self.items.orderedRemove(mid);
                return;
            }
        }
    }
};

pub const TrigramIndex = struct {
    /// trigram → posting list with doc IDs
    index: std.AutoHashMap(Trigram, PostingList),
    /// path → list of trigrams contributed (for cleanup)
    file_trigrams: std.StringHashMap(std.ArrayList(Trigram)),
    /// path → doc_id mapping
    path_to_id: std.StringHashMap(u32),
    /// doc_id → path mapping (may contain "" sentinels for freed slots)
    id_to_path: std.ArrayList([]const u8),
    /// freed doc_id slots available for reuse by getOrCreateDocId
    free_ids: std.ArrayList(u32),
    allocator: std.mem.Allocator,
    /// When true, deinit frees the path keys in file_trigrams (set by readFromDisk).
    owns_paths: bool = false,

    /// Maximum entries per posting list — caps memory for common trigrams.
    /// Trigrams appearing in more files than this are poor discriminators anyway.
    const MAX_POSTINGS: usize = 512;

    pub fn init(allocator: std.mem.Allocator) TrigramIndex {
        return .{
            .index = std.AutoHashMap(Trigram, PostingList).init(allocator),
            .file_trigrams = std.StringHashMap(std.ArrayList(Trigram)).init(allocator),
            .path_to_id = std.StringHashMap(u32).init(allocator),
            .id_to_path = .empty,
            .free_ids = .empty,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *TrigramIndex) void {
        var iter = self.index.iterator();
        while (iter.next()) |entry| {
            entry.value_ptr.deinit(self.allocator);
        }
        self.index.deinit();

        var ft_iter = self.file_trigrams.iterator();
        while (ft_iter.next()) |entry| {
            entry.value_ptr.deinit(self.allocator);
        }
        self.file_trigrams.deinit();

        // When owns_paths=true, paths were dup'd in getOrCreateDocId and stored
        // in id_to_path. Free them here from id_to_path (not file_trigrams),
        // because lean code paths like insertBulkNew never populate file_trigrams,
        // and removeFile clears file_trigrams entries before deinit can see them.
        // id_to_path is the single source of truth for owned path memory.
        // Tombstone slots (empty strings from removeFile) are skipped.
        if (self.owns_paths) {
            for (self.id_to_path.items) |p| {
                if (p.len > 0) self.allocator.free(p);
            }
        }

        self.path_to_id.deinit();
        self.id_to_path.deinit(self.allocator);
        self.free_ids.deinit(self.allocator);
    }

    fn getOrCreateDocId(self: *TrigramIndex, path: []const u8) !u32 {
        if (self.path_to_id.get(path)) |id| return id;
        // owns_paths=true stores a dup so callers can free their source memory.
        const stored_path: []const u8 = if (self.owns_paths)
            try self.allocator.dupe(u8, path)
        else
            path;
        errdefer if (self.owns_paths) self.allocator.free(stored_path);
        const id: u32 = if (self.free_ids.items.len > 0) blk: {
            const freed: u32 = self.free_ids.pop() orelse unreachable;
            self.id_to_path.items[@as(usize, freed)] = stored_path;
            break :blk freed;
        } else blk: {
            const new_id: u32 = @intCast(self.id_to_path.items.len);
            try self.id_to_path.append(self.allocator, stored_path);
            break :blk new_id;
        };
        try self.path_to_id.put(stored_path, id);
        return id;
    }

    pub fn removeFile(self: *TrigramIndex, path: []const u8) void {
        const doc_id = self.path_to_id.get(path) orelse {
            const trigrams = self.file_trigrams.getPtr(path) orelse return;
            trigrams.deinit(self.allocator);
            _ = self.file_trigrams.remove(path);
            return;
        };
        // Always clean path_to_id first, regardless of whether file_trigrams has an entry.
        _ = self.path_to_id.remove(path);
        // Free the doc_id slot for reuse on next indexFile call.
        self.free_ids.append(self.allocator, doc_id) catch {};
        if (self.file_trigrams.getPtr(path)) |trigrams| {
            for (trigrams.items) |tri| {
                if (self.index.getPtr(tri)) |posting_list| {
                    posting_list.removeDocId(doc_id);
                    if (posting_list.items.items.len == 0) {
                        posting_list.deinit(self.allocator);
                        _ = self.index.remove(tri);
                    }
                }
            }
            trigrams.deinit(self.allocator);
            _ = self.file_trigrams.remove(path);
        }
        const old_path = self.id_to_path.items[doc_id];
        if (self.owns_paths and old_path.len > 0) self.allocator.free(old_path);
        self.id_to_path.items[doc_id] = "";
    }

    pub fn indexFile(self: *TrigramIndex, path: []const u8, content: []const u8) !void {
        const id_count_before = self.id_to_path.items.len;
        self.removeFile(path);

        const doc_id = try self.getOrCreateDocId(path);
        // If id_to_path grew, this is a brand-new file (doc_id == max), so append
        // maintains sorted PostingList order.  If it did not grow, a freed slot was
        // reused and we must use sorted insert to preserve the invariant.
        const is_new_doc = self.id_to_path.items.len > id_count_before;

        // Phase 1: accumulate masks locally per trigram (no global index writes)
        var local = std.AutoHashMap(Trigram, PostingMask).init(self.allocator);
        defer local.deinit();
        // Pre-size: a file typically has ~content.len/4 unique trigrams
        const estimated_unique = @max(@as(u32, 64), @as(u32, @intCast(@min(content.len / 4, 65536))));
        local.ensureTotalCapacity(estimated_unique) catch {};

        if (content.len >= 3) {
            for (0..content.len - 2) |i| {
                // Skip trigrams that are pure whitespace (terrible filters, ~12% of all occurrences)
                const c0 = content[i];
                const c1 = content[i + 1];
                const c2 = content[i + 2];
                if ((c0 == ' ' or c0 == '\t' or c0 == '\n' or c0 == '\r') and
                    (c1 == ' ' or c1 == '\t' or c1 == '\n' or c1 == '\r') and
                    (c2 == ' ' or c2 == '\t' or c2 == '\n' or c2 == '\r')) continue;

                const tri = packTrigram(
                    normalizeChar(c0),
                    normalizeChar(c1),
                    normalizeChar(c2),
                );
                const gop = try local.getOrPut(tri);
                if (!gop.found_existing) {
                    gop.value_ptr.* = PostingMask{};
                }
                gop.value_ptr.loc_mask |= @as(u8, 1) << @intCast(i % 8);
                if (i + 3 < content.len) {
                    gop.value_ptr.next_mask |= @as(u8, 1) << @intCast(normalizeChar(content[i + 3]) % 8);
                }
            }
        }

        // Phase 2: bulk-insert one posting per trigram into global index
        var tri_list: std.ArrayList(Trigram) = .empty;
        errdefer tri_list.deinit(self.allocator);

        var local_iter = local.iterator();
        while (local_iter.next()) |entry| {
            const tri = entry.key_ptr.*;
            const mask = entry.value_ptr.*;

            const idx_gop = try self.index.getOrPut(tri);
            if (!idx_gop.found_existing) {
                idx_gop.value_ptr.* = .{ .path_to_id = &self.path_to_id };
            }
            if (is_new_doc) {
                // New doc_id is always max: append maintains sorted PostingList order.
                try idx_gop.value_ptr.items.append(self.allocator, .{
                    .doc_id = doc_id,
                    .next_mask = mask.next_mask,
                    .loc_mask = mask.loc_mask,
                });
            } else {
                // Reused doc_id: sorted insert to maintain PostingList binary-search invariant.
                const posting = try idx_gop.value_ptr.getOrAddPosting(self.allocator, doc_id);
                posting.next_mask = mask.next_mask;
                posting.loc_mask = mask.loc_mask;
            }

            try tri_list.append(self.allocator, tri);
        }
        const stable_path = self.id_to_path.items[doc_id];
        try self.file_trigrams.put(stable_path, tri_list);
    }

    /// Like indexFile but reuses a caller-provided local HashMap to avoid alloc/free per file.
    pub fn indexFileReuse(self: *TrigramIndex, path: []const u8, content: []const u8, local: *std.AutoHashMap(Trigram, PostingMask)) !void {
        const id_count_before = self.id_to_path.items.len;
        self.removeFile(path);
        const doc_id = try self.getOrCreateDocId(path);
        const is_new_doc = self.id_to_path.items.len > id_count_before;

        // Phase 1: accumulate masks in reusable local map
        local.clearRetainingCapacity();
        if (content.len >= 3) {
            for (0..content.len - 2) |i| {
                const c0 = content[i];
                const c1 = content[i + 1];
                const c2 = content[i + 2];
                if ((c0 == ' ' or c0 == '\t' or c0 == '\n' or c0 == '\r') and
                    (c1 == ' ' or c1 == '\t' or c1 == '\n' or c1 == '\r') and
                    (c2 == ' ' or c2 == '\t' or c2 == '\n' or c2 == '\r')) continue;
                const tri = packTrigram(normalizeChar(c0), normalizeChar(c1), normalizeChar(c2));
                const gop = try local.getOrPut(tri);
                if (!gop.found_existing) gop.value_ptr.* = PostingMask{};
                gop.value_ptr.loc_mask |= @as(u8, 1) << @intCast(i % 8);
                if (i + 3 < content.len) {
                    gop.value_ptr.next_mask |= @as(u8, 1) << @intCast(normalizeChar(content[i + 3]) % 8);
                }
            }
        }

        // Phase 2: bulk-insert
        var tri_list: std.ArrayList(Trigram) = .empty;
        errdefer tri_list.deinit(self.allocator);
        var local_iter = local.iterator();
        while (local_iter.next()) |entry| {
            const tri = entry.key_ptr.*;
            const mask = entry.value_ptr.*;
            const idx_gop = try self.index.getOrPut(tri);
            if (!idx_gop.found_existing) {
                idx_gop.value_ptr.* = .{ .path_to_id = &self.path_to_id };
            }
            if (is_new_doc) {
                try idx_gop.value_ptr.items.append(self.allocator, .{
                    .doc_id = doc_id,
                    .next_mask = mask.next_mask,
                    .loc_mask = mask.loc_mask,
                });
            } else {
                const posting = try idx_gop.value_ptr.getOrAddPosting(self.allocator, doc_id);
                posting.next_mask = mask.next_mask;
                posting.loc_mask = mask.loc_mask;
            }
            try tri_list.append(self.allocator, tri);
        }
        try self.file_trigrams.put(path, tri_list);
    }

    /// Extract trigrams from content — thread-safe, no shared state.
    pub fn extractTrigrams(content: []const u8, alloc: std.mem.Allocator) std.AutoHashMap(Trigram, PostingMask) {
        var local = std.AutoHashMap(Trigram, PostingMask).init(alloc);
        const estimated = @max(@as(u32, 64), @as(u32, @intCast(@min(content.len / 4, 65536))));
        local.ensureTotalCapacity(estimated) catch {};
        if (content.len >= 3) {
            for (0..content.len - 2) |i| {
                const c0 = content[i];
                const c1 = content[i + 1];
                const c2 = content[i + 2];
                if ((c0 == ' ' or c0 == '\t' or c0 == '\n' or c0 == '\r') and
                    (c1 == ' ' or c1 == '\t' or c1 == '\n' or c1 == '\r') and
                    (c2 == ' ' or c2 == '\t' or c2 == '\n' or c2 == '\r')) continue;
                const tri = packTrigram(normalizeChar(c0), normalizeChar(c1), normalizeChar(c2));
                const gop = local.getOrPut(tri) catch continue;
                if (!gop.found_existing) gop.value_ptr.* = PostingMask{};
                gop.value_ptr.loc_mask |= @as(u8, 1) << @intCast(i % 8);
                if (i + 3 < content.len) {
                    gop.value_ptr.next_mask |= @as(u8, 1) << @intCast(normalizeChar(content[i + 3]) % 8);
                }
            }
        }
        return local;
    }

    /// Insert pre-extracted trigrams. NOT thread-safe.
    pub fn insertExtracted(self: *TrigramIndex, path: []const u8, local: *std.AutoHashMap(Trigram, PostingMask)) !void {
        self.removeFile(path);
        const doc_id = try self.getOrCreateDocId(path);
        var tri_list: std.ArrayList(Trigram) = .empty;
        errdefer tri_list.deinit(self.allocator);
        var iter = local.iterator();
        while (iter.next()) |entry| {
            const tri = entry.key_ptr.*;
            const mask = entry.value_ptr.*;
            const idx_gop = try self.index.getOrPut(tri);
            if (!idx_gop.found_existing) {
                idx_gop.value_ptr.* = .{ .path_to_id = &self.path_to_id };
            }
            try idx_gop.value_ptr.items.append(self.allocator, .{
                .doc_id = doc_id,
                .next_mask = mask.next_mask,
                .loc_mask = mask.loc_mask,
            });
            try tri_list.append(self.allocator, tri);
        }
        try self.file_trigrams.put(path, tri_list);
    }

    pub const BulkEntry = struct { tri: Trigram, mask: PostingMask };

    /// Lean bulk insert for cold builds: no removeFile, no file_trigrams tracking.
    /// Assumes doc_id is always new (strictly increasing), so uses simple append.
    pub fn insertBulkNew(self: *TrigramIndex, path: []const u8, trigrams: []const BulkEntry) !void {
        const doc_id = try self.getOrCreateDocId(path);
        for (trigrams) |te| {
            const idx_gop = try self.index.getOrPut(te.tri);
            if (!idx_gop.found_existing) {
                idx_gop.value_ptr.* = .{ .path_to_id = &self.path_to_id };
            }
            try idx_gop.value_ptr.items.append(self.allocator, .{
                .doc_id = doc_id,
                .next_mask = te.mask.next_mask,
                .loc_mask = te.mask.loc_mask,
            });
        }
    }

    /// Find candidate files that contain ALL trigrams from the query.
    pub fn candidates(self: *TrigramIndex, query: []const u8, allocator: std.mem.Allocator) ?[]const []const u8 {
        if (query.len < 3) return null;

        const tri_count = query.len - 2;

        var unique = std.AutoHashMap(Trigram, void).init(allocator);
        defer unique.deinit();
        unique.ensureTotalCapacity(@intCast(tri_count)) catch return null;
        for (0..tri_count) |i| {
            const tri = packTrigram(
                normalizeChar(query[i]),
                normalizeChar(query[i + 1]),
                normalizeChar(query[i + 2]),
            );
            _ = unique.getOrPut(tri) catch return null;
        }

        var sets: std.ArrayList(*PostingList) = .empty;
        defer sets.deinit(allocator);
        sets.ensureTotalCapacity(allocator, unique.count()) catch return null;

        var tri_iter = unique.keyIterator();
        while (tri_iter.next()) |tri_ptr| {
            const posting_list = self.index.getPtr(tri_ptr.*) orelse {
                return allocator.alloc([]const u8, 0) catch null;
            };
            sets.appendAssumeCapacity(posting_list);
        }

        if (sets.items.len == 0) {
            return allocator.alloc([]const u8, 0) catch null;
        }

        // Sort posting lists by size (smallest first) for efficient intersection
        std.mem.sort(*PostingList, sets.items, {}, struct {
            fn lt(_: void, a: *PostingList, b: *PostingList) bool {
                return a.items.items.len < b.items.items.len;
            }
        }.lt);

        // Sorted merge intersection: start with smallest list's doc_ids
        var result_ids: std.ArrayList(u32) = .empty;
        defer result_ids.deinit(allocator);

        // Seed with doc_ids from smallest posting list
        result_ids.ensureTotalCapacity(allocator, sets.items[0].items.items.len) catch return null;
        for (sets.items[0].items.items) |p| {
            result_ids.appendAssumeCapacity(p.doc_id);
        }

        // Intersect with each subsequent list (both sorted → merge O(n+m))
        for (sets.items[1..]) |set| {
            var write: usize = 0;
            var si: usize = 0;
            const set_items = set.items.items;
            for (result_ids.items) |id| {
                // Advance set pointer to >= id
                while (si < set_items.len and set_items[si].doc_id < id) : (si += 1) {}
                if (si < set_items.len and set_items[si].doc_id == id) {
                    result_ids.items[write] = id;
                    write += 1;
                    si += 1;
                }
            }
            result_ids.items.len = write;
            if (write == 0) break; // early exit if intersection is empty
        }

        var result: std.ArrayList([]const u8) = .empty;
        errdefer result.deinit(allocator);
        result.ensureTotalCapacity(allocator, result_ids.items.len) catch return null;

        next_cand: for (result_ids.items) |doc_id| {
            // Bloom-filter check for consecutive trigram pairs
            if (tri_count >= 2) {
                for (0..tri_count - 1) |j| {
                    const tri_a = packTrigram(
                        normalizeChar(query[j]),
                        normalizeChar(query[j + 1]),
                        normalizeChar(query[j + 2]),
                    );
                    const tri_b = packTrigram(
                        normalizeChar(query[j + 1]),
                        normalizeChar(query[j + 2]),
                        normalizeChar(query[j + 3]),
                    );
                    const list_a = self.index.getPtr(tri_a) orelse continue;
                    const list_b = self.index.getPtr(tri_b) orelse continue;
                    const mask_a = list_a.getByDocId(doc_id) orelse continue;
                    const mask_b = list_b.getByDocId(doc_id) orelse continue;

                    const next_bit: u8 = @as(u8, 1) << @intCast(normalizeChar(query[j + 3]) % 8);
                    if ((mask_a.next_mask & next_bit) == 0) continue :next_cand;

                    const rotated = (mask_a.loc_mask << 1) | (mask_a.loc_mask >> 7);
                    if ((rotated & mask_b.loc_mask) == 0) continue :next_cand;
                }
            }

            if (doc_id < self.id_to_path.items.len) {
                result.appendAssumeCapacity(self.id_to_path.items[doc_id]);
            }
        }

        return result.toOwnedSlice(allocator) catch {
            result.deinit(allocator);
            return null;
        };
    }

    pub fn candidatesRegex(self: *TrigramIndex, query: *const RegexQuery, allocator: std.mem.Allocator) ?[]const []const u8 {
        if (query.and_trigrams.len == 0 and query.or_groups.len == 0) return null;

        var result_set: ?std.AutoHashMap(u32, void) = null;
        defer if (result_set) |*rs| rs.deinit();

        if (query.and_trigrams.len > 0) {
            for (query.and_trigrams) |tri| {
                const posting_list = self.index.getPtr(tri) orelse {
                    var empty = allocator.alloc([]const u8, 0) catch return null;
                    _ = &empty;
                    return allocator.alloc([]const u8, 0) catch null;
                };
                if (result_set == null) {
                    result_set = std.AutoHashMap(u32, void).init(allocator);
                    for (posting_list.items.items) |p| {
                        result_set.?.put(p.doc_id, {}) catch return null;
                    }
                } else {
                    var to_remove: std.ArrayList(u32) = .empty;
                    defer to_remove.deinit(allocator);
                    var it = result_set.?.keyIterator();
                    while (it.next()) |key| {
                        if (!posting_list.containsDocId(key.*)) {
                            to_remove.append(allocator, key.*) catch return null;
                        }
                    }
                    for (to_remove.items) |key| {
                        _ = result_set.?.remove(key);
                    }
                }
            }
        }

        for (query.or_groups) |group| {
            if (group.len == 0) continue;

            var union_set = std.AutoHashMap(u32, void).init(allocator);
            defer union_set.deinit();
            for (group) |tri| {
                const posting_list = self.index.getPtr(tri) orelse continue;
                for (posting_list.items.items) |p| {
                    union_set.put(p.doc_id, {}) catch return null;
                }
            }

            if (result_set == null) {
                result_set = std.AutoHashMap(u32, void).init(allocator);
                var it = union_set.keyIterator();
                while (it.next()) |key| {
                    result_set.?.put(key.*, {}) catch return null;
                }
            } else {
                var to_remove: std.ArrayList(u32) = .empty;
                defer to_remove.deinit(allocator);
                var it = result_set.?.keyIterator();
                while (it.next()) |key| {
                    if (!union_set.contains(key.*)) {
                        to_remove.append(allocator, key.*) catch return null;
                    }
                }
                for (to_remove.items) |key| {
                    _ = result_set.?.remove(key);
                }
            }
        }

        if (result_set == null) return null;

        var result: std.ArrayList([]const u8) = .empty;
        errdefer result.deinit(allocator);
        result.ensureTotalCapacity(allocator, result_set.?.count()) catch return null;
        var it = result_set.?.keyIterator();
        while (it.next()) |id_ptr| {
            const doc_id = id_ptr.*;
            if (doc_id < self.id_to_path.items.len) {
                result.appendAssumeCapacity(self.id_to_path.items[doc_id]);
            }
        }
        return result.toOwnedSlice(allocator) catch {
            result.deinit(allocator);
            return null;
        };
    }

    // ── Disk persistence ────────────────────────────────────

    pub const POSTINGS_MAGIC = [4]u8{ 'C', 'D', 'B', 'T' };
    pub const LOOKUP_MAGIC = [4]u8{ 'C', 'D', 'B', 'L' };
    pub const FORMAT_VERSION: u16 = 3;

    /// Posting entry for v3+: file_id (u32) + next_mask (u8) + loc_mask (u8) + pad (2 bytes) = 8 bytes
    pub const DiskPosting = extern struct {
        file_id: u32,
        next_mask: u8,
        loc_mask: u8,
        _pad: [2]u8 = .{ 0, 0 },
    };

    /// Posting entry for v1/v2 files: file_id (u16) + next_mask (u8) + loc_mask (u8) = 4 bytes
    pub const OldDiskPosting = extern struct {
        file_id: u16,
        next_mask: u8,
        loc_mask: u8,
    };

    /// Lookup entry: trigram (u32 low 24 bits) + offset (u32) + count (u32) = 12 bytes
    pub const LookupEntry = extern struct {
        trigram: u32,
        offset: u32,
        count: u32,
    };

    /// Write the current in-memory index to disk in a two-file format.
    /// Files are written atomically (write to tmp, then rename).
    pub fn writeToDisk(self: *TrigramIndex, io: std.Io, dir_path: []const u8, git_head: ?[40]u8) !void {
        // Step 1: Build file table from path_to_id (reuse existing doc IDs for consistency)
        var file_table: std.ArrayList([]const u8) = .empty;
        defer file_table.deinit(self.allocator);
        var disk_path_to_id = std.StringHashMap(u32).init(self.allocator);
        defer disk_path_to_id.deinit();

        if (self.file_trigrams.count() > 0) {
            var ft_iter = self.file_trigrams.keyIterator();
            while (ft_iter.next()) |path_ptr| {
                const id: u32 = @intCast(file_table.items.len);
                try file_table.append(self.allocator, path_ptr.*);
                try disk_path_to_id.put(path_ptr.*, id);
            }
        } else {
            for (self.id_to_path.items) |path| {
                if (path.len == 0) continue;
                const id: u32 = @intCast(file_table.items.len);
                try file_table.append(self.allocator, path);
                try disk_path_to_id.put(path, id);
            }
        }

        const file_count: u32 = @intCast(file_table.items.len);

        // Step 2: Collect all trigrams, sort them, serialize postings contiguously
        var trigrams_sorted: std.ArrayList(Trigram) = .empty;
        defer trigrams_sorted.deinit(self.allocator);
        {
            var tri_iter = self.index.keyIterator();
            while (tri_iter.next()) |tri_ptr| {
                try trigrams_sorted.append(self.allocator, tri_ptr.*);
            }
        }
        std.mem.sort(Trigram, trigrams_sorted.items, {}, struct {
            fn lt(_: void, a: Trigram, b: Trigram) bool {
                return a < b;
            }
        }.lt);

        // Step 3: Build postings blob and lookup entries
        var postings_buf: std.ArrayList(DiskPosting) = .empty;
        defer postings_buf.deinit(self.allocator);
        var lookup_entries: std.ArrayList(LookupEntry) = .empty;
        defer lookup_entries.deinit(self.allocator);

        for (trigrams_sorted.items) |tri| {
            const posting_list = self.index.getPtr(tri) orelse continue;
            const offset: u32 = @intCast(postings_buf.items.len);
            var count: u32 = 0;
            for (posting_list.items.items) |p| {
                // Map in-memory doc_id to disk file_id via path lookup
                if (p.doc_id >= self.id_to_path.items.len) continue;
                const path = self.id_to_path.items[p.doc_id];
                const fid = disk_path_to_id.get(path) orelse continue;
                try postings_buf.append(self.allocator, .{
                    .file_id = fid,
                    .next_mask = p.next_mask,
                    .loc_mask = p.loc_mask,
                });
                count += 1;
            }
            try lookup_entries.append(self.allocator, .{
                .trigram = @as(u32, tri),
                .offset = offset,
                .count = count,
            });
        }

        // Step 4: Write postings file atomically (random suffix prevents collisions)
        const post_rand: u64 = cio.randU64();
        const postings_tmp = try std.fmt.allocPrint(self.allocator, "{s}/trigram.postings.{x}.tmp", .{ dir_path, post_rand });
        defer self.allocator.free(postings_tmp);
        const postings_final = try std.fmt.allocPrint(self.allocator, "{s}/trigram.postings", .{dir_path});
        defer self.allocator.free(postings_final);

        {
            const file = try std.Io.Dir.cwd().createFile(io, postings_tmp, .{});
            defer file.close(io);

            var pw_buf: [256 * 1024]u8 = undefined;
            var pw = file.writer(io, &pw_buf);

            // Header v3: magic(4) + version(2) + file_count(4) + head_len(1) + head(40) = 51 bytes
            try pw.interface.writeAll(&POSTINGS_MAGIC);
            var ver_buf: [2]u8 = undefined;
            std.mem.writeInt(u16, &ver_buf, FORMAT_VERSION, .little);
            try pw.interface.writeAll(&ver_buf);
            var fc_buf: [4]u8 = undefined;
            std.mem.writeInt(u32, &fc_buf, file_count, .little);
            try pw.interface.writeAll(&fc_buf);
            // Git HEAD: head_len (1 byte) + head (40 bytes)
            if (git_head) |head| {
                try pw.interface.writeAll(&.{40});
                try pw.interface.writeAll(&head);
            } else {
                try pw.interface.writeAll(&.{0});
                try pw.interface.writeAll(&([_]u8{0} ** 40));
            }

            // File table: for each file, path_len(u16) + path bytes
            for (file_table.items) |path| {
                var pl_buf: [2]u8 = undefined;
                std.mem.writeInt(u16, &pl_buf, @intCast(path.len), .little);
                try pw.interface.writeAll(&pl_buf);
                try pw.interface.writeAll(path);
            }

            // Postings data
            const postings_bytes = std.mem.sliceAsBytes(postings_buf.items);
            try pw.interface.writeAll(postings_bytes);
            try pw.interface.flush();
        }
        try std.Io.Dir.cwd().rename(postings_tmp, std.Io.Dir.cwd(), postings_final, io);

        // Step 5: Write lookup file atomically (random suffix prevents collisions)
        const lk_rand: u64 = cio.randU64();
        const lookup_tmp = try std.fmt.allocPrint(self.allocator, "{s}/trigram.lookup.{x}.tmp", .{ dir_path, lk_rand });
        defer self.allocator.free(lookup_tmp);
        const lookup_final = try std.fmt.allocPrint(self.allocator, "{s}/trigram.lookup", .{dir_path});
        defer self.allocator.free(lookup_final);

        {
            const file = try std.Io.Dir.cwd().createFile(io, lookup_tmp, .{});
            defer file.close(io);

            var lw_buf: [64 * 1024]u8 = undefined;
            var lw = file.writer(io, &lw_buf);

            // Header: magic(4) + version(2) + pad(2) + entry_count(4) = 12 bytes
            try lw.interface.writeAll(&LOOKUP_MAGIC);
            var ver_buf2: [2]u8 = undefined;
            std.mem.writeInt(u16, &ver_buf2, FORMAT_VERSION, .little);
            try lw.interface.writeAll(&ver_buf2);
            var pad_buf: [2]u8 = .{ 0, 0 };
            try lw.interface.writeAll(&pad_buf);
            var ec_buf: [4]u8 = undefined;
            std.mem.writeInt(u32, &ec_buf, @intCast(lookup_entries.items.len), .little);
            try lw.interface.writeAll(&ec_buf);

            // Entries (already aligned at 12 bytes each)
            const entry_bytes = std.mem.sliceAsBytes(lookup_entries.items);
            try lw.interface.writeAll(entry_bytes);
            try lw.interface.flush();
        }
        try std.Io.Dir.cwd().rename(lookup_tmp, std.Io.Dir.cwd(), lookup_final, io);
    }

    /// Load index from disk files into a fresh TrigramIndex.
    /// Returns null if files don't exist or are corrupt/stale.
    pub fn readFromDisk(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) ?TrigramIndex {
        return readFromDiskInner(io, dir_path, allocator) catch null;
    }

    fn readFromDiskInner(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?TrigramIndex {
        const postings_path = try std.fmt.allocPrint(allocator, "{s}/trigram.postings", .{dir_path});
        defer allocator.free(postings_path);
        const lookup_path = try std.fmt.allocPrint(allocator, "{s}/trigram.lookup", .{dir_path});
        defer allocator.free(lookup_path);

        // Read both files
        const postings_data = std.Io.Dir.cwd().readFileAlloc(io, postings_path, allocator, .limited(64 * 1024 * 1024)) catch return null;
        defer allocator.free(postings_data);
        const lookup_data = std.Io.Dir.cwd().readFileAlloc(io, lookup_path, allocator, .limited(64 * 1024 * 1024)) catch return null;
        defer allocator.free(lookup_data);

        // Validate postings header (v1: 8 bytes, v2: 49 bytes, v3: 51 bytes)
        if (postings_data.len < 8) return null;
        if (!std.mem.eql(u8, postings_data[0..4], &POSTINGS_MAGIC)) return null;
        const post_version = std.mem.readInt(u16, postings_data[4..6], .little);
        if (post_version < 1 or post_version > FORMAT_VERSION) return null;
        const file_count: u32 = if (post_version >= 3)
            std.mem.readInt(u32, postings_data[6..10], .little)
        else
            std.mem.readInt(u16, postings_data[6..8], .little);

        const file_table_start: usize = if (post_version >= 3) blk: {
            if (postings_data.len < 51) return null;
            break :blk 51;
        } else if (post_version >= 2) blk: {
            if (postings_data.len < 49) return null;
            break :blk 49;
        } else 8;

        // Parse file table
        var file_paths = try allocator.alloc([]u8, file_count);
        var parsed_files: u32 = 0;
        defer {
            for (0..parsed_files) |i| allocator.free(file_paths[i]);
            allocator.free(file_paths);
        }
        var pos: usize = file_table_start;
        for (0..file_count) |i| {
            if (pos + 2 > postings_data.len) return null;
            const path_len = std.mem.readInt(u16, postings_data[pos..][0..2], .little);
            pos += 2;
            if (pos + path_len > postings_data.len) return null;
            file_paths[i] = try allocator.dupe(u8, postings_data[pos .. pos + path_len]);
            parsed_files += 1;
            pos += path_len;
        }

        // Remaining bytes are DiskPosting entries
        const postings_start = pos;
        const postings_byte_len = postings_data.len - postings_start;
        const posting_size: usize = if (post_version >= 3) @sizeOf(DiskPosting) else @sizeOf(OldDiskPosting);
        if (postings_byte_len % posting_size != 0) return null;
        const total_postings = postings_byte_len / posting_size;

        // Validate lookup header
        if (lookup_data.len < 12) return null;
        if (!std.mem.eql(u8, lookup_data[0..4], &LOOKUP_MAGIC)) return null;
        const lk_version = std.mem.readInt(u16, lookup_data[4..6], .little);
        if (lk_version < 1 or lk_version > FORMAT_VERSION) return null;
        const entry_count = std.mem.readInt(u32, lookup_data[8..12], .little);
        if (lookup_data.len < 12 + entry_count * @sizeOf(LookupEntry)) return null;

        // Build in-memory index
        var result = TrigramIndex.init(allocator);
        result.owns_paths = true;
        errdefer result.deinit();

        // Allocate stable path strings owned by the index and build doc ID mappings
        var stable_paths = try allocator.alloc([]const u8, file_count);
        defer allocator.free(stable_paths);
        for (0..file_count) |i| {
            const duped = try allocator.dupe(u8, file_paths[i]);
            errdefer allocator.free(duped);
            stable_paths[i] = duped;
            try result.file_trigrams.put(duped, .empty);
            try result.path_to_id.put(duped, @intCast(i));
            try result.id_to_path.append(allocator, duped);
        }

        // Parse lookup entries and populate index + file_trigrams
        for (0..entry_count) |e| {
            const entry_off = 12 + e * @sizeOf(LookupEntry);
            const raw = lookup_data[entry_off..][0..@sizeOf(LookupEntry)];
            const entry: *align(1) const LookupEntry = @ptrCast(raw.ptr);

            const tri: Trigram = @intCast(entry.trigram);
            const p_off = entry.offset;
            const p_count = entry.count;

            if (@as(u64, p_off) + @as(u64, p_count) > @as(u64, total_postings)) return error.InvalidData;

            var posting_list: PostingList = .{ .path_to_id = &result.path_to_id };
            errdefer posting_list.deinit(allocator);

            for (0..p_count) |pi| {
                const pb_off = postings_start + (p_off + pi) * posting_size;
                const raw_posting = postings_data[pb_off..][0..posting_size];
                const file_id: u32 = if (post_version >= 3)
                    std.mem.readInt(u32, raw_posting[0..4], .little)
                else
                    std.mem.readInt(u16, raw_posting[0..2], .little);
                const next_mask = raw_posting[if (post_version >= 3) 4 else 2];
                const loc_mask = raw_posting[if (post_version >= 3) 5 else 3];

                if (file_id >= file_count) return error.InvalidData;

                const doc_id: u32 = file_id;
                const posting = try posting_list.getOrAddPosting(allocator, doc_id);
                posting.next_mask |= next_mask;
                posting.loc_mask |= loc_mask;

                // Track trigram in file_trigrams
                const path = stable_paths[file_id];
                if (result.file_trigrams.getPtr(path)) |tri_list| {
                    var found = false;
                    for (tri_list.items) |existing| {
                        if (existing == tri) {
                            found = true;
                            break;
                        }
                    }
                    if (!found) try tri_list.append(allocator, tri);
                }
            }

            try result.index.put(tri, posting_list);
        }

        return result;
    }

    /// Returns the number of indexed files (for staleness checks).
    pub fn fileCount(self: *const TrigramIndex) u32 {
        return @intCast(self.file_trigrams.count());
    }

    /// Shrink all posting lists to their actual length, releasing excess capacity.
    /// Call after bulk indexing to reclaim ArrayList over-allocation (~50% savings).
    pub fn shrinkPostingLists(self: *TrigramIndex) void {
        var iter = self.index.iterator();
        while (iter.next()) |entry| {
            if (entry.value_ptr.items.capacity > entry.value_ptr.items.items.len) {
                entry.value_ptr.items.shrinkAndFree(self.allocator, entry.value_ptr.items.items.len);
            }
        }
        // Also shrink file_trigrams lists
        var ft_iter = self.file_trigrams.iterator();
        while (ft_iter.next()) |entry| {
            if (entry.value_ptr.capacity > entry.value_ptr.items.len) {
                entry.value_ptr.shrinkAndFree(self.allocator, entry.value_ptr.items.len);
            }
        }
    }

    /// Header info that can be read without loading the full index.
    pub const DiskHeader = struct {
        file_count: u32,
        git_head: ?[40]u8,
    };

    /// Read just the postings file header — fast, no full file load.
    /// Returns null if the file doesn't exist or has an unrecognised format.
    pub fn readDiskHeader(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?DiskHeader {
        const postings_path = try std.fmt.allocPrint(allocator, "{s}/trigram.postings", .{dir_path});
        defer allocator.free(postings_path);

        const file = std.Io.Dir.cwd().openFile(io, postings_path, .{}) catch return null;
        defer file.close(io);

        var buf: [51]u8 = undefined;
        const n = file.readPositionalAll(io, &buf, 0) catch return null;
        if (n < 8) return null;
        if (!std.mem.eql(u8, buf[0..4], &POSTINGS_MAGIC)) return null;
        const version = std.mem.readInt(u16, buf[4..6], .little);
        if (version < 1 or version > FORMAT_VERSION) return null;
        const file_count: u32 = if (version >= 3)
            std.mem.readInt(u32, buf[6..10], .little)
        else
            std.mem.readInt(u16, buf[6..8], .little);

        var git_head: ?[40]u8 = null;
        if (version >= 3 and n >= 51) {
            const head_len = buf[10];
            if (head_len == 40) {
                var head: [40]u8 = undefined;
                @memcpy(&head, buf[11..51]);
                git_head = head;
            }
        } else if (version >= 2 and n >= 49) {
            const head_len = buf[8];
            if (head_len == 40) {
                var head: [40]u8 = undefined;
                @memcpy(&head, buf[9..49]);
                git_head = head;
            }
        }
        return DiskHeader{ .file_count = file_count, .git_head = git_head };
    }

    /// Read the git HEAD stored in the disk index header.
    /// Returns null if no git HEAD is stored or the file doesn't exist.
    pub fn readGitHead(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?[40]u8 {
        const header = try readDiskHeader(io, dir_path, allocator) orelse return null;
        return header.git_head;
    }
};

// ── mmap-backed trigram index ───────────────────────────────
// Zero-copy: binary search on mmap'd lookup table, read postings directly.
// Replaces heap-based TrigramIndex after writeToDisk for O(log n) lookups
// with ~0 RSS (data lives in OS page cache).

pub const MmapTrigramIndex = struct {
    const mmap_align = std.heap.page_size_min;
    postings_data: []align(mmap_align) const u8,
    lookup_data: []align(mmap_align) const u8,
    file_table: []const []const u8,
    file_set: std.StringHashMap(void),
    postings_start: usize,
    lookup_entries: usize,
    post_version: u16,
    allocator: std.mem.Allocator,

    pub fn initFromDisk(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) ?MmapTrigramIndex {
        return initFromDiskInner(io, dir_path, allocator) catch null;
    }

    fn initFromDiskInner(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?MmapTrigramIndex {
        const postings_path = try std.fmt.allocPrint(allocator, "{s}/trigram.postings", .{dir_path});
        defer allocator.free(postings_path);
        const lookup_path = try std.fmt.allocPrint(allocator, "{s}/trigram.lookup", .{dir_path});
        defer allocator.free(lookup_path);

        // mmap postings file
        const post_file = std.Io.Dir.cwd().openFile(io, postings_path, .{}) catch return null;
        defer post_file.close(io);
        const post_size = post_file.length(io) catch return null;
        if (post_size < 8) return null;
        const postings_data = std.posix.mmap(
            null,
            post_size,
            .{ .READ = true },
            .{ .TYPE = .SHARED },
            post_file.handle,
            0,
        ) catch return null;
        errdefer std.posix.munmap(postings_data);

        // mmap lookup file
        const lk_file = std.Io.Dir.cwd().openFile(io, lookup_path, .{}) catch {
            std.posix.munmap(postings_data);
            return null;
        };
        defer lk_file.close(io);
        const lk_size = lk_file.length(io) catch {
            std.posix.munmap(postings_data);
            return null;
        };
        if (lk_size < 12) {
            std.posix.munmap(postings_data);
            return null;
        }
        const lookup_data = std.posix.mmap(
            null,
            lk_size,
            .{ .READ = true },
            .{ .TYPE = .SHARED },
            lk_file.handle,
            0,
        ) catch {
            std.posix.munmap(postings_data);
            return null;
        };
        errdefer std.posix.munmap(lookup_data);

        // Validate postings header
        if (!std.mem.eql(u8, postings_data[0..4], &TrigramIndex.POSTINGS_MAGIC)) return null;
        const post_version = std.mem.readInt(u16, postings_data[4..6], .little);
        if (post_version < 1 or post_version > TrigramIndex.FORMAT_VERSION) return null;
        const file_count: u32 = if (post_version >= 3)
            std.mem.readInt(u32, postings_data[6..10], .little)
        else
            std.mem.readInt(u16, postings_data[6..8], .little);

        const file_table_start: usize = if (post_version >= 3) blk: {
            if (postings_data.len < 51) return null;
            break :blk 51;
        } else if (post_version >= 2) blk: {
            if (postings_data.len < 49) return null;
            break :blk 49;
        } else 8;

        // Parse file table (we need owned path strings for lookups)
        var file_table = try allocator.alloc([]const u8, file_count);
        var parsed: u32 = 0;
        errdefer {
            for (0..parsed) |i| allocator.free(file_table[i]);
            allocator.free(file_table);
        }
        var pos: usize = file_table_start;
        for (0..file_count) |i| {
            if (pos + 2 > postings_data.len) return null;
            const path_len = std.mem.readInt(u16, postings_data[pos..][0..2], .little);
            pos += 2;
            if (pos + path_len > postings_data.len) return null;
            file_table[i] = try allocator.dupe(u8, postings_data[pos .. pos + path_len]);
            parsed += 1;
            pos += path_len;
        }

        // Build file_set for containsFile queries
        var file_set = std.StringHashMap(void).init(allocator);
        errdefer file_set.deinit();
        for (file_table[0..parsed]) |p| {
            try file_set.put(p, {});
        }

        const postings_start = pos;

        // Validate lookup header
        if (!std.mem.eql(u8, lookup_data[0..4], &TrigramIndex.LOOKUP_MAGIC)) return null;
        const lk_version = std.mem.readInt(u16, lookup_data[4..6], .little);
        if (lk_version < 1 or lk_version > TrigramIndex.FORMAT_VERSION) return null;
        const entry_count = std.mem.readInt(u32, lookup_data[8..12], .little);
        if (lookup_data.len < 12 + entry_count * @sizeOf(TrigramIndex.LookupEntry)) return null;

        return MmapTrigramIndex{
            .postings_data = postings_data,
            .lookup_data = lookup_data,
            .file_table = file_table,
            .file_set = file_set,
            .postings_start = postings_start,
            .lookup_entries = entry_count,
            .post_version = post_version,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *MmapTrigramIndex) void {
        for (self.file_table) |p| self.allocator.free(p);
        self.allocator.free(self.file_table);
        self.file_set.deinit();
        std.posix.munmap(self.postings_data);
        std.posix.munmap(self.lookup_data);
    }

    pub fn fileCount(self: *const MmapTrigramIndex) u32 {
        return @intCast(self.file_table.len);
    }

    pub fn containsFile(self: *const MmapTrigramIndex, path: []const u8) bool {
        return self.file_set.contains(path);
    }

    fn lookupTrigram(self: *const MmapTrigramIndex, tri_val: u32) ?struct { offset: u32, count: u32 } {
        const entries = self.lookup_entries;
        if (entries == 0) return null;
        var lo: usize = 0;
        var hi: usize = entries;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            const entry_off = 12 + mid * @sizeOf(TrigramIndex.LookupEntry);
            const entry_tri = std.mem.readInt(u32, self.lookup_data[entry_off..][0..4], .little);
            if (entry_tri == tri_val) {
                const offset = std.mem.readInt(u32, self.lookup_data[entry_off + 4 ..][0..4], .little);
                const count = std.mem.readInt(u32, self.lookup_data[entry_off + 8 ..][0..4], .little);
                return .{ .offset = offset, .count = count };
            }
            if (entry_tri < tri_val) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return null;
    }

    fn readPosting(self: *const MmapTrigramIndex, index: usize) ?struct { file_id: u32, next_mask: u8, loc_mask: u8 } {
        const posting_size: usize = if (self.post_version >= 3) @sizeOf(TrigramIndex.DiskPosting) else @sizeOf(TrigramIndex.OldDiskPosting);
        const pb_off = self.postings_start + index * posting_size;
        if (pb_off + posting_size > self.postings_data.len) return null;
        const raw = self.postings_data[pb_off..][0..posting_size];
        const file_id: u32 = if (self.post_version >= 3)
            std.mem.readInt(u32, raw[0..4], .little)
        else
            std.mem.readInt(u16, raw[0..2], .little);
        const next_mask = raw[if (self.post_version >= 3) 4 else 2];
        const loc_mask = raw[if (self.post_version >= 3) 5 else 3];
        return .{ .file_id = file_id, .next_mask = next_mask, .loc_mask = loc_mask };
    }

    pub fn candidates(self: *const MmapTrigramIndex, query: []const u8, allocator: std.mem.Allocator) ?[]const []const u8 {
        if (query.len < 3) return null;

        const tri_count = query.len - 2;

        // Collect unique trigrams
        var unique = std.AutoHashMap(Trigram, void).init(allocator);
        defer unique.deinit();
        unique.ensureTotalCapacity(@intCast(tri_count)) catch return null;
        for (0..tri_count) |i| {
            const tri = packTrigram(
                normalizeChar(query[i]),
                normalizeChar(query[i + 1]),
                normalizeChar(query[i + 2]),
            );
            _ = unique.getOrPut(tri) catch return null;
        }

        // Collect posting ranges for each trigram, sorted by count (smallest first)
        const Range = struct { offset: u32, count: u32 };
        var ranges: std.ArrayList(Range) = .empty;
        defer ranges.deinit(allocator);
        ranges.ensureTotalCapacity(allocator, unique.count()) catch return null;

        var tri_iter = unique.keyIterator();
        while (tri_iter.next()) |tri_ptr| {
            const r = self.lookupTrigram(@as(u32, tri_ptr.*)) orelse {
                return allocator.alloc([]const u8, 0) catch null;
            };
            ranges.appendAssumeCapacity(.{ .offset = r.offset, .count = r.count });
        }

        if (ranges.items.len == 0) {
            return allocator.alloc([]const u8, 0) catch null;
        }

        std.mem.sort(Range, ranges.items, {}, struct {
            fn lt(_: void, a: Range, b: Range) bool {
                return a.count < b.count;
            }
        }.lt);

        // Seed with file_ids from smallest posting range
        var result_ids: std.ArrayList(u32) = .empty;
        defer result_ids.deinit(allocator);
        result_ids.ensureTotalCapacity(allocator, ranges.items[0].count) catch return null;
        for (0..ranges.items[0].count) |pi| {
            const p = self.readPosting(ranges.items[0].offset + pi) orelse continue;
            result_ids.appendAssumeCapacity(p.file_id);
        }

        // Intersect with subsequent ranges (sorted merge)
        for (ranges.items[1..]) |range| {
            var write: usize = 0;
            var si: usize = 0;
            for (result_ids.items) |id| {
                while (si < range.count) {
                    const p = self.readPosting(range.offset + si) orelse break;
                    if (p.file_id >= id) {
                        if (p.file_id == id) {
                            result_ids.items[write] = id;
                            write += 1;
                            si += 1;
                        }
                        break;
                    }
                    si += 1;
                }
            }
            result_ids.items.len = write;
            if (write == 0) break;
        }

        // Bloom filter verification for consecutive trigram pairs
        var result: std.ArrayList([]const u8) = .empty;
        errdefer result.deinit(allocator);
        result.ensureTotalCapacity(allocator, result_ids.items.len) catch return null;

        next_cand: for (result_ids.items) |file_id| {
            if (tri_count >= 2) {
                for (0..tri_count - 1) |j| {
                    const tri_a_val = @as(u32, packTrigram(
                        normalizeChar(query[j]),
                        normalizeChar(query[j + 1]),
                        normalizeChar(query[j + 2]),
                    ));
                    const tri_b_val = @as(u32, packTrigram(
                        normalizeChar(query[j + 1]),
                        normalizeChar(query[j + 2]),
                        normalizeChar(query[j + 3]),
                    ));
                    const range_a = self.lookupTrigram(tri_a_val) orelse continue;
                    const range_b = self.lookupTrigram(tri_b_val) orelse continue;
                    const mask_a = self.findPostingMask(range_a.offset, range_a.count, file_id) orelse continue;
                    const mask_b = self.findPostingMask(range_b.offset, range_b.count, file_id) orelse continue;

                    const next_bit: u8 = @as(u8, 1) << @intCast(normalizeChar(query[j + 3]) % 8);
                    if ((mask_a.next_mask & next_bit) == 0) continue :next_cand;

                    const rotated = (mask_a.loc_mask << 1) | (mask_a.loc_mask >> 7);
                    if ((rotated & mask_b.loc_mask) == 0) continue :next_cand;
                }
            }

            if (file_id < self.file_table.len) {
                result.appendAssumeCapacity(self.file_table[file_id]);
            }
        }

        return result.toOwnedSlice(allocator) catch {
            result.deinit(allocator);
            return null;
        };
    }

    fn findPostingMask(self: *const MmapTrigramIndex, offset: u32, count: u32, file_id: u32) ?PostingMask {
        var lo: usize = 0;
        var hi: usize = count;
        while (lo < hi) {
            const mid = lo + (hi - lo) / 2;
            const p = self.readPosting(offset + mid) orelse return null;
            if (p.file_id == file_id) return PostingMask{ .next_mask = p.next_mask, .loc_mask = p.loc_mask };
            if (p.file_id < file_id) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return null;
    }

    pub fn candidatesRegex(self: *const MmapTrigramIndex, query: *const RegexQuery, allocator: std.mem.Allocator) ?[]const []const u8 {
        if (query.and_trigrams.len == 0 and query.or_groups.len == 0) return null;

        var result_set: ?std.AutoHashMap(u32, void) = null;
        defer if (result_set) |*rs| rs.deinit();

        if (query.and_trigrams.len > 0) {
            for (query.and_trigrams) |tri| {
                const range = self.lookupTrigram(@as(u32, tri)) orelse {
                    return allocator.alloc([]const u8, 0) catch null;
                };
                if (result_set == null) {
                    result_set = std.AutoHashMap(u32, void).init(allocator);
                    for (0..range.count) |pi| {
                        const p = self.readPosting(range.offset + pi) orelse continue;
                        result_set.?.put(p.file_id, {}) catch return null;
                    }
                } else {
                    var to_remove: std.ArrayList(u32) = .empty;
                    defer to_remove.deinit(allocator);
                    var it = result_set.?.keyIterator();
                    while (it.next()) |key| {
                        if (self.findPostingMask(range.offset, range.count, key.*) == null) {
                            to_remove.append(allocator, key.*) catch return null;
                        }
                    }
                    for (to_remove.items) |key| {
                        _ = result_set.?.remove(key);
                    }
                }
            }
        }

        for (query.or_groups) |group| {
            if (group.len == 0) continue;

            var union_set = std.AutoHashMap(u32, void).init(allocator);
            defer union_set.deinit();
            for (group) |tri| {
                const range = self.lookupTrigram(@as(u32, tri)) orelse continue;
                for (0..range.count) |pi| {
                    const p = self.readPosting(range.offset + pi) orelse continue;
                    union_set.put(p.file_id, {}) catch return null;
                }
            }

            if (result_set == null) {
                result_set = std.AutoHashMap(u32, void).init(allocator);
                var it = union_set.keyIterator();
                while (it.next()) |key| {
                    result_set.?.put(key.*, {}) catch return null;
                }
            } else {
                var to_remove: std.ArrayList(u32) = .empty;
                defer to_remove.deinit(allocator);
                var it = result_set.?.keyIterator();
                while (it.next()) |key| {
                    if (!union_set.contains(key.*)) {
                        to_remove.append(allocator, key.*) catch return null;
                    }
                }
                for (to_remove.items) |key| {
                    _ = result_set.?.remove(key);
                }
            }
        }

        if (result_set == null) return null;

        var result: std.ArrayList([]const u8) = .empty;
        errdefer result.deinit(allocator);
        result.ensureTotalCapacity(allocator, result_set.?.count()) catch return null;
        var it = result_set.?.keyIterator();
        while (it.next()) |id_ptr| {
            const doc_id = id_ptr.*;
            if (doc_id < self.file_table.len) {
                result.appendAssumeCapacity(self.file_table[doc_id]);
            }
        }
        return result.toOwnedSlice(allocator) catch {
            result.deinit(allocator);
            return null;
        };
    }
};

pub const AnyTrigramIndex = union(enum) {
    heap: TrigramIndex,
    mmap: MmapTrigramIndex,
    mmap_overlay: MmapOverlay,

    pub const MmapOverlay = struct {
        base: MmapTrigramIndex,
        overlay: TrigramIndex,

        pub fn deinit(self: *MmapOverlay) void {
            self.base.deinit();
            self.overlay.deinit();
        }
    };

    pub fn deinit(self: *AnyTrigramIndex) void {
        switch (self.*) {
            .heap => |*h| h.deinit(),
            .mmap => |*m| m.deinit(),
            .mmap_overlay => |*mo| mo.deinit(),
        }
    }

    pub fn candidates(self: *AnyTrigramIndex, query: []const u8, allocator: std.mem.Allocator) ?[]const []const u8 {
        return switch (self.*) {
            .heap => |*h| h.candidates(query, allocator),
            .mmap => |*m| m.candidates(query, allocator),
            .mmap_overlay => |*mo| blk: {
                const base = mo.base.candidates(query, allocator);
                const over = mo.overlay.candidates(query, allocator);
                if (base == null and over == null) break :blk null;
                if (base == null) break :blk over;
                if (over == null) break :blk base;
                // Merge and dedup — return null on alloc failure (triggers full scan fallback)
                var merged = std.StringHashMap(void).init(allocator);
                defer merged.deinit();
                for (base.?) |p| merged.put(p, {}) catch {
                    allocator.free(base.?);
                    allocator.free(over.?);
                    break :blk null;
                };
                for (over.?) |p| merged.put(p, {}) catch {
                    allocator.free(base.?);
                    allocator.free(over.?);
                    break :blk null;
                };
                allocator.free(base.?);
                allocator.free(over.?);
                var result: std.ArrayList([]const u8) = .empty;
                result.ensureTotalCapacity(allocator, merged.count()) catch break :blk null;
                var it = merged.keyIterator();
                while (it.next()) |k| result.appendAssumeCapacity(k.*);
                break :blk result.toOwnedSlice(allocator) catch {
                    result.deinit(allocator);
                    break :blk null;
                };
            },
        };
    }

    pub fn candidatesRegex(self: *AnyTrigramIndex, query: *const RegexQuery, allocator: std.mem.Allocator) ?[]const []const u8 {
        return switch (self.*) {
            .heap => |*h| h.candidatesRegex(query, allocator),
            .mmap => |*m| m.candidatesRegex(query, allocator),
            .mmap_overlay => |*mo| blk: {
                const base = mo.base.candidatesRegex(query, allocator);
                const over = mo.overlay.candidatesRegex(query, allocator);
                if (base == null and over == null) break :blk null;
                if (base == null) break :blk over;
                if (over == null) break :blk base;
                var merged = std.StringHashMap(void).init(allocator);
                defer merged.deinit();
                for (base.?) |p| merged.put(p, {}) catch {
                    allocator.free(base.?);
                    allocator.free(over.?);
                    break :blk null;
                };
                for (over.?) |p| merged.put(p, {}) catch {
                    allocator.free(base.?);
                    allocator.free(over.?);
                    break :blk null;
                };
                allocator.free(base.?);
                allocator.free(over.?);
                var result: std.ArrayList([]const u8) = .empty;
                result.ensureTotalCapacity(allocator, merged.count()) catch break :blk null;
                var it = merged.keyIterator();
                while (it.next()) |k| result.appendAssumeCapacity(k.*);
                break :blk result.toOwnedSlice(allocator) catch {
                    result.deinit(allocator);
                    break :blk null;
                };
            },
        };
    }

    pub fn containsFile(self: *const AnyTrigramIndex, path: []const u8) bool {
        return switch (self.*) {
            .heap => |*h| h.file_trigrams.contains(path),
            .mmap => |*m| m.containsFile(path),
            .mmap_overlay => |*mo| mo.base.containsFile(path) or mo.overlay.file_trigrams.contains(path),
        };
    }

    pub fn indexFile(self: *AnyTrigramIndex, path: []const u8, content: []const u8) !void {
        switch (self.*) {
            .heap => |*h| try h.indexFile(path, content),
            .mmap => |*m| {
                // Promote to mmap_overlay: keep mmap base, add heap overlay
                const alloc = m.allocator;
                const base = self.mmap;
                self.* = .{ .mmap_overlay = .{
                    .base = base,
                    .overlay = TrigramIndex.init(alloc),
                } };
                try self.mmap_overlay.overlay.indexFile(path, content);
            },
            .mmap_overlay => |*mo| try mo.overlay.indexFile(path, content),
        }
    }

    pub fn removeFile(self: *AnyTrigramIndex, path: []const u8) void {
        switch (self.*) {
            .heap => |*h| h.removeFile(path),
            .mmap => {},
            .mmap_overlay => |*mo| mo.overlay.removeFile(path),
        }
    }

    pub fn writeToDisk(self: *AnyTrigramIndex, io: std.Io, dir_path: []const u8, git_head: ?[40]u8) !void {
        switch (self.*) {
            .heap => |*h| try h.writeToDisk(io, dir_path, git_head),
            .mmap => {},
            .mmap_overlay => {},
        }
    }

    pub fn fileCount(self: *const AnyTrigramIndex) u32 {
        return switch (self.*) {
            .heap => |*h| h.fileCount(),
            .mmap => |*m| m.fileCount(),
            .mmap_overlay => |*mo| mo.base.fileCount() + mo.overlay.fileCount(),
        };
    }

    pub fn asHeap(self: *AnyTrigramIndex) ?*TrigramIndex {
        return switch (self.*) {
            .heap => |*h| h,
            .mmap => null,
            .mmap_overlay => |*mo| &mo.overlay,
        };
    }
};
// ── Regex decomposition ─────────────────────────────────────

pub const RegexQuery = struct {
    and_trigrams: []Trigram,
    or_groups: [][]Trigram,
    allocator: std.mem.Allocator,

    pub fn deinit(self: *RegexQuery) void {
        self.allocator.free(self.and_trigrams);
        for (self.or_groups) |group| {
            self.allocator.free(group);
        }
        self.allocator.free(self.or_groups);
    }
};

/// Parse a regex pattern and extract literal segments that yield trigrams.
/// Handles: . \s \w \d * + ? | [...] \ (escapes)
/// Literal runs >= 3 chars produce AND trigrams.
/// Alternations (foo|bar) produce OR groups.
pub fn decomposeRegex(pattern: []const u8, allocator: std.mem.Allocator) !RegexQuery {
    // First check if this is an alternation at the top level
    // We need to respect grouping: only split on | outside of [...] and (...)
    var top_pipes: std.ArrayList(usize) = .empty;
    defer top_pipes.deinit(allocator);

    {
        var depth: usize = 0;
        var in_bracket = false;
        var i: usize = 0;
        while (i < pattern.len) {
            const c = pattern[i];
            if (c == '\\' and i + 1 < pattern.len) {
                i += 2;
                continue;
            }
            if (c == '[') {
                in_bracket = true;
                i += 1;
                continue;
            }
            if (c == ']') {
                in_bracket = false;
                i += 1;
                continue;
            }
            if (in_bracket) {
                i += 1;
                continue;
            }
            if (c == '(') {
                depth += 1;
                i += 1;
                continue;
            }
            if (c == ')') {
                if (depth > 0) depth -= 1;
                i += 1;
                continue;
            }
            if (c == '|' and depth == 0) {
                try top_pipes.append(allocator, i);
            }
            i += 1;
        }
    }

    if (top_pipes.items.len > 0) {
        // Top-level alternation: merge all branch trigrams into a single OR group.
        // A file matching ANY branch's trigrams is a valid candidate.
        var all_tris: std.ArrayList(Trigram) = .empty;
        errdefer all_tris.deinit(allocator);

        var start: usize = 0;
        for (top_pipes.items) |pipe_pos| {
            const branch = pattern[start..pipe_pos];
            const branch_tris = try extractLiteralTrigrams(branch, allocator);
            defer allocator.free(branch_tris);
            for (branch_tris) |tri| {
                try all_tris.append(allocator, tri);
            }
            start = pipe_pos + 1;
        }
        // Last branch
        const last_branch = pattern[start..];
        const last_tris = try extractLiteralTrigrams(last_branch, allocator);
        defer allocator.free(last_tris);
        for (last_tris) |tri| {
            try all_tris.append(allocator, tri);
        }

        const empty_and = try allocator.alloc(Trigram, 0);
        var or_groups: std.ArrayList([]Trigram) = .empty;
        errdefer or_groups.deinit(allocator);
        if (all_tris.items.len > 0) {
            try or_groups.append(allocator, try all_tris.toOwnedSlice(allocator));
        }
        return RegexQuery{
            .and_trigrams = empty_and,
            .or_groups = try or_groups.toOwnedSlice(allocator),
            .allocator = allocator,
        };
    }

    // No top-level alternation: extract trigrams from literal segments
    const and_tris = try extractLiteralTrigrams(pattern, allocator);
    const empty_or = try allocator.alloc([]Trigram, 0);
    return RegexQuery{
        .and_trigrams = and_tris,
        .or_groups = empty_or,
        .allocator = allocator,
    };
}

/// Extract trigrams from literal runs in a regex fragment (no top-level |).
fn extractLiteralTrigrams(pattern: []const u8, allocator: std.mem.Allocator) ![]Trigram {
    var literals: std.ArrayList(u8) = .empty;
    defer literals.deinit(allocator);

    var trigrams_list: std.ArrayList(Trigram) = .empty;
    errdefer trigrams_list.deinit(allocator);

    // Deduplicate trigrams
    var seen = std.AutoHashMap(Trigram, void).init(allocator);
    defer seen.deinit();

    var i: usize = 0;
    while (i < pattern.len) {
        const c = pattern[i];

        // Escape sequences
        if (c == '\\' and i + 1 < pattern.len) {
            const next = pattern[i + 1];
            switch (next) {
                's', 'S', 'w', 'W', 'd', 'D', 'b', 'B' => {
                    // Character class — breaks literal chain
                    try flushLiterals(allocator, &literals, &trigrams_list, &seen);
                    i += 2;
                    // If followed by quantifier, skip it too
                    if (i < pattern.len and isQuantifier(pattern[i])) i += 1;
                    continue;
                },
                else => {
                    // Escaped literal char (e.g. \. \( \) \\ etc.)
                    try literals.append(allocator, next);
                    i += 2;
                    // Check for quantifier after escaped char
                    if (i < pattern.len and isQuantifier(pattern[i])) {
                        // Quantifier on single char — pop it and flush
                        if (literals.items.len > 0) {
                            _ = literals.pop();
                        }
                        try flushLiterals(allocator, &literals, &trigrams_list, &seen);
                        i += 1;
                    }
                    continue;
                },
            }
        }

        // Character class [...]
        if (c == '[') {
            try flushLiterals(allocator, &literals, &trigrams_list, &seen);
            // Skip to closing ]
            i += 1;
            if (i < pattern.len and pattern[i] == '^') i += 1;
            if (i < pattern.len and pattern[i] == ']') i += 1; // literal ] at start
            while (i < pattern.len and pattern[i] != ']') : (i += 1) {}
            if (i < pattern.len) i += 1; // skip ]
            // Skip quantifier after class
            if (i < pattern.len and isQuantifier(pattern[i])) i += 1;
            continue;
        }

        // Grouping parens — just skip them, process contents
        if (c == '(' or c == ')') {
            try flushLiterals(allocator, &literals, &trigrams_list, &seen);
            i += 1;
            continue;
        }

        // Anchors
        if (c == '^' or c == '$') {
            try flushLiterals(allocator, &literals, &trigrams_list, &seen);
            i += 1;
            continue;
        }

        // Dot — any char, breaks chain
        if (c == '.') {
            try flushLiterals(allocator, &literals, &trigrams_list, &seen);
            i += 1;
            if (i < pattern.len and isQuantifier(pattern[i])) i += 1;
            continue;
        }

        // Quantifiers on previous char
        if (isQuantifier(c)) {
            // Remove last literal (it's now optional/repeated)
            if (literals.items.len > 0) {
                _ = literals.pop();
            }
            try flushLiterals(allocator, &literals, &trigrams_list, &seen);
            // If it's a brace quantifier {n}, {n,m}, {n,}, skip to closing }
            if (c == '{') {
                i += 1;
                while (i < pattern.len and pattern[i] != '}') : (i += 1) {}
                if (i < pattern.len) i += 1; // skip '}'
            } else {
                i += 1;
            }
            continue;
        }

        // Plain literal character
        try literals.append(allocator, c);
        i += 1;
    }

    // Flush remaining literals
    try flushLiterals(allocator, &literals, &trigrams_list, &seen);

    return trigrams_list.toOwnedSlice(allocator);
}

fn isQuantifier(c: u8) bool {
    return c == '*' or c == '+' or c == '?' or c == '{';
}

/// Flush a run of literal characters into trigrams (if >= 3 chars).
fn flushLiterals(
    allocator: std.mem.Allocator,
    literals: *std.ArrayList(u8),
    trigrams_list: *std.ArrayList(Trigram),
    seen: *std.AutoHashMap(Trigram, void),
) !void {
    if (literals.items.len >= 3) {
        for (0..literals.items.len - 2) |j| {
            const tri = packTrigram(
                normalizeChar(literals.items[j]),
                normalizeChar(literals.items[j + 1]),
                normalizeChar(literals.items[j + 2]),
            );
            const gop = try seen.getOrPut(tri);
            if (!gop.found_existing) {
                try trigrams_list.append(allocator, tri);
            }
        }
    }
    literals.clearRetainingCapacity();
}

// ── Tokenizer ───────────────────────────────────────────────

pub const WordTokenizer = struct {
    buf: []const u8,
    pos: usize = 0,

    pub fn next(self: *WordTokenizer) ?[]const u8 {
        // Skip non-word chars
        while (self.pos < self.buf.len and !isWordChar(self.buf[self.pos])) {
            self.pos += 1;
        }
        if (self.pos >= self.buf.len) return null;

        const start = self.pos;
        while (self.pos < self.buf.len and isWordChar(self.buf[self.pos])) {
            self.pos += 1;
        }
        return self.buf[start..self.pos];
    }
};

fn isWordChar(c: u8) bool {
    return std.ascii.isAlphanumeric(c) or c == '_';
}

pub fn normalizeChar(c: u8) u8 {
    // Lowercase for case-insensitive trigram matching
    return if (c >= 'A' and c <= 'Z') c + 32 else c;
}

fn emitSubToken(seg: []const u8, out: *std.ArrayList([]const u8), arena: std.mem.Allocator) !void {
    if (seg.len < 2) return;
    const lower = try arena.alloc(u8, seg.len);
    for (seg, 0..) |c, j| lower[j] = normalizeChar(c);
    try out.append(arena, lower);
}

/// Split an identifier into lowercased sub-tokens.
/// Handles: snake_case, camelCase, SCREAMING_CASE, HTTPHandler-style acronyms.
/// Appends to `out`; strings are allocated from `arena`.
pub fn splitIdentifier(token: []const u8, out: *std.ArrayList([]const u8), arena: std.mem.Allocator) !void {
    if (token.len < 2) return;

    var start: usize = 0;
    var i: usize = 1;
    while (i < token.len) : (i += 1) {
        const c = token[i];
        const prev = token[i - 1];
        var split = false;

        if (c == '_') {
            try emitSubToken(token[start..i], out, arena);
            start = i + 1;
            continue;
        }
        if (prev == '_') continue;

        // CamelCase: lowercase → uppercase boundary
        if (c >= 'A' and c <= 'Z' and prev >= 'a' and prev <= 'z') split = true;

        // Digit/letter transition
        if (!split) {
            const c_d = c >= '0' and c <= '9';
            const p_d = prev >= '0' and prev <= '9';
            if (c_d != p_d) split = true;
        }

        // Acronym end: HTTPHandler → HTTP|Handler
        if (!split and c >= 'A' and c <= 'Z' and prev >= 'A' and prev <= 'Z') {
            if (i + 1 < token.len and token[i + 1] >= 'a' and token[i + 1] <= 'z') split = true;
        }

        if (split) {
            try emitSubToken(token[start..i], out, arena);
            start = i;
        }
    }
    try emitSubToken(token[start..token.len], out, arena);
}
// ── Sparse N-gram index ───────────────────────────────────────────────────────

pub const MAX_NGRAM_LEN: usize = 16;

/// Comptime character-pair frequency table for source code.
/// Common pairs → LOW weight (they stay interior to n-grams).
/// Rare pairs   → HIGH weight (they become n-gram boundaries).
/// All unspecified pairs default to 0xFE00 (rare = high weight).
pub const default_pair_freq: [256][256]u16 = blk: {
    var table: [256][256]u16 = .{.{0xFE00} ** 256} ** 256;
    // English bigrams (lowercase) — common in identifiers and prose
    table['t']['h'] = 0x1000;
    table['h']['e'] = 0x1000;
    table['i']['n'] = 0x1000;
    table['e']['r'] = 0x1000;
    table['a']['n'] = 0x1000;
    table['r']['e'] = 0x1000;
    table['o']['n'] = 0x1000;
    table['e']['n'] = 0x1000;
    table['s']['t'] = 0x1000;
    table['e']['s'] = 0x1000;
    table['a']['t'] = 0x1000;
    table['i']['o'] = 0x1000;
    table['t']['e'] = 0x1000;
    table['o']['r'] = 0x1000;
    table['t']['i'] = 0x1000;
    table['a']['r'] = 0x1000;
    table['a']['l'] = 0x1000;
    table['l']['e'] = 0x1000;
    table['n']['t'] = 0x1000;
    table['e']['d'] = 0x1000;
    table['n']['d'] = 0x1000;
    table['o']['u'] = 0x1000;
    table['e']['a'] = 0x1000;
    table['f']['o'] = 0x1000;
    // Common code keyword fragments
    table['f']['n'] = 0x1000;
    table['i']['f'] = 0x1000;
    table['r']['n'] = 0x1000;
    table['t']['u'] = 0x1000;
    table['p']['u'] = 0x1000;
    table['b']['l'] = 0x1000;
    table['c']['o'] = 0x1000;
    table['n']['s'] = 0x1000;
    table['t']['r'] = 0x1000;
    table['u']['e'] = 0x1000;
    // Common operator / punctuation pairs
    table['('][')'] = 0x0800;
    table['{']['}'] = 0x0800;
    table['['][']'] = 0x0800;
    table['/']['/'] = 0x0800;
    table['-']['>'] = 0x0800;
    table['=']['>'] = 0x0800;
    table[':'][':'] = 0x0800;
    table['!']['='] = 0x0800;
    table['=']['='] = 0x0800;
    table['<']['='] = 0x0800;
    table['>']['='] = 0x0800;
    table['&']['&'] = 0x0800;
    table['|']['|'] = 0x0800;
    // Whitespace / structural pairs
    table[' '][' '] = 0x0800;
    table['\t'][' '] = 0x0800;
    table[' ']['('] = 0x0800;
    table[' ']['{'] = 0x0800;
    table[';'][' '] = 0x0800;
    table[':'][' '] = 0x0800;
    table['='][' '] = 0x0800;
    table[' ']['='] = 0x0800;
    table[','][' '] = 0x0800;
    table['.']['.'] = 0x0800;
    table['\n'][' '] = 0x0800;
    table['\n']['\t'] = 0x0800;
    break :blk table;
};

/// Active frequency table — points to the comptime default or a runtime
/// per-project table.  Swap only before indexing starts (not thread-safe).
pub var active_pair_freq: *const [256][256]u16 = &default_pair_freq;
var loaded_freq_table: [256][256]u16 = undefined;

/// Deterministic weight for a character pair, used to place content-defined
/// boundaries between n-grams.  Frequency-weighted: common source-code pairs
/// get LOW weight (they stay interior to n-grams); rare pairs get HIGH weight
/// (they become boundaries).  A small hash jitter (0-255) breaks ties
/// deterministically between pairs in the same frequency tier.
pub fn pairWeight(a: u8, b: u8) u16 {
    const freq_weight = active_pair_freq[a][b];
    const pair = [2]u8{ a, b };
    const jitter: u16 = @truncate(std.hash.Wyhash.hash(0, &pair) & 0xFF);
    return freq_weight +| jitter;
}

/// Swap in a custom frequency table.  Call before indexing; not thread-safe.
pub fn setFrequencyTable(table: *const [256][256]u16) void {
    loaded_freq_table = table.*;
    active_pair_freq = &loaded_freq_table;
}

/// Revert to the built-in comptime frequency table.
pub fn resetFrequencyTable() void {
    active_pair_freq = &default_pair_freq;
}

/// Build a per-project frequency table by counting byte-pair occurrences in
/// `content`, then inverting counts to weights (common → low, rare → high).
pub fn buildFrequencyTable(content: []const u8) [256][256]u16 {
    var counts: [256][256]u64 = .{.{0} ** 256} ** 256;
    if (content.len >= 2) {
        for (0..content.len - 1) |i| {
            counts[content[i]][content[i + 1]] += 1;
        }
    }
    return finishFrequencyTable(&counts);
}

/// Build a frequency table by streaming over multiple content slices.
/// Zero extra memory — counts pairs within each slice, skipping cross-slice
/// boundaries (negligible loss for large corpora).
pub fn buildFrequencyTableFromSlices(slices: []const []const u8) [256][256]u16 {
    var counts: [256][256]u64 = .{.{0} ** 256} ** 256;
    for (slices) |content| {
        if (content.len < 2) continue;
        for (0..content.len - 1) |i| {
            counts[content[i]][content[i + 1]] += 1;
        }
    }
    return finishFrequencyTable(&counts);
}

/// Build a frequency table by streaming over a StringHashMap of content.
/// Iterates file-by-file — no concatenation, zero extra memory.
pub fn buildFrequencyTableFromMap(contents: *ContentCache) [256][256]u16 {
    var counts: [256][256]u64 = .{.{0} ** 256} ** 256;
    var iter = contents.iterator();
    while (iter.next()) |entry| {
        const content = entry.value_ptr.*;
        if (content.len < 2) continue;
        for (0..content.len - 1) |i| {
            counts[content[i]][content[i + 1]] += 1;
        }
    }
    return finishFrequencyTable(&counts);
}

fn finishFrequencyTable(counts: *const [256][256]u64) [256][256]u16 {
    var max_count: u64 = 1;
    for (counts) |row| {
        for (row) |c| {
            if (c > max_count) max_count = c;
        }
    }
    // Invert: count 0 → 0xFE00 (rare, high); max_count → 0x1000 (common, low).
    var table: [256][256]u16 = .{.{0xFE00} ** 256} ** 256;
    for (0..256) |a| {
        for (0..256) |b| {
            const c = counts[a][b];
            if (c == 0) continue;
            const span: u64 = 0xFE00 - 0x1000;
            const w: u64 = 0xFE00 - (c * span / max_count);
            table[a][b] = @intCast(@min(w, 0xFE00));
        }
    }
    return table;
}

/// Persist a frequency table as a raw binary blob to `<dir_path>/pair_freq.bin`.
/// Uses tmp+rename for atomic writes.
pub fn writeFrequencyTable(io: std.Io, table: *const [256][256]u16, dir_path: []const u8) !void {
    var dir = try std.Io.Dir.cwd().openDir(io, dir_path, .{});
    defer dir.close(io);
    {
        const tmp = try dir.createFile(io, "pair_freq.bin.tmp", .{});
        defer tmp.close(io);
        var tw_buf: [64 * 1024]u8 = undefined;
        var tw = tmp.writer(io, &tw_buf);
        var row_buf: [256 * 2]u8 = undefined;
        for (table) |row| {
            for (row, 0..) |val, j| {
                std.mem.writeInt(u16, row_buf[j * 2 ..][0..2], val, .little);
            }
            try tw.interface.writeAll(&row_buf);
        }
        try tw.interface.flush();
    }
    try dir.rename("pair_freq.bin.tmp", dir, "pair_freq.bin", io);
}

/// Load a frequency table from `<dir_path>/pair_freq.bin`.
/// Returns null if the file does not exist or has the wrong size.
/// Caller owns the returned allocation.
pub fn readFrequencyTable(io: std.Io, dir_path: []const u8, allocator: std.mem.Allocator) !?*[256][256]u16 {
    const path = try std.fmt.allocPrint(allocator, "{s}/pair_freq.bin", .{dir_path});
    defer allocator.free(path);
    const file = std.Io.Dir.cwd().openFile(io, path, .{}) catch |err| switch (err) {
        error.FileNotFound => return null,
        else => return err,
    };
    defer file.close(io);
    const expected_size = 256 * 256 * @sizeOf(u16);
    const size = try file.length(io);
    if (size != expected_size) return null;
    const result = try allocator.create([256][256]u16);
    errdefer allocator.destroy(result);
    var read_pos: u64 = 0;
    var row_buf: [256 * 2]u8 = undefined;
    for (result) |*row| {
        const n = try file.readPositionalAll(io, &row_buf, read_pos);
        if (n != row_buf.len) {
            allocator.destroy(result);
            return null;
        }
        read_pos += row_buf.len;
        for (row, 0..) |*val, j| {
            val.* = std.mem.readInt(u16, row_buf[j * 2 ..][0..2], .little);
        }
    }
    return result;
}

/// A single sparse n-gram extracted from a string.
pub const SparseNgram = struct {
    hash: u64, // Wyhash of the normalized (lowercased) n-gram bytes
    pos: usize, // byte offset in the source string
    len: usize, // byte length of the n-gram
};

fn makeNgram(content: []const u8, pos: usize, len: usize) SparseNgram {
    var buf: [MAX_NGRAM_LEN]u8 = undefined;
    for (0..len) |k| buf[k] = normalizeChar(content[pos + k]);
    return .{
        .hash = std.hash.Wyhash.hash(0, buf[0..len]),
        .pos = pos,
        .len = len,
    };
}

/// Extract sparse n-grams from `content` using content-defined boundaries.
///
/// Boundaries are placed at strict local maxima of pairWeight over the
/// normalized character pairs.  N-grams span consecutive boundaries; spans
/// wider than MAX_NGRAM_LEN are force-split into MAX_NGRAM_LEN chunks.
/// Minimum n-gram length is 3 (same as a trigram).
///
/// Caller owns the returned slice.
pub fn extractSparseNgrams(content: []const u8, allocator: std.mem.Allocator) ![]SparseNgram {
    const MIN_LEN = 3;
    if (content.len < MIN_LEN) return try allocator.alloc(SparseNgram, 0);

    const pair_count = content.len - 1;

    // Compute pair weights.
    const weights = try allocator.alloc(u16, pair_count);
    defer allocator.free(weights);
    for (0..pair_count) |i| {
        weights[i] = pairWeight(normalizeChar(content[i]), normalizeChar(content[i + 1]));
    }

    // Collect boundary pair-positions: always include 0 and pair_count-1,
    // plus any interior strict local maximum.
    var bounds: std.ArrayList(usize) = .empty;
    defer bounds.deinit(allocator);

    try bounds.append(allocator, 0);
    if (pair_count >= 3) {
        for (1..pair_count - 1) |i| {
            if (weights[i] > weights[i - 1] and weights[i] > weights[i + 1]) {
                try bounds.append(allocator, i);
            }
        }
    }
    try bounds.append(allocator, pair_count - 1);

    // Emit n-grams spanning consecutive boundary positions.
    // N-gram for boundary pair at position p covers content[p .. p+2].
    var result: std.ArrayList(SparseNgram) = .empty;
    errdefer result.deinit(allocator);

    var b: usize = 0;
    while (b + 1 < bounds.items.len) : (b += 1) {
        const start = bounds.items[b];
        const end_pair = bounds.items[b + 1];
        // The right-hand boundary pair covers content[end_pair .. end_pair+2].
        const ngram_end = end_pair + 2;
        const ngram_len = ngram_end - start;

        if (ngram_len < MIN_LEN) continue;

        if (ngram_len <= MAX_NGRAM_LEN) {
            try result.append(allocator, makeNgram(content, start, ngram_len));
        } else {
            // Force-split into MAX_NGRAM_LEN-sized chunks.
            var off = start;
            while (off + MAX_NGRAM_LEN <= ngram_end) {
                try result.append(allocator, makeNgram(content, off, MAX_NGRAM_LEN));
                off += MAX_NGRAM_LEN;
            }
            const rem = ngram_end - off;
            if (rem >= MIN_LEN) {
                try result.append(allocator, makeNgram(content, off, rem));
            } else if (rem > 0) {
                // Tail is too short for its own ngram.  Overlap with the
                // previous chunk by backing up to ngram_end - MIN_LEN so
                // every byte in the span is covered.
                try result.append(allocator, makeNgram(content, ngram_end - MIN_LEN, MIN_LEN));
            }
        }
    }

    return result.toOwnedSlice(allocator);
}

/// Build the covering set of n-gram hashes for a query using a sliding window.
/// Extracts every substring of the query with length in [3, MAX_NGRAM_LEN] so
/// that file boundary-based n-grams overlapping the query are matched regardless
/// of where content-defined boundaries fall in the indexed file.
/// Caller owns the returned slice.
pub fn buildCoveringSet(query: []const u8, allocator: std.mem.Allocator) ![]SparseNgram {
    const MIN_LEN = 3;
    if (query.len < MIN_LEN) return try allocator.alloc(SparseNgram, 0);

    var result: std.ArrayList(SparseNgram) = .empty;
    errdefer result.deinit(allocator);

    // Slide a window of every length [MIN_LEN, MAX_NGRAM_LEN] across the query.
    // This avoids boundary-misalignment false negatives when a query substring
    // appears in the indexed file as a content-defined boundary n-gram.
    var len: usize = MIN_LEN;
    while (len <= @min(MAX_NGRAM_LEN, query.len)) : (len += 1) {
        var pos: usize = 0;
        while (pos + len <= query.len) : (pos += 1) {
            try result.append(allocator, makeNgram(query, pos, len));
        }
    }

    return result.toOwnedSlice(allocator);
}

/// In-memory sparse n-gram index.  Mirrors the TrigramIndex API so it can
/// be used as a drop-in acceleration layer alongside the trigram index.
pub const SparseNgramIndex = struct {
    /// ngram hash → set of file paths that contain the n-gram
    index: std.AutoHashMap(u64, std.StringHashMap(void)),
    /// path → list of ngram hashes contributed (for cleanup on re-index)
    file_ngrams: std.StringHashMap(std.ArrayList(u64)),
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator) SparseNgramIndex {
        return .{
            .index = std.AutoHashMap(u64, std.StringHashMap(void)).init(allocator),
            .file_ngrams = std.StringHashMap(std.ArrayList(u64)).init(allocator),
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *SparseNgramIndex) void {
        var iter = self.index.iterator();
        while (iter.next()) |entry| {
            entry.value_ptr.deinit();
        }
        self.index.deinit();

        var fn_iter = self.file_ngrams.iterator();
        while (fn_iter.next()) |entry| {
            entry.value_ptr.deinit(self.allocator);
        }
        self.file_ngrams.deinit();
    }

    pub fn removeFile(self: *SparseNgramIndex, path: []const u8) void {
        const ngrams = self.file_ngrams.getPtr(path) orelse return;
        for (ngrams.items) |hash| {
            if (self.index.getPtr(hash)) |file_set| {
                _ = file_set.remove(path);
                if (file_set.count() == 0) {
                    file_set.deinit();
                    _ = self.index.remove(hash);
                }
            }
        }
        ngrams.deinit(self.allocator);
        _ = self.file_ngrams.remove(path);
    }

    pub fn indexFile(self: *SparseNgramIndex, path: []const u8, content: []const u8) !void {
        self.removeFile(path);

        const ngrams = try extractSparseNgrams(content, self.allocator);
        defer self.allocator.free(ngrams);

        // Deduplicate hashes so the cleanup list stays compact.
        var seen = std.AutoHashMap(u64, void).init(self.allocator);
        defer seen.deinit();

        for (ngrams) |ng| {
            const gop = try self.index.getOrPut(ng.hash);
            if (!gop.found_existing) {
                gop.value_ptr.* = std.StringHashMap(void).init(self.allocator);
            }
            _ = try gop.value_ptr.getOrPut(path);
            _ = try seen.getOrPut(ng.hash);
        }

        var hash_list: std.ArrayList(u64) = .empty;
        errdefer hash_list.deinit(self.allocator);
        var seen_iter = seen.keyIterator();
        while (seen_iter.next()) |h| {
            try hash_list.append(self.allocator, h.*);
        }
        try self.file_ngrams.put(path, hash_list);
    }

    /// Find candidate files that may contain the query string.
    /// Uses the sliding-window covering set from buildCoveringSet and returns
    /// the UNION of all matching posting lists — a superset of true matches,
    /// to be verified by content search.  Returns null when the query is too
    /// short.  Caller must free the returned slice.
    pub fn candidates(self: *SparseNgramIndex, query: []const u8, allocator: std.mem.Allocator) ?[]const []const u8 {
        const ngrams = buildCoveringSet(query, allocator) catch return null;
        defer allocator.free(ngrams);

        if (ngrams.len == 0) return null;

        // Union posting sets for all sliding-window n-gram hashes.
        // A file is a candidate if it shares any substring with the query.
        var seen_files = std.StringHashMap(void).init(allocator);
        defer seen_files.deinit();

        for (ngrams) |ng| {
            const file_set = self.index.getPtr(ng.hash) orelse continue;
            var it = file_set.keyIterator();
            while (it.next()) |path_ptr| {
                seen_files.put(path_ptr.*, {}) catch return null;
            }
        }

        if (seen_files.count() == 0) {
            return allocator.alloc([]const u8, 0) catch null;
        }

        var result: std.ArrayList([]const u8) = .empty;
        errdefer result.deinit(allocator);
        result.ensureTotalCapacity(allocator, seen_files.count()) catch return null;

        var file_it = seen_files.keyIterator();
        while (file_it.next()) |path_ptr| {
            result.appendAssumeCapacity(path_ptr.*);
        }

        return result.toOwnedSlice(allocator) catch {
            result.deinit(allocator);
            return null;
        };
    }

    pub fn fileCount(self: *SparseNgramIndex) u32 {
        return @intCast(self.file_ngrams.count());
    }
};
