//! MongoDB command dispatch. Each command fills `reply` with its result;
//! unknown commands and failures produce error replies with real codes.

const std = @import("std");
const builtin = @import("builtin");
const bson = @import("bson.zig");
const wire = @import("wire.zig");
const db = @import("db.zig");
const query = @import("query.zig");
const update = @import("update.zig");
const index = @import("index.zig");

pub const Context = struct {
    gpa: std.mem.Allocator,
    io: std.Io,
    oid_gen: *bson.ObjectIdGen,
    connection_id: u32,
    client_desc: []const u8,
    engine: *db.Engine,
    server_start: std.Io.Timestamp,
};

pub const ErrorCode = enum(i32) {
    command_not_found = 59,
    bad_value = 2,
    invalid_argument = 72,
    namespace_not_found = 26,
    index_not_found = 27,
    duplicate_key = 11000,
    namespace_exists = 48,
    failed_to_parse = 9,
    internal_error = 1,
    invalid_pipeline_operator = 40324,
    index_options_conflict = 85,
    cannot_create_index = 67,
    invalid_index_specification_option = 197,
};

/// Which lock (if any) a command needs on the engine. Contract: only
/// `.write` commands may call engine mutation functions (insert, replace,
/// remove, drop*, get_or_create_collection, compact); `.read` commands may
/// only read (`get_collection`, `database_names`, `collection_names`);
/// `.none` commands must not touch the engine at all.
const CommandKind = enum { none, read, write };

const Command = struct {
    name: []const u8,
    kind: CommandKind,
    handler: *const fn (*Context, *wire.Message, *wire.Reply) anyerror!void,
};

/// The one place a command exists: name, lock class, and handler declared
/// together so a new command cannot be given a handler but no lock.
const command_table = [_]Command{
    // Handshake, admin info, no-ops: never touch the engine.
    .{ .name = "hello", .kind = .none, .handler = cmd_hello },
    .{ .name = "isMaster", .kind = .none, .handler = cmd_is_master },
    .{ .name = "ismaster", .kind = .none, .handler = cmd_is_master },
    .{ .name = "ping", .kind = .none, .handler = cmd_ping },
    .{ .name = "buildInfo", .kind = .none, .handler = cmd_build_info },
    .{ .name = "getParameter", .kind = .none, .handler = cmd_get_parameter },
    .{ .name = "whatsmyuri", .kind = .none, .handler = cmd_whatsmyuri },
    .{ .name = "hostInfo", .kind = .none, .handler = cmd_host_info },
    .{ .name = "getCmdLineOpts", .kind = .none, .handler = cmd_get_cmd_line_opts },
    .{ .name = "serverStatus", .kind = .none, .handler = cmd_server_status },
    .{ .name = "endSessions", .kind = .none, .handler = cmd_end_sessions },
    .{ .name = "connectionStatus", .kind = .none, .handler = cmd_connection_status },
    .{ .name = "getMore", .kind = .none, .handler = cmd_get_more },
    .{ .name = "killCursors", .kind = .none, .handler = cmd_kill_cursors },
    // Read-only: scan the engine without mutating it.
    .{ .name = "find", .kind = .read, .handler = cmd_find },
    .{ .name = "count", .kind = .read, .handler = cmd_count },
    .{ .name = "aggregate", .kind = .read, .handler = cmd_aggregate },
    .{ .name = "listDatabases", .kind = .read, .handler = cmd_list_databases },
    .{ .name = "listCollections", .kind = .read, .handler = cmd_list_collections },
    // Writes: exclusive, totally ordered.
    .{ .name = "create", .kind = .write, .handler = cmd_create },
    .{ .name = "drop", .kind = .write, .handler = cmd_drop },
    .{ .name = "dropDatabase", .kind = .write, .handler = cmd_drop_database },
    .{ .name = "createIndexes", .kind = .write, .handler = cmd_create_indexes },
    .{ .name = "dropIndexes", .kind = .write, .handler = cmd_drop_indexes },
    .{ .name = "insert", .kind = .write, .handler = cmd_insert },
    .{ .name = "update", .kind = .write, .handler = cmd_update },
    .{ .name = "delete", .kind = .write, .handler = cmd_delete },
    .{ .name = "findAndModify", .kind = .write, .handler = cmd_find_and_modify },
    .{ .name = "listIndexes", .kind = .read, .handler = cmd_list_indexes },
};

/// Command name to its index in `command_table`, resolved at comptime so a
/// request costs one hash instead of a walk down the whole table comparing
/// strings.
const command_index = blk: {
    var kvs: [command_table.len]struct { []const u8, usize } = undefined;
    for (&command_table, 0..) |c, i| kvs[i] = .{ c.name, i };
    break :blk std.StaticStringMap(usize).initComptime(kvs);
};

pub fn dispatch(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const name = msg.command_name();
    const cmd = if (command_index.get(name)) |i| &command_table[i] else {
        var buf: [256]u8 = undefined;
        const errmsg = try std.fmt.bufPrint(&buf, "no such command: '{s}'", .{name});
        return reply.put_error(@intFromEnum(ErrorCode.command_not_found), "CommandNotFound", errmsg);
    };

    switch (cmd.kind) {
        .none => return cmd.handler(ctx, msg, reply),
        .read => {
            try ctx.engine.lock_read();
            // Defers are block-scoped: this one is registered in the prong
            // block, so it runs when the prong exits — after the handler
            // returns. The shared lock is thus held for the whole command.
            defer ctx.engine.unlock_read();
            return cmd.handler(ctx, msg, reply);
        },
        .write => {
            try ctx.engine.lock();
            defer ctx.engine.unlock();
            return cmd.handler(ctx, msg, reply);
        },
    }
}

// ---------------------------------------------------------------------------
// Handshake / administration
// ---------------------------------------------------------------------------

fn add_server_info(ctx: *Context, reply: *wire.Reply) !void {
    try reply.put("isWritablePrimary", .{ .bool = true });
    try reply.put("maxBsonObjectSize", .{ .int32 = wire.max_bson_object_size });
    try reply.put("maxMessageSizeBytes", .{ .int32 = 48000000 });
    try reply.put("maxWriteBatchSize", .{ .int32 = 100000 });
    const now = std.Io.Timestamp.now(ctx.io, .real);
    try reply.put("localTime", .{ .datetime = now.toMilliseconds() });
    try reply.put("logicalSessionTimeoutMinutes", .{ .int32 = 30 });
    try reply.put("connectionId", .{ .int32 = @intCast(ctx.connection_id) });
    try reply.put("minWireVersion", .{ .int32 = 0 });
    try reply.put("maxWireVersion", .{ .int32 = 8 });
    try reply.put("readOnly", .{ .bool = false });

    // Deliberately no `topologyVersion`. A driver treats its presence as
    // "this server supports the streaming (awaitable) hello protocol" and
    // switches monitoring to an exhaust hello: it sends one hello with
    // maxAwaitTimeMS and the OP_MSG exhaustAllowed flag, then expects a
    // stream of unsolicited replies each carrying moreToCome. We answer
    // once with moreToCome clear and go back to reading, so the driver
    // fails the heartbeat ("Server ended moreToCome unexpectedly"), drops
    // the connection and resets its pool — a connect/disconnect loop once
    // per heartbeat, which is what MongoDB Compass showed. Omitting the
    // field keeps monitoring on the polling path, which we do implement,
    // and matches maxWireVersion 8: streaming hello arrived in wire 9.
}

fn cmd_hello(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try add_server_info(ctx, reply);
    try reply.put_ok();
}

fn cmd_is_master(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try add_server_info(ctx, reply);
    try reply.put("ismaster", .{ .bool = true });
    try reply.put("helloOk", .{ .bool = true });
    try reply.put_ok();
}

fn cmd_ping(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try reply.put_ok();
}

fn cmd_build_info(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try reply.put("version", .{ .string = "4.4.0" });
    try reply.put("gitVersion", .{ .string = "mongo-lite" });
    try reply.put("versionArray", .{ .array = try int_array(reply, &.{ 4, 4, 0, 0 }) });
    try reply.put("openssl", .{ .doc = &.{} });
    try reply.put("loaderFlags", .{ .string = "" });
    try reply.put("compilerInfo", .{ .string = "zig 0.16.0" });
    try reply.put("allocator", .{ .string = "system" });
    try reply.put("javascriptEngine", .{ .string = "none" });
    try reply.put("bits", .{ .int32 = 64 });
    try reply.put("debug", .{ .bool = false });
    try reply.put("maxBsonObjectSize", .{ .int32 = wire.max_bson_object_size });
    try reply.put("storageEngines", .{ .array = try str_array(reply, &.{"wiredTiger"}) });
    try reply.put_ok();
}

fn cmd_get_parameter(_: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    // mongosh probes featureCompatibilityVersion; respond per requested key.
    const params = doc_arg(msg.body.get("getParameter")) orelse
        return invalid_arg(reply, "getParameter requires a document");
    if (bson.get_pair(params, "featureCompatibilityVersion") == null) {
        return invalid_arg(reply, "no option found to get");
    }
    const fcv = try reply.arena_alloc().alloc(bson.Pair, 1);
    fcv[0] = .{ .key = "version", .value = .{ .string = "4.4" } };
    try reply.put("featureCompatibilityVersion", .{ .doc = fcv });
    try reply.put_ok();
}

fn cmd_whatsmyuri(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try reply.put("you", .{ .string = ctx.client_desc });
    try reply.put_ok();
}

fn cmd_host_info(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    const now = std.Io.Timestamp.now(ctx.io, .real);
    const cpu_count: usize = std.Thread.getCpuCount() catch 1;

    const system = try reply.arena_alloc().alloc(bson.Pair, 6);
    system[0] = .{ .key = "currentTime", .value = .{ .datetime = now.toMilliseconds() } };
    system[1] = .{ .key = "hostname", .value = .{ .string = "localhost" } };
    system[2] = .{ .key = "cpuAddrSize", .value = .{ .int32 = 64 } };
    system[3] = .{ .key = "memSizeMB", .value = .{ .int32 = 0 } };
    system[4] = .{ .key = "numCores", .value = .{ .int32 = @intCast(cpu_count) } };
    system[5] = .{ .key = "cpuArch", .value = .{ .string = @tagName(builtin.cpu.arch) } };
    try reply.put("system", .{ .doc = system });

    const os = try reply.arena_alloc().alloc(bson.Pair, 3);
    os[0] = .{ .key = "type", .value = .{ .string = @tagName(builtin.os.tag) } };
    os[1] = .{ .key = "name", .value = .{ .string = @tagName(builtin.os.tag) } };
    os[2] = .{ .key = "version", .value = .{ .string = "unknown" } };
    try reply.put("os", .{ .doc = os });

    try reply.put("extra", .{ .doc = &.{} });
    try reply.put_ok();
}

fn cmd_get_cmd_line_opts(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    const argv = try reply.arena_alloc().alloc(bson.Pair, 2);
    argv[0] = .{ .key = "dbpath", .value = .{ .string = "mongo-lite.log" } };
    argv[1] = .{ .key = "port", .value = .{ .int32 = 27017 } };
    try reply.put("argv", .{ .array = &.{} });
    try reply.put("parsed", .{ .doc = argv });
    try reply.put_ok();
}

fn cmd_server_status(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    const now = std.Io.Timestamp.now(ctx.io, .real);
    const uptime: i64 = std.Io.Timestamp.durationTo(ctx.server_start, now).toSeconds();

    try reply.put("host", .{ .string = "localhost" });
    try reply.put("version", .{ .string = "4.4.0" });
    try reply.put("process", .{ .string = "mongod" });
    try reply.put("uptime", .{ .double = @floatFromInt(uptime) });
    try reply.put("localTime", .{ .datetime = now.toMilliseconds() });
    const connections = try reply.arena_alloc().alloc(bson.Pair, 1);
    connections[0] = .{ .key = "current", .value = .{ .int32 = @intCast(ctx.connection_id) } };
    try reply.put("connections", .{ .doc = connections });
    try reply.put_ok();
}

