//! Update operators: $set, $unset, $inc, $push, $pull, $rename with
//! dot-path navigation. Mutates the document's pairs in place, allocating
//! from the document's own arena.

const std = @import("std");
const bson = @import("bson.zig");
const query = @import("query.zig");

pub const UpdateError = error{ ImmutableId, InvalidUpdate, OutOfMemory };

const max_path_segments = 16;

/// Apply an update document (whose fields are operator documents) to `doc`.
pub fn apply(doc: *bson.Document, update: *const bson.Document) UpdateError!void {
    const arena = doc.arena.allocator();
    var pairs = try copy_pairs_to_list(arena, doc.pairs);
    for (update.pairs) |op| {
        if (op.key.len == 0 or op.key[0] != '$') return error.InvalidUpdate;
        try apply_operator(arena, &pairs, op.key, op.value);
    }
    doc.pairs = try pairs.toOwnedSlice(arena);
}

fn apply_operator(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), op: []const u8, value: bson.Value) UpdateError!void {
    if (std.mem.eql(u8, op, "$set")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            if (std.mem.eql(u8, p.key, "_id")) return error.ImmutableId;
            var segs: [max_path_segments][]const u8 = undefined;
            const n = split_path(p.key, &segs) orelse return error.InvalidUpdate;
            try set_path(arena, pairs, segs[0..n], try bson.copy_value(arena, p.value));
        }
        return;
    }
    if (std.mem.eql(u8, op, "$unset")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            var segs: [max_path_segments][]const u8 = undefined;
            const n = split_path(p.key, &segs) orelse continue;
            unset_path(arena, pairs, segs[0..n]);
        }
        return;
    }
    if (std.mem.eql(u8, op, "$inc")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            var segs: [max_path_segments][]const u8 = undefined;
            const n = split_path(p.key, &segs) orelse return error.InvalidUpdate;
            const current = get_value(pairs.items, segs[0..n]) orelse bson.Value{ .int32 = 0 };
            if (!current.is_number() or !p.value.is_number()) return error.InvalidUpdate;
            const sum = try numeric_add(arena, current, p.value);
            try set_path(arena, pairs, segs[0..n], sum);
        }
        return;
    }
    if (std.mem.eql(u8, op, "$push")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            var segs: [max_path_segments][]const u8 = undefined;
            const n = split_path(p.key, &segs) orelse return error.InvalidUpdate;
            const current_opt = get_value(pairs.items, segs[0..n]);
            var items: std.ArrayListUnmanaged(bson.Value) = .empty;
            defer items.deinit(arena);
            if (current_opt) |current| {
                switch (current) {
                    .array => |arr| try items.appendSlice(arena, arr),
                    .null => {},
                    else => return error.InvalidUpdate, // non-array field
                }
            }
            if (p.value == .doc) {
                if (bson.get_pair(p.value.doc, "$each")) |each| {
                    const arr = switch (each) {
                        .array => |a| a,
                        else => return error.InvalidUpdate,
                    };
                    for (arr) |item| try items.append(arena, try bson.copy_value(arena, item));
                    try set_path(arena, pairs, segs[0..n], .{ .array = try items.toOwnedSlice(arena) });
                    continue;
                }
            }
            try items.append(arena, try bson.copy_value(arena, p.value));
            try set_path(arena, pairs, segs[0..n], .{ .array = try items.toOwnedSlice(arena) });
        }
        return;
    }
    if (std.mem.eql(u8, op, "$pull")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            var segs: [max_path_segments][]const u8 = undefined;
            const n = split_path(p.key, &segs) orelse return error.InvalidUpdate;
            const current = get_value(pairs.items, segs[0..n]) orelse continue;
            const arr = switch (current) {
                .array => |a| a,
                else => return error.InvalidUpdate,
            };
            var items: std.ArrayListUnmanaged(bson.Value) = .empty;
            defer items.deinit(arena);
            for (arr) |elem| {
                if (!pull_matches(arena, p.value, elem)) {
                    try items.append(arena, elem);
                }
            }
            try set_path(arena, pairs, segs[0..n], .{ .array = try items.toOwnedSlice(arena) });
        }
        return;
    }
    if (std.mem.eql(u8, op, "$rename")) {
        const ops = doc_pairs(value) orelse return error.InvalidUpdate;
        for (ops) |p| {
            if (p.value != .string) return error.InvalidUpdate;
            if (std.mem.eql(u8, p.key, "_id") or std.mem.eql(u8, p.value.string, "_id")) return error.ImmutableId;
            var old_segs: [max_path_segments][]const u8 = undefined;
            const old_n = split_path(p.key, &old_segs) orelse return error.InvalidUpdate;
            const v = get_value(pairs.items, old_segs[0..old_n]) orelse continue; // no-op when absent
            unset_path(arena, pairs, old_segs[0..old_n]);
            var new_segs: [max_path_segments][]const u8 = undefined;
            const new_n = split_path(p.value.string, &new_segs) orelse return error.InvalidUpdate;
            try set_path(arena, pairs, new_segs[0..new_n], v);
        }
        return;
    }
    return error.InvalidUpdate;
}

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