fn cmd_end_sessions(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try reply.put_ok();
}

fn cmd_connection_status(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    const auth_info = try reply.arena_alloc().alloc(bson.Pair, 2);
    auth_info[0] = .{ .key = "authenticatedUsers", .value = .{ .array = &.{} } };
    auth_info[1] = .{ .key = "authenticatedUserRoles", .value = .{ .array = &.{} } };
    try reply.put("authInfo", .{ .doc = auth_info });
    try reply.put_ok();
}

fn cmd_list_databases(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    var names: std.ArrayListUnmanaged([]const u8) = .empty;
    defer names.deinit(ctx.gpa);
    try ctx.engine.database_names(&names);

    const values = try reply.arena_alloc().alloc(bson.Value, names.items.len);
    for (names.items, 0..) |n, i| {
        const entry = try reply.arena_alloc().alloc(bson.Pair, 3);
        entry[0] = .{ .key = "name", .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } };
        entry[1] = .{ .key = "sizeOnDisk", .value = .{ .double = 0 } };
        entry[2] = .{ .key = "empty", .value = .{ .bool = true } };
        values[i] = .{ .doc = entry };
    }
    try reply.put("databases", .{ .array = values });
    try reply.put("totalSize", .{ .int32 = 0 });
    try reply.put("totalSizeMb", .{ .int32 = 0 });
    try reply.put_ok();
}

fn cmd_list_collections(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "listCollections requires $db");
    var names: std.ArrayListUnmanaged([]const u8) = .empty;
    defer names.deinit(ctx.gpa);
    try ctx.engine.collection_names(db_name, &names);

    const values = try reply.arena_alloc().alloc(bson.Value, names.items.len);
    for (names.items, 0..) |n, i| {
        const entry = try reply.arena_alloc().alloc(bson.Pair, 3);
        entry[0] = .{ .key = "name", .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } };
        entry[1] = .{ .key = "type", .value = .{ .string = "collection" } };
        entry[2] = .{ .key = "options", .value = .{ .doc = &.{} } };
        values[i] = .{ .doc = entry };
    }
    try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, ""), "firstBatch", values) });
    try reply.put_ok();
}

fn cmd_create(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "create requires $db");
    const coll_name = str_arg(msg.body.get("create")) orelse return bad_value(reply, "create requires a collection name");
    _ = try ctx.engine.get_or_create_collection(db_name, coll_name);
    try reply.put_ok();
}

fn cmd_drop(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "drop requires $db");
    const coll_name = str_arg(msg.body.get("drop")) orelse return bad_value(reply, "drop requires a collection name");
    if (!try ctx.engine.drop_collection(db_name, coll_name)) {
        return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found");
    }
    try reply.put_ok();
}

fn cmd_drop_database(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "dropDatabase requires $db");
    _ = try ctx.engine.drop_database(db_name);
    try reply.put("dropped", .{ .string = try reply.arena_alloc().dupe(u8, db_name) });
    try reply.put_ok();
}

// ---------------------------------------------------------------------------
// Indexes
// ---------------------------------------------------------------------------

fn cmd_create_indexes(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "createIndexes requires $db");
    const coll_name = str_arg(msg.body.get("createIndexes")) orelse return bad_value(reply, "createIndexes requires a collection name");
    const indexes = switch (msg.body.get("indexes") orelse return bad_value(reply, "createIndexes requires indexes")) {
        .array => |a| a,
        else => return bad_value(reply, "indexes must be an array"),
    };

    const existed = ctx.engine.get_collection(db_name, coll_name) != null;
    const coll = try ctx.engine.get_or_create_collection(db_name, coll_name);
    const num_before: i32 = @intCast(coll.indexes.items.len + 1); // + the _id_ index

    for (indexes) |spec_v| {
        const spec = switch (spec_v) {
            .doc => |p| p,
            else => return bad_value(reply, "indexes must be documents"),
        };
        const key_value = bson.get_pair(spec, "key") orelse return bad_value(reply, "index spec requires key");
        const key_pairs = switch (key_value) {
            .doc => |p| p,
            else => return bad_value(reply, "key must be a document"),
        };
        if (key_pairs.len == 0) return bad_value(reply, "cannot create index with an empty key");

        // {_id: 1} is the implicit index: an idempotent no-op. Any other
        // secondary index touching _id is rejected.
        var has_id = false;
        for (key_pairs) |p| {
            if (std.mem.eql(u8, p.key, "_id")) has_id = true;
        }
        if (has_id) {
            const only = key_pairs[0].value;
            const is_id_index = key_pairs.len == 1 and ((only == .int32 and only.int32 == 1) or
                (only == .int64 and only.int64 == 1) or
                (only == .double and only.double == 1.0));
            if (is_id_index) {
                // The no-op must not swallow options that would change what
                // the index does — MongoDB rejects a TTL _id index rather
                // than quietly ignoring the expiry.
                if (bson.get_pair(spec, "expireAfterSeconds") != null) {
                    return reply.put_error(
                        @intFromEnum(ErrorCode.invalid_index_specification_option),
                        "InvalidIndexSpecificationOption",
                        "the field 'expireAfterSeconds' is not valid for an _id index specification",
                    );
                }
                continue;
            }
            return bad_value(reply, "cannot create a secondary index on the _id field");
        }
        const name = bson.get_pair(spec, "name") orelse bson.Value.null;
        if (name == .string and std.mem.eql(u8, name.string, "_id_")) {
            return bad_value(reply, "cannot create index with name '_id_'");
        }

        const spec_doc = bson.Document{ .arena = undefined, .pairs = spec };
        _ = ctx.engine.create_index(db_name, coll_name, &spec_doc) catch |err| switch (err) {
            error.InvalidIndexSpec => return bad_value(reply, "invalid index spec"),
            error.TtlOnCompoundIndex => return reply.put_error(
                @intFromEnum(ErrorCode.cannot_create_index),
                "CannotCreateIndex",
                "TTL indexes are single-field indexes, compound indexes do not support TTL",
            ),
            error.InvalidExpireAfterSeconds => return reply.put_error(
                @intFromEnum(ErrorCode.cannot_create_index),
                "CannotCreateIndex",
                "TTL index 'expireAfterSeconds' option must be a whole number between 0 and 2147483647",
            ),
            error.IndexOptionsConflict => return reply.put_error(@intFromEnum(ErrorCode.index_options_conflict), "IndexOptionsConflict", "index already exists with a different specification"),
            error.DuplicateKeyIndex => {
                const ix_name = if (name == .string) name.string else "index";
                const msg_text = try e11000_message(reply, db_name, coll_name, ix_name, try render_spec_key(reply, key_pairs));
                return reply.put_error(@intFromEnum(ErrorCode.duplicate_key), "DuplicateKey", msg_text);
            },
            error.ParallelArrays => return bad_value(reply, "cannot index parallel arrays"),
            else => return err,
        };
    }

    try reply.put("createdCollectionAutomatically", .{ .bool = !existed });
    try reply.put("numIndexesBefore", .{ .int32 = num_before });
    try reply.put("numIndexesAfter", .{ .int32 = @intCast(coll.indexes.items.len + 1) });
    try reply.put_ok();
}

fn cmd_list_indexes(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "listIndexes requires $db");
    const coll_name = str_arg(msg.body.get("listIndexes")) orelse return bad_value(reply, "listIndexes requires a collection name");
    const coll = ctx.engine.get_collection(db_name, coll_name) orelse
        return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found");

    // The _id_ index first, then the secondaries.
    const n = coll.indexes.items.len + 1;
    const values = try reply.arena_alloc().alloc(bson.Value, n);
    const id_pairs = try reply.arena_alloc().alloc(bson.Pair, 2);
    id_pairs[0] = .{ .key = "v", .value = .{ .int32 = 2 } };
    id_pairs[1] = .{ .key = "key", .value = .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} } };
    values[0] = .{ .doc = try index_pairs_append(reply, id_pairs, "_id_") };
    for (coll.indexes.items, 0..) |*ix, i| {
        // The pairs live in the reply arena (freed with it); the values
        // array below references them.
        var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
        try ix.spec_pairs(reply.arena_alloc(), &pairs);
        values[1 + i] = .{ .doc = pairs.items };
    }
    try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) });
    try reply.put_ok();
}

/// The _id_ index entry: {v, key: {_id: 1}, name: "_id_"}.
fn index_pairs_append(reply: *wire.Reply, pairs: []const bson.Pair, name: []const u8) ![]const bson.Pair {
    const arena = reply.arena_alloc();
    const with_name = try arena.alloc(bson.Pair, pairs.len + 1);
    @memcpy(with_name[0..pairs.len], pairs);
    with_name[pairs.len] = .{ .key = "name", .value = .{ .string = try arena.dupe(u8, name) } };
    return with_name;
}

fn cmd_drop_indexes(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "dropIndexes requires $db");
    const coll_name = str_arg(msg.body.get("dropIndexes")) orelse return bad_value(reply, "dropIndexes requires a collection name");
    const coll = ctx.engine.get_collection(db_name, coll_name) orelse
        return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found");
    const n_indexes_was: i32 = @intCast(coll.indexes.items.len + 1);

    const arg = msg.body.get("index") orelse return bad_value(reply, "dropIndexes requires index");
    if (arg == .string and std.mem.eql(u8, arg.string, "*")) {
        // Drop every secondary index. Copy the names first: each drop
        // mutates the collection's index list.
        var names: std.ArrayListUnmanaged([]const u8) = .empty;
        defer names.deinit(ctx.gpa);
        for (coll.indexes.items) |ix| try names.append(ctx.gpa, ix.name);
        for (names.items) |nm| _ = try ctx.engine.drop_index(db_name, coll_name, nm);
    } else if (arg == .string) {
        if (std.mem.eql(u8, arg.string, "_id_")) {
            return invalid_arg(reply, "cannot drop the _id index");
        }
        if (!try ctx.engine.drop_index(db_name, coll_name, arg.string)) {
            return reply.put_error(@intFromEnum(ErrorCode.index_not_found), "IndexNotFound", "index not found with name");
        }
    } else if (arg == .doc) {
        // A key document: drop the index with a matching key pattern.
        const key_value = bson.get_pair(arg.doc, "key") orelse bson.Value{ .doc = arg.doc };
        const key_pairs = switch (key_value) {
            .doc => |p| p,
            else => return bad_value(reply, "dropIndexes index must be a name, key document, or '*'"),
        };
        const target = index.find_by_key_pattern(coll.indexes.items, key_pairs) orelse
            return reply.put_error(@intFromEnum(ErrorCode.index_not_found), "IndexNotFound", "index not found with key pattern");
        _ = try ctx.engine.drop_index(db_name, coll_name, target.name);
    } else {
        return bad_value(reply, "dropIndexes index must be a name, key document, or '*'");
    }

    try reply.put("nIndexesWas", .{ .int32 = n_indexes_was });
    try reply.put_ok();
}

/// The E11000 text drivers parse. One definition for both create-time and
/// write-time conflicts; they differ only in how the dup key is rendered.
fn e11000_message(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, index_name: []const u8, key_text: []const u8) ![]const u8 {
    return std.fmt.allocPrint(
        reply.arena_alloc(),
        "E11000 duplicate key error collection: {s}.{s} index: {s} dup key: {s}",
        .{ db_name, coll_name, index_name, key_text },
    );
}

/// Render the key pattern with placeholder values for a createIndexes
/// duplicate-key error (no specific document is involved).
fn render_spec_key(reply: *wire.Reply, key_pairs: []const bson.Pair) ![]const u8 {
    const arena = reply.arena_alloc();
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(arena);
    try out.append(arena, '{');
    for (key_pairs, 0..) |p, i| {
        if (i > 0) try out.appendSlice(arena, ", ");
        try out.appendSlice(arena, p.key);
        try out.appendSlice(arena, ": ?");
    }
    try out.append(arena, '}');
    return out.toOwnedSlice(arena);
}

// ---------------------------------------------------------------------------
// CRUD
// ---------------------------------------------------------------------------