fn split_path(path: []const u8, out: *[max_path_segments][]const u8) ?usize {
    var n: usize = 0;
    var it = std.mem.splitScalar(u8, path, '.');
    while (it.next()) |seg| {
        if (n >= max_path_segments) return null;
        out[n] = seg;
        n += 1;
    }
    return n;
}

fn parse_index(seg: []const u8) ?usize {
    return std.fmt.parseInt(usize, seg, 10) catch null;
}

fn copy_pairs_to_list(arena: std.mem.Allocator, pairs: []const bson.Pair) UpdateError!std.ArrayListUnmanaged(bson.Pair) {
    var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
    errdefer out.deinit(arena);
    try out.appendSlice(arena, pairs);
    return out;
}

fn find_pair(pairs: []const bson.Pair, key: []const u8) ?usize {
    for (pairs, 0..) |p, i| {
        if (std.mem.eql(u8, p.key, key)) return i;
    }
    return null;
}

fn get_value(pairs: []const bson.Pair, segs: []const []const u8) ?bson.Value {
    const idx = find_pair(pairs, segs[0]) orelse return null;
    if (segs.len == 1) return pairs[idx].value;
    return switch (pairs[idx].value) {
        .doc => |sub| get_value(sub, segs[1..]),
        .array => |arr| blk: {
            const index = parse_index(segs[1]) orelse break :blk null;
            if (index >= arr.len) break :blk null;
            if (segs.len == 2) break :blk arr[index];
            break :blk switch (arr[index]) {
                .doc => |sub| get_value(sub, segs[2..]),
                else => null,
            };
        },
        else => null,
    };
}

fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), segs: []const []const u8, value: bson.Value) UpdateError!void {
    if (segs.len == 1) {
        if (find_pair(pairs.items, segs[0])) |idx| {
            pairs.items[idx].value = value;
        } else {
            try pairs.append(arena, .{ .key = try arena.dupe(u8, segs[0]), .value = value });
        }
        return;
    }
    const idx = find_pair(pairs.items, segs[0]) orelse {
        const is_array = parse_index(segs[1]) != null;
        try pairs.append(arena, .{ .key = try arena.dupe(u8, segs[0]), .value = if (is_array) .{ .array = &.{} } else .{ .doc = &.{} } });
        return set_path(arena, pairs, segs, value);
    };
    switch (pairs.items[idx].value) {
        .doc => |sub| {
            var sub_pairs = try copy_pairs_to_list(arena, sub);
            defer sub_pairs.deinit(arena);
            try set_path(arena, &sub_pairs, segs[1..], value);
            pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
        },
        .array => |arr| {
            const index = parse_index(segs[1]) orelse {
                // treat as non-array: replace with a doc
                var sub_pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
                defer sub_pairs.deinit(arena);
                try set_path(arena, &sub_pairs, segs[1..], value);
                pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
                return;
            };
            var items = try copy_array_to_list(arena, arr);
            defer items.deinit(arena);
            if (index >= items.items.len) {
                try items.appendNTimes(arena, .null, index + 1 - items.items.len);
            }
            if (segs.len == 2) {
                items.items[index] = value;
            } else {
                switch (items.items[index]) {
                    .doc => |sub| {
                        var sub_pairs = try copy_pairs_to_list(arena, sub);
                        defer sub_pairs.deinit(arena);
                        try set_path(arena, &sub_pairs, segs[2..], value);
                        items.items[index] = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
                    },
                    else => {
                        var sub_pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
                        defer sub_pairs.deinit(arena);
                        try set_path(arena, &sub_pairs, segs[2..], value);
                        items.items[index] = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
                    },
                }
            }
            pairs.items[idx].value = .{ .array = try items.toOwnedSlice(arena) };
        },
        else => {
            var sub_pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
            defer sub_pairs.deinit(arena);
            try set_path(arena, &sub_pairs, segs[1..], value);
            pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
        },
    }
}