fn cmd_insert(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "insert requires $db");
    const coll_name = str_arg(msg.body.get("insert")) orelse return bad_value(reply, "insert requires a collection name");

    const docs = try batch_arg(msg, reply, "insert", "documents") orelse return;

    var inserted: i64 = 0;
    var write_errors: std.ArrayListUnmanaged(bson.Value) = .empty;
    defer write_errors.deinit(reply.arena_alloc());

    // Group commit: one fsync for the whole batch instead of one per
    // document. end_batch runs on every return path, so even a failed doc
    // (writeErrors) or a hard error still syncs what was appended.
    ctx.engine.begin_batch();
    defer ctx.engine.end_batch() catch {};

    for (docs, 0..) |*doc, i| {
        if (ctx.engine.insert(db_name, coll_name, doc, ctx.oid_gen)) |_| {
            inserted += 1;
        } else |err| {
            switch (err) {
                error.DuplicateKey, error.DuplicateKeyIndex => {
                    const e = try reply.arena_alloc().alloc(bson.Pair, 3);
                    e[0] = .{ .key = "index", .value = .{ .int32 = @intCast(i) } };
                    e[1] = .{ .key = "code", .value = .{ .int32 = @intFromEnum(ErrorCode.duplicate_key) } };
                    e[2] = .{ .key = "errmsg", .value = .{ .string = try duplicate_key_message(ctx, reply, db_name, coll_name, doc) } };
                    try write_errors.append(reply.arena_alloc(), .{ .doc = e });
                },
                else => return err,
            }
        }
    }

    try reply.put("n", .{ .int32 = @intCast(inserted) });
    if (write_errors.items.len > 0) {
        // Copy into the reply arena: write_errors is freed when this command
        // returns, before the reply is serialized.
        const arr = try reply.arena_alloc().alloc(bson.Value, write_errors.items.len);
        @memcpy(arr, write_errors.items);
        try reply.put("writeErrors", .{ .array = arr });
    }
    try reply.put_ok();
}

fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "find requires $db");
    const coll_name = str_arg(msg.body.get("find")) orelse return bad_value(reply, "find requires a collection name");

    const filter = doc_arg(msg.body.get("filter")) orelse return bad_value(reply, "filter must be a document");

    const sort_keys = try parse_sort_keys(reply, msg.body.get("sort"));
    const proj_pairs = doc_arg(msg.body.get("projection"));
    const skip: u64 = int_arg(msg.body.get("skip")) orelse 0;
    // A negative limit means "return this many in a single batch"; we always
    // reply with one batch, so only the magnitude matters.
    const limit: usize = @abs(int_value(msg.body.get("limit")) orelse 0);

    var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
    defer matched.deinit(ctx.gpa);
    // Documents needed to fill the page, counting the skipped prefix; 0
    // means unbounded.
    const page_end: usize = if (limit == 0) 0 else blk: {
        const skip_usize = std.math.cast(usize, skip) orelse break :blk 0;
        break :blk skip_usize +| limit;
    };
    // Stop scanning once the page is filled. Only sound without a sort,
    // which has to see every match before it can tell which ones the page
    // contains.
    _ = try scan_matching(ctx, db_name, coll_name, filter, if (sort_keys.len > 0) 0 else page_end, &matched);

    if (sort_keys.len > 0) {
        // Selecting the page is much cheaper than ordering everything when
        // the page is a small fraction of the matches. Above that fraction
        // the heap's bookkeeping stops paying for itself.
        if (page_end > 0 and page_end *| 4 <= matched.items.len) {
            try query.sort_docs_top_k(reply.arena_alloc(), matched.items, sort_keys, page_end);
        } else {
            try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys);
        }
    }
    const rest = if (skip < matched.items.len) matched.items[skip..] else &.{};
    const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest;
    try emit_docs(reply, db_name, coll_name, proj_pairs, page);
    try reply.put_ok();
}

/// Collect the documents in `db_name.coll_name` matching `filter`, stopping
/// after `limit` matches (0 = unlimited). Returns the number matched; `out`
/// may be null when only the count is wanted. Every command that scans a
/// collection goes through here, so an index only needs this one call site.
///
/// Candidate generation order: the _id_ fast path (docs map lookup), then a
/// secondary-index plan, then a plain scan. Every candidate is re-checked
/// with the unchanged filter, so an index that over-approximates is merely
/// slow — never wrong. With no index created and no usable _id clause, the
/// plain-scan path is the only one reached.
fn scan_matching(
    ctx: *Context,
    db_name: []const u8,
    coll_name: []const u8,
    filter: []const bson.Pair,
    limit: usize,
    out: ?*std.ArrayListUnmanaged(*const bson.Document),
) !usize {
    const coll = ctx.engine.get_collection(db_name, coll_name) orelse return 0;
    const filter_doc = bson.Document{ .arena = undefined, .pairs = filter };
    var n: usize = 0;

    // _id_ fast path: the docs map is the _id index. Skipped when the
    // queried value's compare-equivalence class is serialization-ambiguous
    // (see index.plan_id).
    if (try index.plan_id(ctx.gpa, filter)) |id_plan| {
        var plan = id_plan;
        defer plan.deinit(ctx.gpa);
        // One scratch key, rebuilt per value: a key is never needed past its
        // own lookup.
        var key: std.ArrayListUnmanaged(u8) = .empty;
        defer key.deinit(ctx.gpa);
        for (plan.values) |v| {
            key.clearRetainingCapacity();
            try bson.write_serialized_value(v, ctx.gpa, &key);
            const doc = coll.docs.get(key.items) orelse continue;
            if (!try query.matches(ctx.gpa, &filter_doc, doc)) continue;
            if (out) |list| try list.append(ctx.gpa, doc);
            n += 1;
            if (limit != 0 and n >= limit) break;
        }
        return n;
    }

    // Secondary-index plan: candidates in index order, re-filtered. The
    // returned ids alias the docs map keys, valid under the read lock.
    if (try index.plan(ctx.gpa, coll.indexes.items, filter)) |p| {
        var plan = p;
        defer plan.deinit(ctx.gpa);
        var ids: std.ArrayListUnmanaged([]const u8) = .empty;
        defer ids.deinit(ctx.gpa);
        try plan.search(ctx.gpa, &ids);
        for (ids.items) |id| {
            const doc = coll.docs.get(id) orelse continue;
            if (!try query.matches(ctx.gpa, &filter_doc, doc)) continue;
            if (out) |list| try list.append(ctx.gpa, doc);
            n += 1;
            if (limit != 0 and n >= limit) break;
        }
        return n;
    }

    var it = coll.docs.iterator();
    while (it.next()) |entry| {
        if (!try query.matches(ctx.gpa, &filter_doc, entry.value_ptr.*)) continue;
        if (out) |list| try list.append(ctx.gpa, entry.value_ptr.*);
        n += 1;
        if (limit != 0 and n >= limit) break;
    }
    return n;
}

fn emit_docs(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, proj_pairs: ?[]const bson.Pair, docs: []const *const bson.Document) !void {
    const values = try reply.arena_alloc().alloc(bson.Value, docs.len);
    for (docs, 0..) |d, i| {
        values[i] = try project_doc(reply, d, proj_pairs);
    }
    try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) });
}

/// Project a stored doc (or deep-copy it) into the reply arena.
fn project_doc(reply: *wire.Reply, doc: *const bson.Document, proj_pairs: ?[]const bson.Pair) !bson.Value {
    if (proj_pairs) |pp| {
        var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
        errdefer out.deinit(reply.arena_alloc());
        try query.project(reply.arena_alloc(), doc, &.{ .arena = undefined, .pairs = pp }, &out);
        return .{ .doc = out.items };
    }
    return .{ .doc = try bson.copy_pairs(reply.arena_alloc(), doc.pairs) };
}

fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "update requires $db");
    const coll_name = str_arg(msg.body.get("update")) orelse return bad_value(reply, "update requires a collection name");

    const specs = try batch_arg(msg, reply, "update", "updates") orelse return;

    var n_matched: i64 = 0;
    var n_modified: i64 = 0;
    var upserted: std.ArrayListUnmanaged(bson.Pair) = .empty;
    defer upserted.deinit(reply.arena_alloc());
    var write_errors: std.ArrayListUnmanaged(bson.Value) = .empty;
    defer write_errors.deinit(reply.arena_alloc());

    // Group commit for multi-document updates: one fsync per command.
    ctx.engine.begin_batch();
    defer ctx.engine.end_batch() catch {};

    for (specs, 0..) |*spec, si| {
        const q = doc_arg(spec.get("q")) orelse return bad_value(reply, "update spec requires q");
        const u_doc = doc_arg(spec.get("u")) orelse return bad_value(reply, "update spec requires u");
        const multi = bool_arg(spec.get("multi")) orelse false;
        const upsert = bool_arg(spec.get("upsert")) orelse false;

        var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
        defer matched.deinit(ctx.gpa);
        _ = try scan_matching(ctx, db_name, coll_name, q, if (multi) 0 else 1, &matched);

        if (matched.items.len == 0) {
            if (upsert) {
                const new_doc = try build_upsert_doc(reply, q, u_doc);
                ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) {
                    error.DuplicateKey, error.DuplicateKeyIndex => return duplicate_key_error(ctx, reply, db_name, coll_name, new_doc),
                    else => return err,
                };
                const id = new_doc.get("_id") orelse bson.Value.null;
                const u = try reply.arena_alloc().alloc(bson.Pair, 2);
                u[0] = .{ .key = "index", .value = .{ .int32 = @intCast(si) } };
                u[1] = .{ .key = "_id", .value = try bson.copy_value(reply.arena_alloc(), id) };
                try upserted.append(reply.arena_alloc(), .{ .key = "u", .value = .{ .doc = u } });
                n_matched += 1;
            }
            continue;
        }

        n_matched += @intCast(matched.items.len);
        for (matched.items) |doc| {
            // Work on a copy: the log write must precede any visible change,
            // and a rejected update must not corrupt the stored document.
            const copy = try clone_doc(reply, doc);
            update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }) catch |err| switch (err) {
                error.ImmutableId, error.InvalidUpdate => return bad_value(reply, "bad update"),
                else => return err,
            };
            ctx.engine.replace(db_name, coll_name, copy, ctx.oid_gen) catch |err| switch (err) {
                error.DuplicateKey, error.DuplicateKeyIndex => {
                    const e = try reply.arena_alloc().alloc(bson.Pair, 3);
                    e[0] = .{ .key = "index", .value = .{ .int32 = @intCast(si) } };
                    e[1] = .{ .key = "code", .value = .{ .int32 = @intFromEnum(ErrorCode.duplicate_key) } };
                    e[2] = .{ .key = "errmsg", .value = .{ .string = try duplicate_key_message(ctx, reply, db_name, coll_name, copy) } };
                    try write_errors.append(reply.arena_alloc(), .{ .doc = e });
                    continue;
                },
                else => return err,
            };
            n_modified += 1;
        }
    }

    try reply.put("n", .{ .int32 = @intCast(n_matched) });
    try reply.put("nModified", .{ .int32 = @intCast(n_modified) });
    if (write_errors.items.len > 0) {
        const arr = try reply.arena_alloc().alloc(bson.Value, write_errors.items.len);
        @memcpy(arr, write_errors.items);
        try reply.put("writeErrors", .{ .array = arr });
    }
    if (upserted.items.len > 0) {
        const arr = try reply.arena_alloc().alloc(bson.Value, upserted.items.len);
        for (upserted.items, 0..) |u, i| arr[i] = u.value;
        try reply.put("upserted", .{ .array = arr });
    }
    try reply.put_ok();
}