fn copy_array_to_list(arena: std.mem.Allocator, arr: []const bson.Value) UpdateError!std.ArrayListUnmanaged(bson.Value) {
    var out: std.ArrayListUnmanaged(bson.Value) = .empty;
    errdefer out.deinit(arena);
    try out.appendSlice(arena, arr);
    return out;
}

fn unset_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), segs: []const []const u8) void {
    if (segs.len == 1) {
        if (find_pair(pairs.items, segs[0])) |idx| {
            _ = pairs.orderedRemove(idx);
        }
        return;
    }
    const idx = find_pair(pairs.items, segs[0]) orelse return;
    switch (pairs.items[idx].value) {
        .doc => |sub| {
            var sub_pairs = copy_pairs_to_list(arena, sub) catch return;
            unset_path(arena, &sub_pairs, segs[1..]);
            pairs.items[idx].value = .{ .doc = sub_pairs.items };
        },
        else => {},
    }
}

fn pull_matches(arena: std.mem.Allocator, condition: bson.Value, elem: bson.Value) bool {
    switch (condition) {
        .doc => |cond_pairs| {
            const elem_doc = switch (elem) {
                .doc => |pairs| pairs,
                else => return false,
            };
            var all_operators = cond_pairs.len > 0;
            for (cond_pairs) |p| {
                if (p.key.len == 0 or p.key[0] != '$') {
                    all_operators = false;
                    break;
                }
            }
            if (all_operators) {
                // Operator condition against the element's value at each
                // operator's field — treat element doc as the doc.
                var ok = true;
                for (cond_pairs) |p| {
                    const actuals = bson.get_pair(elem_doc, p.key[1..]);
                    const a: bson.Value = actuals orelse .null;
                    if (!(query.value_matches_operator(arena, p.key, p.value, a) catch false)) ok = false;
                }
                return ok;
            }
            return (query.matches(arena, &.{ .arena = undefined, .pairs = cond_pairs }, &.{ .arena = undefined, .pairs = elem_doc }) catch false);
        },
        else => return bson.compare(condition, elem) == .eq,
    }
}

fn numeric_add(arena: std.mem.Allocator, a: bson.Value, b: bson.Value) UpdateError!bson.Value {
    _ = arena;
    if (a == .double or b == .double) {
        const sum: f64 = @floatCast(a.as_f128() + b.as_f128());
        return .{ .double = sum };
    }
    if (a == .int64 or b == .int64) {
        const av: i64 = switch (a) {
            .int32 => |i| i,
            .int64 => |i| i,
            else => unreachable,
        };
        const bv: i64 = switch (b) {
            .int32 => |i| i,
            .int64 => |i| i,
            else => unreachable,
        };
        const sum = std.math.add(i64, av, bv) catch return error.InvalidUpdate;
        return .{ .int64 = sum };
    }
    const sum: i64 = @as(i64, a.int32) + b.int32;
    if (sum >= std.math.minInt(i32) and sum <= std.math.maxInt(i32)) {
        return .{ .int32 = @intCast(sum) };
    }
    return .{ .int64 = sum };
}

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

const testing = std.testing;

fn doc_of(pairs: []const bson.Pair) bson.Document {
    return .{ .arena = undefined, .pairs = pairs };
}