fn cmd_delete(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "delete requires $db");
    const coll_name = str_arg(msg.body.get("delete")) orelse return bad_value(reply, "delete requires a collection name");

    const specs = try batch_arg(msg, reply, "delete", "deletes") orelse return;

    var n_deleted: i64 = 0;
    // Group commit for multi-document deletes: one fsync per command.
    ctx.engine.begin_batch();
    defer ctx.engine.end_batch() catch {};

    for (specs) |*spec| {
        const q = doc_arg(spec.get("q")) orelse return bad_value(reply, "delete spec requires q");
        const limit = int_value(spec.get("limit")) orelse 1;

        var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
        defer matched.deinit(ctx.gpa);
        _ = try scan_matching(ctx, db_name, coll_name, q, if (limit == 1) 1 else 0, &matched);
        for (matched.items) |doc| {
            const id = doc.get("_id") orelse continue;
            if (try ctx.engine.remove_by_id(db_name, coll_name, id)) n_deleted += 1;
        }
    }

    try reply.put("n", .{ .int32 = @intCast(n_deleted) });
    try reply.put_ok();
}

fn cmd_find_and_modify(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "findAndModify requires $db");
    const coll_name = str_arg(msg.body.get("findAndModify")) orelse return bad_value(reply, "findAndModify requires a collection name");

    const q = doc_arg(msg.body.get("query")) orelse &.{};
    const sort_keys = try parse_sort_keys(reply, msg.body.get("sort"));
    const remove = bool_arg(msg.body.get("remove")) orelse false;
    const do_update = msg.body.get("update") != null;
    const upsert = bool_arg(msg.body.get("upsert")) orelse false;
    const ret_new = bool_arg(msg.body.get("new")) orelse false;
    const proj_pairs = doc_arg(msg.body.get("fields"));

    if (remove and do_update) return bad_value(reply, "remove and update are mutually exclusive");
    if (!remove and !do_update) return bad_value(reply, "must specify update or remove");

    var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
    defer matched.deinit(ctx.gpa);
    // Without a sort, only the first match is ever used.
    _ = try scan_matching(ctx, db_name, coll_name, q, if (sort_keys.len > 0) 0 else 1, &matched);
    if (sort_keys.len > 0) {
        try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys);
    }

    const arena = reply.arena_alloc();
    const target = if (matched.items.len > 0) matched.items[0] else null;

    // Each branch decides what the reply says; the tail below emits it once.
    var n: i32 = 0;
    var updated_existing = false;
    var upserted_id: ?bson.Value = null;
    var value: bson.Value = .null;

    if (target == null and do_update and upsert) {
        const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document");
        const new_doc = try build_upsert_doc(reply, q, u_doc);
        ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) {
            error.DuplicateKey, error.DuplicateKeyIndex => return duplicate_key_error(ctx, reply, db_name, coll_name, new_doc),
            else => return err,
        };
        n = 1;
        upserted_id = try bson.copy_value(arena, new_doc.get("_id") orelse bson.Value.null);
        if (ret_new) value = try project_doc(reply, new_doc, proj_pairs);
    } else if (target != null and remove) {
        n = 1;
        // Project before removing: this reads the stored document.
        value = try project_doc(reply, target.?, proj_pairs);
        _ = try ctx.engine.remove_by_id(db_name, coll_name, target.?.get("_id") orelse unreachable);
    } else if (target != null and do_update) {
        const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document");
        const before = try bson.copy_pairs(arena, target.?.pairs);
        const copy = try clone_doc(reply, target.?);
        update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }) catch |err| switch (err) {
            error.ImmutableId, error.InvalidUpdate => return bad_value(reply, "bad update"),
            else => return err,
        };
        try ctx.engine.replace(db_name, coll_name, copy, ctx.oid_gen);
        n = 1;
        updated_existing = true;
        value = if (ret_new) try project_doc(reply, copy, proj_pairs) else .{ .doc = before };
    } // else: no match and no upsert — an empty result

    var leo: std.ArrayListUnmanaged(bson.Pair) = .empty;
    defer leo.deinit(arena);
    try leo.append(arena, .{ .key = "n", .value = .{ .int32 = n } });
    try leo.append(arena, .{ .key = "updatedExisting", .value = .{ .bool = updated_existing } });
    if (upserted_id) |id| try leo.append(arena, .{ .key = "upserted", .value = id });
    try reply.put("value", value);
    try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(arena) });
    try reply.put_ok();
}

fn cmd_count(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "count requires $db");
    const coll_name = str_arg(msg.body.get("count")) orelse return bad_value(reply, "count requires a collection name");
    const q = doc_arg(msg.body.get("query")) orelse &.{};

    const n = try scan_matching(ctx, db_name, coll_name, q, 0, null);
    try reply.put("n", .{ .int32 = @intCast(n) });
    try reply.put_ok();
}

fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
    const db_name = msg.db_name() orelse return invalid_arg(reply, "aggregate requires $db");
    const coll_name = str_arg(msg.body.get("aggregate")) orelse return bad_value(reply, "aggregate requires a collection name");

    const pipeline_value = msg.body.get("pipeline") orelse return bad_value(reply, "aggregate requires pipeline");
    var stages = switch (pipeline_value) {
        .array => |a| a,
        else => return bad_value(reply, "pipeline must be an array"),
    };

    // countDocuments() reaches us as [{$match: F}?, {$group: {_id: <literal>,
    // n: {$sum: 1}}}]. The general path answers that by materializing every
    // matching document and then throwing them all away, so recognize the
    // shape and answer it from a counting scan instead.
    if (try count_only_pipeline(reply, stages)) |shape| {
        const n = try scan_matching(ctx, db_name, coll_name, shape.filter, 0, null);
        // No documents means no groups at all, not a group holding zero —
        // same as the general path, which builds groups per document.
        var docs: []const *const bson.Document = &.{};
        if (n > 0) {
            const arena = reply.arena_alloc();
            const pairs = try arena.alloc(bson.Pair, 1 + shape.accs.len);
            pairs[0] = .{ .key = "_id", .value = try bson.copy_value(arena, shape.id_value) };
            for (shape.accs, 0..) |acc, i| {
                // Mirrors run_group's coercion exactly: an integral sum in
                // int32 range comes back as int32, otherwise a double.
                const sum: f64 = @as(f64, @floatFromInt(n)) * acc.term;
                pairs[1 + i] = .{
                    .key = acc.key,
                    .value = if (sum == @floor(sum) and sum <= 2_147_483_647 and sum >= -2_147_483_648)
                        .{ .int32 = @intFromFloat(sum) }
                    else
                        .{ .double = sum },
                };
            }
            const doc = try arena.create(bson.Document);
            doc.* = bson.Document{ .arena = undefined, .pairs = pairs };
            const one = try arena.alloc(*const bson.Document, 1);
            one[0] = doc;
            docs = one;
        }
        try emit_docs(reply, db_name, coll_name, null, docs);
        return reply.put_ok();
    }

    // The pipeline operates on a stream of documents; each stage transforms
    // the current window [start, end) of `stream`, and $group replaces the
    // stream entirely (so $sort/$limit after it apply to the groups).
    var stream: std.ArrayListUnmanaged(*const bson.Document) = .empty;
    defer stream.deinit(ctx.gpa);
    // A leading $match is pushed down into an indexed candidate scan; the
    // stage is then dropped from the pipeline so it is not applied twice.
    if (stages.len > 0 and stages[0] == .doc and stages[0].doc.len > 0 and std.mem.eql(u8, stages[0].doc[0].key, "$match")) {
        const filter = doc_arg(stages[0].doc[0].value) orelse return bad_value(reply, "$match requires a document");
        _ = try scan_matching(ctx, db_name, coll_name, filter, 0, &stream);
        stages = stages[1..];
    } else if (ctx.engine.get_collection(db_name, coll_name)) |coll| {
        var it = coll.docs.iterator();
        while (it.next()) |entry| try stream.append(ctx.gpa, entry.value_ptr.*);
    }

    var start: usize = 0;
    var end: usize = stream.items.len;
    var count_stage: ?[]const u8 = null;
    var proj_pairs: ?[]const bson.Pair = null;

    for (stages) |stage_v| {
        const stage = switch (stage_v) {
            .doc => |pairs| pairs,
            else => return bad_value(reply, "pipeline stages must be documents"),
        };
        if (stage.len == 0) continue;
        const stage_name = stage[0].key;
        if (std.mem.eql(u8, stage_name, "$match")) {
            const filter = doc_arg(stage[0].value) orelse return bad_value(reply, "$match requires a document");
            var kept: std.ArrayListUnmanaged(*const bson.Document) = .empty;
            defer kept.deinit(ctx.gpa);
            for (stream.items[start..end]) |d| {
                if (try query.matches(ctx.gpa, &.{ .arena = undefined, .pairs = filter }, d)) {
                    try kept.append(ctx.gpa, d);
                }
            }
            stream.deinit(ctx.gpa);
            stream = kept;
            kept = .empty;
            start = 0;
            end = stream.items.len;
        } else if (std.mem.eql(u8, stage_name, "$sort")) {
            const keys = try parse_sort_keys(reply, stage[0].value);
            if (keys.len > 0) {
                try query.sort_docs(reply.arena_alloc(), stream.items[start..end], keys);
            }
        } else if (std.mem.eql(u8, stage_name, "$skip")) {
            const n = try stage_count(reply, stage[0].value, "$skip") orelse return;
            start = @min(start + n, end);
        } else if (std.mem.eql(u8, stage_name, "$limit")) {
            const n = try stage_count(reply, stage[0].value, "$limit") orelse return;
            end = @min(end, start + n);
        } else if (std.mem.eql(u8, stage_name, "$project")) {
            proj_pairs = doc_arg(stage[0].value);
        } else if (std.mem.eql(u8, stage_name, "$group")) {
            const gp = doc_arg(stage[0].value) orelse return bad_value(reply, "$group requires a document");
            const grouped_opt = try run_group(ctx, reply, gp, stream.items[start..end]);
            const grouped = grouped_opt orelse return;
            // Group results replace the stream: later stages see groups.
            stream.deinit(ctx.gpa);
            stream = grouped;
            start = 0;
            end = stream.items.len;
        } else if (std.mem.eql(u8, stage_name, "$count")) {
            count_stage = switch (stage[0].value) {
                .string => |s| s,
                else => return bad_value(reply, "$count requires a string"),
            };
        } else {
            const msg_text = try std.fmt.allocPrint(reply.arena_alloc(), "Unrecognized pipeline stage name: '{s}'", .{stage_name});
            return reply.put_error(@intFromEnum(ErrorCode.invalid_pipeline_operator), "InvalidPipelineOperator", msg_text);
        }
    }

    const slice = stream.items[start..end];

    if (count_stage) |name| {
        const c = try reply.arena_alloc().alloc(bson.Pair, 1);
        c[0] = .{ .key = try reply.arena_alloc().dupe(u8, name), .value = .{ .int32 = @intCast(slice.len) } };
        const values = try reply.arena_alloc().alloc(bson.Value, 1);
        values[0] = .{ .doc = c };
        try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) });
    } else {
        try emit_docs(reply, db_name, coll_name, proj_pairs, slice);
    }
    try reply.put_ok();
}

/// A pipeline whose whole answer is the number of matching documents.
const CountShape = struct {
    filter: []const bson.Pair,
    /// The literal every document groups under.
    id_value: bson.Value,
    accs: []const Acc,

    const Acc = struct { key: []const u8, term: f64 };
};

/// Recognize `[{$match: F}?, {$group: {_id: <literal>, k: {$sum: <number>}}}]`
/// — the shape a driver sends for countDocuments().
///
/// Deliberately conservative: a `_id` of `"$field"`, an accumulator over a
/// field, or any other stage needs the documents themselves, so anything
/// that is not exactly this shape returns null and takes the general path.
fn count_only_pipeline(reply: *wire.Reply, stages: []const bson.Value) !?CountShape {
    if (stages.len == 0 or stages.len > 2) return null;

    var filter: []const bson.Pair = &.{};
    if (stages.len == 2) {
        const first = switch (stages[0]) {
            .doc => |p| p,
            else => return null,
        };
        if (first.len != 1 or !std.mem.eql(u8, first[0].key, "$match")) return null;
        filter = doc_arg(first[0].value) orelse return null;
    }

    const last = switch (stages[stages.len - 1]) {
        .doc => |p| p,
        else => return null,
    };
    if (last.len != 1 or !std.mem.eql(u8, last[0].key, "$group")) return null;
    const gp = doc_arg(last[0].value) orelse return null;

    const id_value = bson.get_pair(gp, "_id") orelse return null;
    switch (id_value) {
        // A field path or a computed id groups per document.
        .string => |s| if (s.len > 0 and s[0] == '$') return null,
        .doc, .array => return null,
        else => {},
    }

    var accs: std.ArrayListUnmanaged(CountShape.Acc) = .empty;
    for (gp) |p| {
        if (std.mem.eql(u8, p.key, "_id")) continue;
        const spec = switch (p.value) {
            .doc => |d| d,
            else => return null,
        };
        if (spec.len != 1 or !std.mem.eql(u8, spec[0].key, "$sum")) return null;
        const term: f64 = switch (spec[0].value) {
            .int32 => |i| @floatFromInt(i),
            .int64 => |i| @floatFromInt(i),
            .double => |d| d,
            // $sum over a field depends on the documents.
            else => return null,
        };
        try accs.append(reply.arena_alloc(), .{ .key = p.key, .term = term });
    }

    return .{ .filter = filter, .id_value = id_value, .accs = accs.items };
}

/// Minimal $group: supports `_id` of null/literal/"$field" and `$sum`
/// accumulators (constant or "$field").
fn run_group(ctx: *Context, reply: *wire.Reply, group_pairs: []const bson.Pair, docs: []const *const bson.Document) !?std.ArrayListUnmanaged(*const bson.Document) {
    const arena = reply.arena_alloc();
    const id_expr = bson.get_pair(group_pairs, "_id") orelse {
        try bad_value(reply, "$group requires _id");
        return null;
    };

    var accs: std.ArrayListUnmanaged(bson.Pair) = .empty;
    defer accs.deinit(arena);
    for (group_pairs) |p| {
        if (std.mem.eql(u8, p.key, "_id")) continue;
        try accs.append(arena, p);
    }

    const Group = struct {
        id_value: bson.Value,
        sums: []f64,
    };
    var groups: std.StringHashMapUnmanaged(Group) = .empty;
    defer groups.deinit(ctx.gpa);
    // StringHashMapUnmanaged does not copy keys; keep them alive until done.
    var keys_owned: std.ArrayListUnmanaged([]u8) = .empty;
    defer {
        for (keys_owned.items) |k| ctx.gpa.free(k);
        keys_owned.deinit(ctx.gpa);
    }

    var id_key_buf: std.ArrayListUnmanaged(u8) = .empty;
    defer id_key_buf.deinit(ctx.gpa);
    for (docs) |doc| {
        const id_value: bson.Value = switch (id_expr) {
            .string => |s| if (s.len > 0 and s[0] == '$') query_path_value(doc, s[1..]) orelse .null else id_expr,
            else => id_expr,
        };
        id_key_buf.clearRetainingCapacity();
        try bson.write_value(id_value, ctx.gpa, &id_key_buf);

        const gop = try groups.getOrPut(ctx.gpa, id_key_buf.items);
        if (!gop.found_existing) {
            // Only a new group needs an owned copy of the key; the map
            // borrows it, so keys_owned keeps it alive until we are done.
            const key = try ctx.gpa.dupe(u8, id_key_buf.items);
            try keys_owned.append(ctx.gpa, key);
            gop.key_ptr.* = key;
            const sums = try ctx.gpa.alloc(f64, accs.items.len);
            @memset(sums, 0);
            gop.value_ptr.* = .{ .id_value = id_value, .sums = sums };
        }
        for (accs.items, 0..) |acc, i| {
            var expr = acc.value;
            // Unwrap {$sum: <expr>} accumulator documents.
            if (expr == .doc) {
                if (bson.get_pair(expr.doc, "$sum")) |inner| {
                    expr = inner;
                } else continue;
            }
            const term: f64 = switch (expr) {
                .int32 => |n| @floatFromInt(n),
                .int64 => |n| @floatFromInt(n),
                .double => |n| n,
                .string => |s| if (s.len > 0 and s[0] == '$')
                    switch (query_path_value(doc, s[1..]) orelse .null) {
                        .int32 => |n| @floatFromInt(n),
                        .int64 => |n| @floatFromInt(n),
                        .double => |n| n,
                        else => 0,
                    }
                else
                    0,
                else => 0,
            };
            gop.value_ptr.sums[i] += term;
        }
    }

    var out: std.ArrayListUnmanaged(*const bson.Document) = .empty;
    errdefer out.deinit(ctx.gpa);
    var it = groups.iterator();
    // free sum arrays
    defer {
        var git = groups.iterator();
        while (git.next()) |e| ctx.gpa.free(e.value_ptr.sums);
    }
    while (it.next()) |entry| {
        const npairs = 1 + accs.items.len;
        const pairs = try arena.alloc(bson.Pair, npairs);
        pairs[0] = .{ .key = "_id", .value = try bson.copy_value(arena, entry.value_ptr.id_value) };
        for (accs.items, 0..) |acc, i| {
            const sum: f64 = entry.value_ptr.sums[i];
            const sum_value: bson.Value = if (sum == @floor(sum) and sum <= 2_147_483_647 and sum >= -2_147_483_648)
                .{ .int32 = @intFromFloat(sum) }
            else
                .{ .double = sum };
            pairs[1 + i] = .{ .key = acc.key, .value = sum_value };
        }
        const doc = try arena.create(bson.Document);
        doc.* = bson.Document{ .arena = undefined, .pairs = pairs };
        try out.append(ctx.gpa, doc);
    }
    return out;
}

/// Resolve a simple "$field" path expression inside a document.
fn query_path_value(doc: *const bson.Document, path: []const u8) ?bson.Value {
    var cur: bson.Value = undefined;
    var it = std.mem.splitScalar(u8, path, '.');
    const first = it.next() orelse return null;
    cur = bson.get_pair(doc.pairs, first) orelse return null;
    while (it.next()) |seg| {
        cur = switch (cur) {
            .doc => |pairs| bson.get_pair(pairs, seg) orelse return null,
            else => return null,
        };
    }
    return cur;
}

fn cmd_get_more(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    // Cursors are never left open, so getMore always yields an empty batch.
    try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, "test.$cmd", "nextBatch", &.{}) });
    try reply.put_ok();
}