test "$set, $inc, $unset, $rename" {
    var doc = bson.Document{ .arena = std.heap.ArenaAllocator.init(testing.allocator), .pairs = &.{} };
    defer doc.arena.deinit();
    doc.pairs = try doc.arena.allocator().dupe(bson.Pair, &.{
        .{ .key = "a", .value = .{ .int32 = 1 } },
        .{ .key = "user", .value = .{ .doc = &.{
            .{ .key = "name", .value = .{ .string = "bob" } },
            .{ .key = "age", .value = .{ .int32 = 30 } },
        } } },
        .{ .key = "gone", .value = .{ .int32 = 9 } },
    });

    try apply(&doc, &doc_of(&.{
        .{ .key = "$set", .value = .{ .doc = &.{
            .{ .key = "user.name", .value = .{ .string = "alice" } },
            .{ .key = "user.city", .value = .{ .string = "NYC" } },
            .{ .key = "new", .value = .{ .int32 = 5 } },
        } } },
        .{ .key = "$inc", .value = .{ .doc = &.{.{ .key = "user.age", .value = .{ .int32 = 2 } }} } },
        .{ .key = "$unset", .value = .{ .doc = &.{.{ .key = "gone", .value = .{ .string = "" } }} } },
        .{ .key = "$rename", .value = .{ .doc = &.{.{ .key = "new", .value = .{ .string = "renamed" } }} } },
    }));

    const user = bson.get_pair(doc.pairs, "user").?;
    try testing.expectEqualStrings("alice", user.doc[0].value.string);
    try testing.expectEqual(@as(i64, 32), bson.get_pair(user.doc, "age").?.int32);
    try testing.expectEqualStrings("NYC", bson.get_pair(user.doc, "city").?.string);
    try testing.expect(bson.get_pair(doc.pairs, "gone") == null);
    try testing.expect(bson.get_pair(doc.pairs, "new") == null);
    try testing.expectEqual(@as(i64, 5), bson.get_pair(doc.pairs, "renamed").?.int32);
}

test "$push and $pull" {
    var doc = bson.Document{ .arena = std.heap.ArenaAllocator.init(testing.allocator), .pairs = &.{} };
    defer doc.arena.deinit();
    doc.pairs = try doc.arena.allocator().dupe(bson.Pair, &.{
        .{ .key = "tags", .value = .{ .array = &.{ .{ .string = "a" }, .{ .string = "b" } } } },
    });

    try apply(&doc, &doc_of(&.{
        .{ .key = "$push", .value = .{ .doc = &.{.{ .key = "tags", .value = .{ .string = "c" } }} } },
    }));
    try testing.expectEqual(@as(usize, 3), bson.get_pair(doc.pairs, "tags").?.array.len);
    try testing.expectEqualStrings("c", bson.get_pair(doc.pairs, "tags").?.array[2].string);

    try apply(&doc, &doc_of(&.{
        .{ .key = "$pull", .value = .{ .doc = &.{.{ .key = "tags", .value = .{ .string = "b" } }} } },
    }));
    const tags = bson.get_pair(doc.pairs, "tags").?.array;
    try testing.expectEqual(@as(usize, 2), tags.len);
    try testing.expectEqualStrings("a", tags[0].string);
    try testing.expectEqualStrings("c", tags[1].string);

    try apply(&doc, &doc_of(&.{
        .{ .key = "$push", .value = .{ .doc = &.{.{ .key = "tags", .value = .{ .doc = &.{.{ .key = "$each", .value = .{ .array = &.{ .{ .string = "x" }, .{ .string = "y" } } } }} } }} } },
    }));
    try testing.expectEqual(@as(usize, 4), bson.get_pair(doc.pairs, "tags").?.array.len);
}

test "$set nested creation and _id protection" {
    var doc = bson.Document{ .arena = std.heap.ArenaAllocator.init(testing.allocator), .pairs = &.{} };
    defer doc.arena.deinit();
    doc.pairs = try doc.arena.allocator().dupe(bson.Pair, &.{
        .{ .key = "_id", .value = .{ .int32 = 1 } },
    });

    try apply(&doc, &doc_of(&.{
        .{ .key = "$set", .value = .{ .doc = &.{
            .{ .key = "a.b.c", .value = .{ .int32 = 42 } },
            .{ .key = "arr.1", .value = .{ .string = "x" } },
        } } },
    }));
    const a = bson.get_pair(doc.pairs, "a").?;
    const b = bson.get_pair(a.doc, "b").?;
    try testing.expectEqual(@as(i64, 42), bson.get_pair(b.doc, "c").?.int32);
    const arr = bson.get_pair(doc.pairs, "arr").?.array;
    try testing.expectEqual(@as(usize, 2), arr.len);
    try testing.expectEqualStrings("x", arr[1].string);

    try testing.expectError(error.ImmutableId, apply(&doc, &doc_of(&.{
        .{ .key = "$set", .value = .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 2 } }} } },
    })));
}

test "non-operator update rejected" {
    var doc = bson.Document{ .arena = std.heap.ArenaAllocator.init(testing.allocator), .pairs = &.{} };
    defer doc.arena.deinit();
    try testing.expectError(error.InvalidUpdate, apply(&doc, &doc_of(&.{
        .{ .key = "plain", .value = .{ .int32 = 1 } },
    })));
}