fn cmd_kill_cursors(_: *Context, _: *wire.Message, reply: *wire.Reply) !void {
    try reply.put("cursorsKilled", .{ .array = &.{} });
    try reply.put("cursorsNotFound", .{ .array = &.{} });
    try reply.put("cursorsAlive", .{ .array = &.{} });
    try reply.put_ok();
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Deep-copy a stored document into the reply arena so updates can be
/// applied off the live doc.
fn clone_doc(reply: *wire.Reply, doc: *const bson.Document) !*bson.Document {
    const arena = reply.arena_alloc();
    const owned = try arena.create(bson.Document);
    owned.* = .{
        .arena = std.heap.ArenaAllocator.init(arena),
        .pairs = try bson.copy_pairs(arena, doc.pairs),
    };
    return owned;
}

/// Build the document for an upsert: equality fields from the filter, then
/// the update operators applied. Owned by the reply arena.
fn build_upsert_doc(reply: *wire.Reply, q: []const bson.Pair, u_doc: []const bson.Pair) !*bson.Document {
    const arena = reply.arena_alloc();
    var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
    defer pairs.deinit(arena);
    for (q) |p| {
        const is_operator = p.key.len > 0 and p.key[0] == '$';
        const is_embedded_operator = p.value == .doc and query.all_operator_keys(p.value.doc);
        if (!is_operator and !is_embedded_operator) {
            try pairs.append(arena, .{ .key = try arena.dupe(u8, p.key), .value = try bson.copy_value(arena, p.value) });
        }
    }
    const owned = try arena.create(bson.Document);
    owned.* = bson.Document{ .arena = std.heap.ArenaAllocator.init(arena), .pairs = try pairs.toOwnedSlice(arena) };
    // Apply update operators to build the final doc; _id handled by insert.
    update.apply(owned, &.{ .arena = undefined, .pairs = u_doc }) catch |err| switch (err) {
        error.ImmutableId, error.InvalidUpdate => return error.InvalidUpdate,
        else => return err,
    };
    return owned;
}

fn parse_sort_keys(reply: *wire.Reply, value: ?bson.Value) ![]const query.SortKey {
    const pairs = doc_arg(value) orelse return &.{};
    const out = try reply.arena_alloc().alloc(query.SortKey, pairs.len);
    for (pairs, 0..) |p, i| {
        const descending = switch (p.value) {
            .int32 => |n| n < 0,
            .int64 => |n| n < 0,
            .double => |n| n < 0,
            .string => |s| std.mem.eql(u8, s, "desc"),
            else => false,
        };
        out[i] = .{
            .path = try reply.arena_alloc().dupe(u8, p.key),
            .descending = descending,
        };
    }
    return out;
}

fn doc_arg(v: ?bson.Value) ?[]const bson.Pair {
    return switch (v orelse return null) {
        .doc => |pairs| pairs,
        else => null,
    };
}

fn str_arg(v: ?bson.Value) ?[]const u8 {
    return switch (v orelse return null) {
        .string => |s| s,
        else => null,
    };
}

fn bool_arg(v: ?bson.Value) ?bool {
    return switch (v orelse return null) {
        .bool => |b| b,
        .int32 => |i| i != 0,
        else => null,
    };
}

fn int_arg(v: ?bson.Value) ?u64 {
    const i = int_value(v) orelse return null;
    return if (i < 0) null else @intCast(i);
}

/// Numeric argument as a signed integer, with no clamping — callers that care
/// about the sign (find's limit, delete's limit, $skip) apply their own rule.
fn int_value(v: ?bson.Value) ?i64 {
    return switch (v orelse return null) {
        .int32 => |i| i,
        .int64 => |i| i,
        // lossyCast saturates instead of trapping on out-of-range doubles.
        .double => |d| std.math.lossyCast(i64, d),
        else => null,
    };
}

/// A `$skip`/`$limit` stage operand: a non-negative document count. Writes
/// the error reply and returns null when the operand is not a number.
fn stage_count(reply: *wire.Reply, v: bson.Value, stage: []const u8) !?usize {
    const n = int_value(v) orelse {
        const text = try std.fmt.allocPrint(reply.arena_alloc(), "{s} requires a number", .{stage});
        try bad_value(reply, text);
        return null;
    };
    return @intCast(@max(0, n));
}

/// Fetch a batch argument, writing the standard error reply and returning
/// null when it is missing or malformed.
fn batch_arg(msg: *wire.Message, reply: *wire.Reply, cmd: []const u8, name: []const u8) !?[]const bson.Document {
    return msg.batch(name) catch |err| {
        const arena = reply.arena_alloc();
        const text = switch (err) {
            error.MissingBatch => try std.fmt.allocPrint(arena, "{s} requires {s}", .{ cmd, name }),
            error.BatchNotArray => try std.fmt.allocPrint(arena, "{s} must be an array", .{name}),
            error.BatchElementNotDoc => try std.fmt.allocPrint(arena, "{s} must be documents", .{name}),
            error.OutOfMemory => return error.OutOfMemory,
        };
        try bad_value(reply, text);
        return null;
    };
}

fn invalid_arg(reply: *wire.Reply, msg: []const u8) !void {
    return reply.put_error(@intFromEnum(ErrorCode.invalid_argument), "InvalidArgument", msg);
}

fn bad_value(reply: *wire.Reply, msg: []const u8) !void {
    return reply.put_error(@intFromEnum(ErrorCode.bad_value), "BadValue", msg);
}

/// The E11000 text, shared by the top-level error reply and the per-document
/// `writeErrors` entries of a batch insert. Uses engine.dup_index (set by a
/// rejected unique-index write) when the conflict came from a secondary
/// index; otherwise it is the _id_ index.
fn duplicate_key_message(ctx: *Context, reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document) ![]const u8 {
    var index_name: []const u8 = "_id_";
    var key_text: []const u8 = undefined;
    if (ctx.engine.dup_index) |name| {
        index_name = name;
        key_text = try render_dup_key(ctx, reply, db_name, coll_name, name, doc);
    } else {
        key_text = try serialize_value_compact(reply, doc.get("_id") orelse bson.Value.null);
    }
    return e11000_message(reply, db_name, coll_name, index_name, key_text);
}

/// Render the dup key of a secondary index from the offending document: the
/// document's values for the index key pattern.
fn render_dup_key(ctx: *Context, reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, index_name: []const u8, doc: *const bson.Document) ![]const u8 {
    const arena = reply.arena_alloc();
    const coll = ctx.engine.get_collection(db_name, coll_name) orelse
        // Index not found (defensive): fall back to the document's _id.
        return serialize_value_compact(reply, doc.get("_id") orelse bson.Value.null);
    const ix = coll.find_index(index_name) orelse
        return serialize_value_compact(reply, doc.get("_id") orelse bson.Value.null);

    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(arena);
    try out.append(arena, '{');
    for (ix.keys, 0..) |k, i| {
        if (i > 0) try out.appendSlice(arena, ", ");
        try out.appendSlice(arena, k.path);
        try out.appendSlice(arena, ": ");
        var values: std.ArrayListUnmanaged(bson.Value) = .empty;
        defer values.deinit(arena);
        try query.collect_values(arena, doc.pairs, k.path, &values, 0);
        const v: bson.Value = if (values.items.len > 0) values.items[0] else .null;
        try out.appendSlice(arena, try serialize_value_compact(reply, v));
    }
    try out.append(arena, '}');
    return out.toOwnedSlice(arena);
}

fn duplicate_key_error(ctx: *Context, reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document) !void {
    const msg_text = try duplicate_key_message(ctx, reply, db_name, coll_name, doc);
    return reply.put_error(@intFromEnum(ErrorCode.duplicate_key), "DuplicateKey", msg_text);
}

/// Compact extended-JSON-ish rendering of a value for error messages.
fn serialize_value_compact(reply: *wire.Reply, v: bson.Value) ![]const u8 {
    const arena = reply.arena_alloc();
    return switch (v) {
        .int32 => |i| try std.fmt.allocPrint(arena, "{d}", .{i}),
        .int64 => |i| try std.fmt.allocPrint(arena, "{d}", .{i}),
        .double => |d| try std.fmt.allocPrint(arena, "{d}", .{d}),
        .string => |s| try std.fmt.allocPrint(arena, "'{s}'", .{s}),
        .bool => |b| try std.fmt.allocPrint(arena, "{}", .{b}),
        .object_id => |oid| blk: {
            const hex = std.fmt.bytesToHex(oid[0..], .lower);
            break :blk try std.fmt.allocPrint(arena, "ObjectId('{s}')", .{hex});
        },
        else => "{ ... }",
    };
}

/// Build a { id: <n>, ns: "...", <batch_key>: [...] } cursor document.
pub fn cursor_doc(reply: *wire.Reply, cursor_id: i64, ns: []const u8, batch_key: []const u8, docs: []const bson.Value) ![]const bson.Pair {
    const c = try reply.arena_alloc().alloc(bson.Pair, 3);
    c[0] = .{ .key = "id", .value = .{ .int64 = cursor_id } };
    c[1] = .{ .key = "ns", .value = .{ .string = ns } };
    c[2] = .{ .key = batch_key, .value = .{ .array = docs } };
    return c;
}

fn format_namespace(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8) ![]const u8 {
    return std.fmt.allocPrint(reply.arena_alloc(), "{s}.{s}", .{ db_name, coll_name });
}

fn int_array(reply: *wire.Reply, values: []const i32) ![]const bson.Value {
    const out = try reply.arena_alloc().alloc(bson.Value, values.len);
    for (values, 0..) |v, i| out[i] = .{ .int32 = v };
    return out;
}

fn str_array(reply: *wire.Reply, values: []const []const u8) ![]const bson.Value {
    const out = try reply.arena_alloc().alloc(bson.Value, values.len);
    for (values, 0..) |v, i| out[i] = .{ .string = try reply.arena_alloc().dupe(u8, v) };
    return out;
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

const testing = std.testing;

/// Temp-log-backed engine plus the Context a command test dispatches with.
const TestDb = struct {
    tmp: std.testing.TmpDir,
    path: []u8,
    engine: db.Engine,
    gen: bson.ObjectIdGen,

    fn init(io: std.Io) !TestDb {
        const tmp = std.testing.tmpDir(.{});
        const path = try std.fmt.allocPrint(testing.allocator, ".zig-cache/tmp/{s}/cmd.log", .{tmp.sub_path});
        return .{
            .tmp = tmp,
            .path = path,
            .engine = try db.Engine.open(testing.allocator, io, path),
            .gen = bson.ObjectIdGen.init(io),
        };
    }

    fn deinit(self: *TestDb) void {
        self.engine.deinit();
        self.tmp.cleanup();
        testing.allocator.free(self.path);
    }

    fn ctx(self: *TestDb, io: std.Io) Context {
        return test_ctx(io, &self.engine, &self.gen, 1);
    }
};

fn test_ctx(io: std.Io, engine: *db.Engine, gen: *bson.ObjectIdGen, connection_id: u32) Context {
    return .{
        .gpa = testing.allocator,
        .io = io,
        .oid_gen = gen,
        .connection_id = connection_id,
        .client_desc = "127.0.0.1:0",
        .engine = engine,
        .server_start = std.Io.Timestamp.now(io, .real),
    };
}

test "ping and hello replies parse" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();
    var ctx = tdb.ctx(io);

    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    var msg = try parse_fake_msg("ping", .null, &.{});
    defer msg.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(f64, 1.0), reply.pairs.items[0].value.double);

    var reply2 = wire.Reply.init(testing.allocator);
    defer reply2.deinit();
    var msg2 = try parse_fake_msg("hello", .null, &.{});
    defer msg2.deinit();
    try dispatch(&ctx, &msg2, &reply2);
    const ok = bson.get_pair(reply2.pairs.items, "ok").?;
    try testing.expectEqual(@as(f64, 1.0), ok.double);
    const primary = bson.get_pair(reply2.pairs.items, "isWritablePrimary").?;
    try testing.expect(primary.bool);
    try testing.expectEqual(@as(i32, 8), bson.get_pair(reply2.pairs.items, "maxWireVersion").?.int32);
}

test "handshake does not advertise the streaming hello protocol" {
    // `topologyVersion` in a hello/isMaster reply tells the driver it may
    // monitor with an exhaust hello and expect moreToCome replies we never
    // send; the driver then kills the connection every heartbeat.
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();
    var ctx = tdb.ctx(io);

    for ([_][]const u8{ "hello", "isMaster", "ismaster" }) |cmd| {
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        var msg = try parse_fake_msg(cmd, .null, &.{});
        defer msg.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
        try testing.expect(bson.get_pair(reply.pairs.items, "topologyVersion") == null);
    }
}

test "unknown command gives CommandNotFound" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();
    var ctx = tdb.ctx(io);
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    var msg = try parse_fake_msg("nonsenseCmd", .null, &.{});
    defer msg.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(f64, 0.0), bson.get_pair(reply.pairs.items, "ok").?.double);
    try testing.expectEqual(@as(i32, 59), bson.get_pair(reply.pairs.items, "code").?.int32);
    try testing.expectEqualStrings("CommandNotFound", bson.get_pair(reply.pairs.items, "codeName").?.string);
}

test "getParameter responds for featureCompatibilityVersion" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();
    var ctx = tdb.ctx(io);
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    var fcv_doc = try testing.allocator.alloc(bson.Pair, 1);
    defer testing.allocator.free(fcv_doc);
    fcv_doc[0] = .{ .key = "featureCompatibilityVersion", .value = .{ .int32 = 1 } };
    var msg = try parse_fake_msg("getParameter", .{ .doc = fcv_doc }, &.{});
    defer msg.deinit();
    try dispatch(&ctx, &msg, &reply);
    const fcv = bson.get_pair(reply.pairs.items, "featureCompatibilityVersion").?;
    try testing.expectEqualStrings("4.4", fcv.doc[0].value.string);
}

test "insert counts only successful inserts" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();
    var ctx = tdb.ctx(io);

    // documents: [{_id:1}, {_id:1} (dup), {_id:2}] → n: 2 + one writeError.
    const docs = [_]bson.Value{
        .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} },
        .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} },
        .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 2 } }} },
    };
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    var msg = try parse_fake_msg("insert", .{ .string = "users" }, &.{
        .{ .key = "documents", .value = .{ .array = &docs } },
    });
    defer msg.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(i64, 2), bson.get_pair(reply.pairs.items, "n").?.int32);
    const errors = bson.get_pair(reply.pairs.items, "writeErrors").?;
    try testing.expectEqual(@as(usize, 1), errors.array.len);
    const first = errors.array[0].doc[0];
    try testing.expectEqualStrings("index", first.key);
    try testing.expectEqual(@as(i64, 1), first.value.int32);
}

fn parse_fake_msg(name: []const u8, value: bson.Value, extra: []const bson.Pair) !wire.Message {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(testing.allocator);
    const pairs = try testing.allocator.alloc(bson.Pair, 2 + extra.len);
    defer testing.allocator.free(pairs);
    pairs[0] = .{ .key = name, .value = value };
    pairs[1] = .{ .key = "$db", .value = .{ .string = "test" } };
    @memcpy(pairs[2..], extra);
    try bson.write_doc(pairs, testing.allocator, &out);
    var msg: std.ArrayListUnmanaged(u8) = .empty;
    defer msg.deinit(testing.allocator);
    try msg.appendSlice(testing.allocator, &[_]u8{ 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0xDD, 0x07, 0, 0, 0, 0, 0, 0 });
    try msg.append(testing.allocator, 0x00);
    try msg.appendSlice(testing.allocator, out.items);
    std.mem.writeInt(u32, msg.items[0..4], @intCast(msg.items.len), .little);
    std.mem.writeInt(i32, msg.items[12..16], wire.op_code_msg, .little);
    return wire.Message.parse(testing.allocator, msg.items);
}

test "concurrent insert/find commands on a threaded Io" {
    // Exercises dispatch's lock classification end-to-end: writer fibers run
    // `insert` under the exclusive lock, reader fibers run `count`/`find`
    // under the shared lock. Every committed insert must be visible once all
    // writers finish, and no reader may observe more docs than can exist.
    var threaded: std.Io.Threaded = .init(testing.allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var gen = bson.ObjectIdGen.init(io);

    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    const path = try std.fmt.allocPrint(testing.allocator, ".zig-cache/tmp/{s}/conc.log", .{tmp.sub_path});
    defer testing.allocator.free(path);
    var engine = try db.Engine.open(testing.allocator, io, path);
    defer engine.deinit();

    const writers = 4;
    const readers = 4;
    const per_writer: i32 = 150;
    const total: i32 = writers * per_writer;
    var next_id = std.atomic.Value(i32).init(1);
    var remaining = std.atomic.Value(usize).init(@intCast(total));

    const Worker = struct {
        fn writer(iow: std.Io, eng: *db.Engine, id_counter: *std.atomic.Value(i32), pending: *std.atomic.Value(usize), fiber_id: u32, total_writes: i32) error{Canceled}!void {
            var wgen = bson.ObjectIdGen.init(iow);
            var ctx = test_ctx(iow, eng, &wgen, fiber_id);
            while (true) {
                const id = id_counter.fetchAdd(1, .monotonic);
                if (id > total_writes) return;
                const docs = [_]bson.Value{.{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = id } }} }};
                var msg = parse_fake_msg("insert", .{ .string = "users" }, &.{
                    .{ .key = "documents", .value = .{ .array = &docs } },
                }) catch return error.Canceled;
                defer msg.deinit();
                var reply = wire.Reply.init(testing.allocator);
                defer reply.deinit();
                dispatch(&ctx, &msg, &reply) catch return error.Canceled;
                _ = pending.fetchSub(1, .monotonic);
            }
        }

        fn reader(iow: std.Io, eng: *db.Engine, pending: *std.atomic.Value(usize), fiber_id: u32, total_writes: i32) error{Canceled}!void {
            var rgen = bson.ObjectIdGen.init(iow);
            var ctx = test_ctx(iow, eng, &rgen, fiber_id);
            while (pending.load(.acquire) > 0) {
                var msg = parse_fake_msg("count", .{ .string = "users" }, &.{}) catch return error.Canceled;
                defer msg.deinit();
                var reply = wire.Reply.init(testing.allocator);
                defer reply.deinit();
                dispatch(&ctx, &msg, &reply) catch return error.Canceled;
                const n = bson.get_pair(reply.pairs.items, "n") orelse return error.Canceled;
                if (n.int32 > total_writes) return error.Canceled;
            }
        }
    };

    var group: std.Io.Group = .init;
    defer group.cancel(io);
    for (0..readers) |i| group.async(io, Worker.reader, .{ io, &engine, &remaining, @intCast(i + 1), total });
    for (0..writers) |i| group.async(io, Worker.writer, .{ io, &engine, &next_id, &remaining, @intCast(i + 1), total });
    try group.await(io);

    var ctx = test_ctx(io, &engine, &gen, 0);
    var msg = try parse_fake_msg("count", .{ .string = "users" }, &.{});
    defer msg.deinit();
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(total, bson.get_pair(reply.pairs.items, "n").?.int32);
}

// -- index command tests ----------------------------------------------------

/// Dispatch an insert of the given documents (each a bson.Value .doc).
fn dispatch_insert(tdb: *TestDb, io: std.Io, coll_name: []const u8, docs: []const bson.Value) !void {
    var ctx = tdb.ctx(io);
    var msg = try parse_fake_msg("insert", .{ .string = coll_name }, &.{
        .{ .key = "documents", .value = .{ .array = docs } },
    });
    defer msg.deinit();
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
}

/// Dispatch createIndexes for one spec.
fn dispatch_create_index(tdb: *TestDb, io: std.Io, coll_name: []const u8, spec: bson.Value) !void {
    var ctx = tdb.ctx(io);
    const specs = [_]bson.Value{spec};
    var msg = try parse_fake_msg("createIndexes", .{ .string = coll_name }, &.{
        .{ .key = "indexes", .value = .{ .array = &specs } },
    });
    defer msg.deinit();
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
}

/// Dispatch find and append the serialized _id of every result to `out`.
fn dispatch_find_ids(tdb: *TestDb, io: std.Io, coll_name: []const u8, filter: []const bson.Pair, out: *std.ArrayListUnmanaged([]u8)) !void {
    var ctx = tdb.ctx(io);
    var msg = try parse_fake_msg("find", .{ .string = coll_name }, &.{
        .{ .key = "filter", .value = .{ .doc = filter } },
    });
    defer msg.deinit();
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    try dispatch(&ctx, &msg, &reply);
    try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
    const cursor = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult;
    const batch = switch (bson.get_pair(cursor.doc, "firstBatch") orelse return error.TestUnexpectedResult) {
        .array => |a| a,
        else => return error.TestUnexpectedResult,
    };
    for (batch) |d| {
        const idv = bson.get_pair(d.doc, "_id") orelse continue;
        try out.append(testing.allocator, try bson.serialize_value(testing.allocator, idv));
    }
    std.mem.sort([]u8, out.items, {}, less_u8);
}

fn less_u8(_: void, a: []u8, b: []u8) bool {
    return std.mem.order(u8, a, b) == .lt;
}

test "createIndexes, listIndexes, dropIndexes, and idempotent re-create" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();

    // createIndexes builds the index immediately.
    {
        var ctx = tdb.ctx(io);
        const specs = [_]bson.Value{.{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "email", .value = .{ .int32 = 1 } }} } },
            .{ .key = "name", .value = .{ .string = "email_1" } },
        } }};
        var msg = try parse_fake_msg("createIndexes", .{ .string = "users" }, &.{
            .{ .key = "indexes", .value = .{ .array = &specs } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
        try testing.expectEqual(@as(i64, 1), bson.get_pair(reply.pairs.items, "numIndexesBefore").?.int32);
        try testing.expectEqual(@as(i64, 2), bson.get_pair(reply.pairs.items, "numIndexesAfter").?.int32);
    }

    // Idempotent re-create of the same spec.
    try dispatch_create_index(&tdb, io, "users", .{ .doc = &.{
        .{ .key = "key", .value = .{ .doc = &.{.{ .key = "email", .value = .{ .int32 = 1 } }} } },
        .{ .key = "name", .value = .{ .string = "email_1" } },
    } });

    // listIndexes: _id_ first, then the secondary.
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("listIndexes", .{ .string = "users" }, &.{});
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
        const cursor = bson.get_pair(reply.pairs.items, "cursor").?;
        const batch = bson.get_pair(cursor.doc, "firstBatch").?.array;
        try testing.expectEqual(@as(usize, 2), batch.len);
        try testing.expectEqualStrings("_id_", bson.get_pair(batch[0].doc, "name").?.string);
        try testing.expectEqualStrings("email_1", bson.get_pair(batch[1].doc, "name").?.string);
        const key_pairs = bson.get_pair(batch[1].doc, "key").?.doc;
        try testing.expect(bson.get_pair(key_pairs, "email") != null);
    }

    // listIndexes on a missing collection: NamespaceNotFound.
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("listIndexes", .{ .string = "nope" }, &.{});
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(i32, 26), bson.get_pair(reply.pairs.items, "code").?.int32);
    }

    // dropIndexes("*") removes the secondary but keeps _id_.
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("dropIndexes", .{ .string = "users" }, &.{
            .{ .key = "index", .value = .{ .string = "*" } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double);
        try testing.expectEqual(@as(i64, 2), bson.get_pair(reply.pairs.items, "nIndexesWas").?.int32);
    }
    try testing.expectEqual(@as(usize, 0), tdb.engine.get_collection("test", "users").?.indexes.items.len);

    // dropIndexes("_id_") errors.
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("dropIndexes", .{ .string = "users" }, &.{
            .{ .key = "index", .value = .{ .string = "_id_" } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(@as(i32, 72), bson.get_pair(reply.pairs.items, "code").?.int32);
    }
}

test "TTL index round-trips through createIndexes/listIndexes; bad specs give 67" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();

    // A driver sends expireAfterSeconds as a double.
    try dispatch_create_index(&tdb, io, "sessions", .{ .doc = &.{
        .{ .key = "key", .value = .{ .doc = &.{.{ .key = "expireAt", .value = .{ .int32 = 1 } }} } },
        .{ .key = "name", .value = .{ .string = "expireAt_1" } },
        .{ .key = "expireAfterSeconds", .value = .{ .double = 60.0 } },
    } });
    try testing.expectEqual(@as(?i64, 60), tdb.engine.get_collection("test", "sessions").?.indexes.items[0].ttl);

    // listIndexes reports it back.
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("listIndexes", .{ .string = "sessions" }, &.{});
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        const cursor = bson.get_pair(reply.pairs.items, "cursor").?;
        const batch = bson.get_pair(cursor.doc, "firstBatch").?.array;
        try testing.expectEqual(@as(usize, 2), batch.len);
        try testing.expectEqual(@as(i32, 60), bson.get_pair(batch[1].doc, "expireAfterSeconds").?.int32);
        // The _id_ entry never carries one.
        try testing.expect(bson.get_pair(batch[0].doc, "expireAfterSeconds") == null);
    }

    // Same name and key, different expiry: IndexOptionsConflict, as in
    // MongoDB (changing it is a collMod, which this server does not have).
    const bad = [_]struct { spec: bson.Value, code: i32 }{
        .{ .code = 85, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "expireAt", .value = .{ .int32 = 1 } }} } },
            .{ .key = "name", .value = .{ .string = "expireAt_1" } },
            .{ .key = "expireAfterSeconds", .value = .{ .int32 = 90 } },
        } } },
        // TTL on a compound key.
        .{ .code = 67, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{
                .{ .key = "a", .value = .{ .int32 = 1 } },
                .{ .key = "b", .value = .{ .int32 = 1 } },
            } } },
            .{ .key = "expireAfterSeconds", .value = .{ .int32 = 60 } },
        } } },
        // Negative and non-numeric expiries.
        .{ .code = 67, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} } },
            .{ .key = "expireAfterSeconds", .value = .{ .int32 = -1 } },
        } } },
        .{ .code = 67, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} } },
            .{ .key = "expireAfterSeconds", .value = .{ .string = "60" } },
        } } },
        // Past MongoDB's 2147483647 bound.
        .{ .code = 67, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} } },
            .{ .key = "expireAfterSeconds", .value = .{ .int64 = index.max_expire_after_seconds + 1 } },
        } } },
        // {_id: 1} is otherwise an idempotent no-op, but an expiry on it
        // would be silently dropped, so it is rejected instead.
        .{ .code = 197, .spec = .{ .doc = &.{
            .{ .key = "key", .value = .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} } },
            .{ .key = "expireAfterSeconds", .value = .{ .int32 = 60 } },
        } } },
    };
    for (bad) |case| {
        var ctx = tdb.ctx(io);
        const specs = [_]bson.Value{case.spec};
        var msg = try parse_fake_msg("createIndexes", .{ .string = "sessions" }, &.{
            .{ .key = "indexes", .value = .{ .array = &specs } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
        try testing.expectEqual(case.code, bson.get_pair(reply.pairs.items, "code").?.int32);
    }
    // Nothing partial was registered by the rejected specs.
    try testing.expectEqual(@as(usize, 1), tdb.engine.get_collection("test", "sessions").?.indexes.items.len);
}

test "unique index constraint returns 11000 through insert and update" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();

    try dispatch_create_index(&tdb, io, "users", .{ .doc = &.{
        .{ .key = "key", .value = .{ .doc = &.{.{ .key = "email", .value = .{ .int32 = 1 } }} } },
        .{ .key = "name", .value = .{ .string = "email_1" } },
        .{ .key = "unique", .value = .{ .bool = true } },
    } });

    const docs = [_]bson.Value{
        .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "email", .value = .{ .string = "a@x.io" } } } },
        .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 2 } }, .{ .key = "email", .value = .{ .string = "a@x.io" } } } },
        .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 3 } }, .{ .key = "email", .value = .{ .string = "b@x.io" } } } },
    };
    var ctx = tdb.ctx(io);
    var msg = try parse_fake_msg("insert", .{ .string = "users" }, &.{
        .{ .key = "documents", .value = .{ .array = &docs } },
    });
    defer msg.deinit();
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();
    try dispatch(&ctx, &msg, &reply);
    // 2 inserted, 1 writeError with code 11000 naming the index.
    try testing.expectEqual(@as(i64, 2), bson.get_pair(reply.pairs.items, "n").?.int32);
    const errors = bson.get_pair(reply.pairs.items, "writeErrors").?.array;
    try testing.expectEqual(@as(usize, 1), errors.len);
    try testing.expectEqual(@as(i64, 11000), bson.get_pair(errors[0].doc, "code").?.int32);
    const errmsg = bson.get_pair(errors[0].doc, "errmsg").?.string;
    try testing.expect(std.mem.indexOf(u8, errmsg, "email_1") != null);
    try testing.expect(std.mem.indexOf(u8, errmsg, "E11000") != null);

    // An update that collides is reported as a writeError with code 11000.
    const updates = [_]bson.Value{.{ .doc = &.{
        .{ .key = "q", .value = .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 3 } }} } },
        .{ .key = "u", .value = .{ .doc = &.{.{ .key = "$set", .value = .{ .doc = &.{.{ .key = "email", .value = .{ .string = "a@x.io" } }} } }} } },
    } }};
    var msg2 = try parse_fake_msg("update", .{ .string = "users" }, &.{
        .{ .key = "updates", .value = .{ .array = &updates } },
    });
    defer msg2.deinit();
    var reply2 = wire.Reply.init(testing.allocator);
    defer reply2.deinit();
    try dispatch(&ctx, &msg2, &reply2);
    try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply2.pairs.items, "ok").?.double);
    try testing.expectEqual(@as(i64, 0), bson.get_pair(reply2.pairs.items, "nModified").?.int32);
    const up_errs = bson.get_pair(reply2.pairs.items, "writeErrors").?.array;
    try testing.expectEqual(@as(usize, 1), up_errs.len);
    try testing.expectEqual(@as(i64, 11000), bson.get_pair(up_errs[0].doc, "code").?.int32);
}


/// Free a list of serialized ids (each element is gpa-owned).
fn free_id_list(gpa: std.mem.Allocator, list: *std.ArrayListUnmanaged([]u8)) void {
    for (list.items) |id| gpa.free(id);
    list.deinit(gpa);
}

/// Free the serialized ids and reset the list, keeping capacity.
fn clear_id_list(gpa: std.mem.Allocator, list: *std.ArrayListUnmanaged([]u8)) void {
    for (list.items) |id| gpa.free(id);
    list.clearRetainingCapacity();
}

test "count_only_pipeline accepts only shapes a count can answer" {
    // The fast path skips materializing documents, so mis-accepting a
    // pipeline would silently return a wrong aggregate rather than a slow
    // one. Pin exactly which shapes it claims.
    var reply = wire.Reply.init(testing.allocator);
    defer reply.deinit();

    const group_count = bson.Value{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
        .{ .key = "_id", .value = .null },
        .{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } },
    } } }} };
    const match_k = bson.Value{ .doc = &.{.{ .key = "$match", .value = .{ .doc = &.{.{ .key = "k", .value = .{ .int32 = 2 } }} } }} };

    // Accepted: the two shapes countDocuments() produces.
    try testing.expect(try count_only_pipeline(&reply, &.{group_count}) != null);
    const with_match = try count_only_pipeline(&reply, &.{ match_k, group_count });
    try testing.expect(with_match != null);
    try testing.expectEqual(@as(usize, 1), with_match.?.filter.len);
    try testing.expectEqualStrings("k", with_match.?.filter[0].key);

    // Rejected: grouping by a field value needs the documents.
    try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
        .{ .key = "_id", .value = .{ .string = "$k" } },
        .{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } },
    } } }} }}) == null);

    // Rejected: summing a field, not a constant.
    try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
        .{ .key = "_id", .value = .null },
        .{ .key = "t", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .string = "$x" } }} } },
    } } }} }}) == null);

    // Rejected: an accumulator we do not model at all.
    try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
        .{ .key = "_id", .value = .null },
        .{ .key = "m", .value = .{ .doc = &.{.{ .key = "$max", .value = .{ .string = "$x" } }} } },
    } } }} }}) == null);

    // Rejected: any extra stage, since it could reshape the result.
    try testing.expect(try count_only_pipeline(&reply, &.{ match_k, group_count, .{ .doc = &.{.{ .key = "$limit", .value = .{ .int32 = 1 } }} } }) == null);
    // Rejected: a leading stage that is not $match.
    try testing.expect(try count_only_pipeline(&reply, &.{ .{ .doc = &.{.{ .key = "$sort", .value = .{ .doc = &.{.{ .key = "k", .value = .{ .int32 = 1 } }} } }} }, group_count }) == null);
    // Rejected: empty pipeline.
    try testing.expect(try count_only_pipeline(&reply, &.{}) == null);
}

test "indexed queries are equivalent to scans over a mixed corpus" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    const io = threaded.io();
    var tdb = try TestDb.init(io);
    defer tdb.deinit();

    // Corpus deliberately mixes numeric _id encodings (int32 1 and int64 1
    // compare equal but hash differently), arrays, nested docs, missing
    // fields, explicit nulls, and duplicate values.
    const corpus = [_]bson.Value{
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 1 } },
            .{ .key = "a", .value = .{ .int32 = 10 } },
            .{ .key = "b", .value = .{ .string = "x" } },
            .{ .key = "tags", .value = .{ .array = &.{ .{ .string = "a" }, .{ .string = "b" } } } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int64 = 1 } },
            .{ .key = "a", .value = .{ .int32 = 20 } },
            .{ .key = "b", .value = .{ .string = "y" } },
            .{ .key = "tags", .value = .{ .array = &.{} } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 3 } },
            .{ .key = "a", .value = .{ .double = 30.0 } },
            .{ .key = "b", .value = .null },
            .{ .key = "tags", .value = .{ .array = &.{ .{ .string = "a" } } } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .string = "s4" } },
            .{ .key = "a", .value = .{ .int32 = 10 } },
            .{ .key = "b", .value = .{ .string = "z" } },
            .{ .key = "tags", .value = .{ .array = &.{ .{ .int32 = 10 }, .{ .int32 = 20 } } } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 5 } },
            .{ .key = "a", .value = .null },
            .{ .key = "b", .value = .{ .string = "x" } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 6 } },
            .{ .key = "a", .value = .{ .doc = &.{.{ .key = "n", .value = .{ .int32 = 1 } }} } },
            .{ .key = "b", .value = .{ .string = "w" } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 7 } },
            .{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 2 } } } },
            .{ .key = "b", .value = .{ .string = "q" } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 8 } },
            .{ .key = "c", .value = .{ .int32 = 99 } },
        } },
        .{ .doc = &.{
            .{ .key = "_id", .value = .{ .int32 = 9 } },
            .{ .key = "a", .value = .{ .int32 = 40 } },
            .{ .key = "b", .value = .{ .array = &.{ .{ .int32 = 5 }, .{ .int32 = 6 } } } },
        } },
    };
    try dispatch_insert(&tdb, io, "eq", &corpus);

    const filters = [_]struct { pairs: []const bson.Pair }{
        .{ .pairs = &.{} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .int32 = 10 } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$eq", .value = .{ .int32 = 20 } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gt", .value = .{ .int32 = 15 } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gte", .value = .{ .int32 = 20 } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$lt", .value = .{ .int32 = 30 } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$lte", .value = .{ .int32 = 30 } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &.{ .{ .int32 = 10 }, .{ .int32 = 30 } } } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{ .{ .key = "$gt", .value = .{ .int32 = 5 } }, .{ .key = "$lt", .value = .{ .int32 = 25 } } } }} } },
        .{ .pairs = &.{.{ .key = "b", .value = .{ .string = "x" } }} },
        .{ .pairs = &.{.{ .key = "b", .value = .null }} },
        .{ .pairs = &.{ .{ .key = "a", .value = .{ .int32 = 10 } }, .{ .key = "b", .value = .{ .string = "x" } } } },
        .{ .pairs = &.{.{ .key = "tags", .value = .{ .string = "a" } }} },
        .{ .pairs = &.{.{ .key = "tags", .value = .{ .array = &.{ .{ .string = "a" }, .{ .string = "b" } } } }} },
        .{ .pairs = &.{.{ .key = "tags", .value = .{ .doc = &.{.{ .key = "$all", .value = .{ .array = &.{ .{ .string = "a" } } } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$regex", .value = .{ .string = "^1" } }} } }} },
        .{ .pairs = &.{.{ .key = "c", .value = .{ .doc = &.{.{ .key = "$exists", .value = .{ .bool = false } }} } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .null }} },
        .{ .pairs = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} },
        .{ .pairs = &.{.{ .key = "_id", .value = .{ .string = "s4" } }} },
        .{ .pairs = &.{.{ .key = "_id", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .string = "s4" } } } }} } }} },
        .{ .pairs = &.{.{ .key = "$and", .value = .{ .array = &.{
            .{ .doc = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gte", .value = .{ .int32 = 10 } }} } }} },
            .{ .doc = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$lte", .value = .{ .int32 = 30 } }} } }} },
        } } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 2 } } } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gt", .value = .{ .int32 = 100 } }} } }} },
        .{ .pairs = &.{.{ .key = "$or", .value = .{ .array = &.{
            .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 10 } }} },
            .{ .doc = &.{.{ .key = "b", .value = .{ .string = "y" } }} },
        } } }} },
        .{ .pairs = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gt", .value = .null }} } }} },
        .{ .pairs = &.{ .{ .key = "a", .value = .{ .int32 = 10 } }, .{ .key = "tags", .value = .{ .int32 = 10 } } } },
    };

    // First with a compound (a, b) index, then after dropping it — and the
    // _id fast path is exercised by the _id filters in both runs.
    try dispatch_create_index(&tdb, io, "eq", .{ .doc = &.{
        .{ .key = "key", .value = .{ .doc = &.{
            .{ .key = "a", .value = .{ .int32 = 1 } },
            .{ .key = "b", .value = .{ .int32 = 1 } },
        } } },
        .{ .key = "name", .value = .{ .string = "a_1_b_1" } },
    } });

    // Run every filter through the index (the (a, b) plan plus the _id fast
    // path), then drop the index and require identical results from the
    // scan. The corpus's int32/int64 _id pair exercises the fast-path guard
    // (numbers fall back to a scan in both runs).
    var indexed_results: std.ArrayListUnmanaged(std.ArrayListUnmanaged([]u8)) = .empty;
    defer {
        for (indexed_results.items) |*l| free_id_list(testing.allocator, l);
        indexed_results.deinit(testing.allocator);
    }
    for (filters) |f| {
        var list: std.ArrayListUnmanaged([]u8) = .empty;
        errdefer free_id_list(testing.allocator, &list);
        try dispatch_find_ids(&tdb, io, "eq", f.pairs, &list);
        try indexed_results.append(testing.allocator, list);
    }
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("dropIndexes", .{ .string = "eq" }, &.{
            .{ .key = "index", .value = .{ .string = "*" } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
    }
    var scanned: std.ArrayListUnmanaged([]u8) = .empty;
    defer free_id_list(testing.allocator, &scanned);
    for (filters, 0..) |f, fi| {
        clear_id_list(testing.allocator, &scanned);
        try dispatch_find_ids(&tdb, io, "eq", f.pairs, &scanned);
        const indexed = indexed_results.items[fi];
        try testing.expectEqual(scanned.items.len, indexed.items.len);
        for (scanned.items, indexed.items) |a, b| try testing.expectEqualSlices(u8, a, b);
    }

    // Same corpus with a sparse (a, b) index: the sparse/null bail keeps
    // {a: null} and null-range filters correct.
    try dispatch_create_index(&tdb, io, "eq", .{ .doc = &.{
        .{ .key = "key", .value = .{ .doc = &.{
            .{ .key = "a", .value = .{ .int32 = 1 } },
            .{ .key = "b", .value = .{ .int32 = 1 } },
        } } },
        .{ .key = "name", .value = .{ .string = "a_1_b_1_sparse" } },
        .{ .key = "sparse", .value = .{ .bool = true } },
    } });
    for (indexed_results.items) |*l| free_id_list(testing.allocator, l);
    indexed_results.clearRetainingCapacity();
    for (filters) |f| {
        var list: std.ArrayListUnmanaged([]u8) = .empty;
        errdefer free_id_list(testing.allocator, &list);
        try dispatch_find_ids(&tdb, io, "eq", f.pairs, &list);
        try indexed_results.append(testing.allocator, list);
    }
    {
        var ctx = tdb.ctx(io);
        var msg = try parse_fake_msg("dropIndexes", .{ .string = "eq" }, &.{
            .{ .key = "index", .value = .{ .string = "*" } },
        });
        defer msg.deinit();
        var reply = wire.Reply.init(testing.allocator);
        defer reply.deinit();
        try dispatch(&ctx, &msg, &reply);
    }
    for (filters, 0..) |f, fi| {
        clear_id_list(testing.allocator, &scanned);
        try dispatch_find_ids(&tdb, io, "eq", f.pairs, &scanned);
        const indexed = indexed_results.items[fi];
        try testing.expectEqual(scanned.items.len, indexed.items.len);
        for (scanned.items, indexed.items) |a, b| try testing.expectEqualSlices(u8, a, b);
    }
}