const std = @import("std");
const dbmod = @import("db.zig");
const batchmod = @import("batch.zig");
const txnmod = @import("txn.zig");
const txntable = @import("txn_table.zig");
const querymod = @import("query.zig");

const Allocator = std.mem.Allocator;
const Db = dbmod.Db;
const Transaction = txnmod.Transaction;
const TableTxn = txntable.TableTxn;
const WriteBatch = batchmod.WriteBatch;

/// A thread-safe facade over `*Db`. `Db` itself keeps no internal locking -
/// a fact documented (and worked around ad hoc) at every call site that
/// already needs concurrent access from more than one OS thread:
/// server.zig takes its own `mutex` before touching `db` on every request,
/// and every session note about txn_table.zig/query.zig/collections.zig
/// repeats the same known gap for the same reason. `GuardedDb` is that
/// exact pattern promoted to one reusable, independently tested primitive
/// instead of being reimplemented by hand wherever it is next needed.
///
/// This does not make `Db` itself thread-safe: it makes "every access to
/// this particular `Db` goes through this particular lock" enforceable,
/// which is the same guarantee `Server` already relies on. Two independent
/// `GuardedDb` wrapping the same underlying `*Db` would NOT be safe - the
/// mutex only serializes callers that agree to go through it.
pub const GuardedDb = struct {
    db: *Db,
    mutex: std.Thread.Mutex = .{},

    pub fn init(db: *Db) GuardedDb {
        return .{ .db = db };
    }

    /// Exclusive handle returned by `acquire()`. Holds the lock until
    /// `release()` runs, so a caller that needs several `Db` calls to
    /// happen as one atomic unit (a read-modify-write, or several puts
    /// that must never be observed half-applied by another guarded
    /// reader) can acquire once and issue every call through `.db`
    /// without releasing in between - the same shape as
    /// `Server.dispatch` locking `self.mutex` once per request no matter
    /// how many `Db` calls that request needs.
    pub const Guard = struct {
        db: *Db,
        mutex: *std.Thread.Mutex,

        pub fn release(self: Guard) void {
            self.mutex.unlock();
        }
    };

    pub fn acquire(self: *GuardedDb) Guard {
        self.mutex.lock();
        return .{ .db = self.db, .mutex = &self.mutex };
    }

    pub fn put(self: *GuardedDb, key: []const u8, value: []const u8) !void {
        const g = self.acquire();
        defer g.release();
        try g.db.put(key, value);
    }

    pub fn delete(self: *GuardedDb, key: []const u8) !void {
        const g = self.acquire();
        defer g.release();
        try g.db.delete(key);
    }

    pub fn get(self: *GuardedDb, allocator: Allocator, key: []const u8) !?[]u8 {
        const g = self.acquire();
        defer g.release();
        return g.db.get(allocator, key);
    }

    pub fn apply(self: *GuardedDb, batch: *WriteBatch) !void {
        const g = self.acquire();
        defer g.release();
        try g.db.apply(batch);
    }

    pub fn flush(self: *GuardedDb) !void {
        const g = self.acquire();
        defer g.release();
        try g.db.flush();
    }

    pub fn compact(self: *GuardedDb) !void {
        const g = self.acquire();
        defer g.release();
        try g.db.compact();
    }

    /// Atomic read-modify-write counter: decodes the current value at
    /// `key` as a little-endian i64 (0 if absent or the wrong width),
    /// adds `delta`, writes the new value back - the whole read, decide,
    /// write sequence runs under one lock acquisition, so no other
    /// guarded caller can observe or clobber the value partway through.
    /// A bare `db.get` followed later by a separate `db.put` cannot make
    /// that guarantee, which is exactly the lost-update hazard this
    /// exists to close (see the test below that proves it with real
    /// concurrent OS threads, not a simulated interleaving).
    pub fn incrBy(self: *GuardedDb, allocator: Allocator, key: []const u8, delta: i64) !i64 {
        const g = self.acquire();
        defer g.release();
        var current: i64 = 0;
        if (try g.db.get(allocator, key)) |raw| {
            defer allocator.free(raw);
            if (raw.len == 8) current = std.mem.readInt(i64, raw[0..8], .little);
        }
        const next = current + delta;
        var buf: [8]u8 = undefined;
        std.mem.writeInt(i64, &buf, next, .little);
        try g.db.put(key, &buf);
        return next;
    }
};

fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
    const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
    defer allocator.free(dbPath);
    return Db.open(allocator, dbPath, .{});
}

test "put/get/delete round-trip through the guard like a plain Db call" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();
    var guarded = GuardedDb.init(&db);

    try guarded.put("k", "v1");
    {
        const got = try guarded.get(allocator, "k");
        defer if (got) |g| allocator.free(g);
        try std.testing.expectEqualStrings("v1", got.?);
    }

    try guarded.delete("k");
    const afterDelete = try guarded.get(allocator, "k");
    try std.testing.expectEqual(@as(?[]u8, null), afterDelete);
}

test "incrBy on a fresh key starts from zero and returns the running total" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();
    var guarded = GuardedDb.init(&db);

    try std.testing.expectEqual(@as(i64, 5), try guarded.incrBy(allocator, "ctr", 5));
    try std.testing.expectEqual(@as(i64, 3), try guarded.incrBy(allocator, "ctr", -2));
    try std.testing.expectEqual(@as(i64, 103), try guarded.incrBy(allocator, "ctr", 100));
}

const IncrWorker = struct {
    guarded: *GuardedDb,
    key: []const u8,
    iterations: usize,

    fn run(self: IncrWorker) void {
        const allocator = std.heap.page_allocator;
        var i: usize = 0;
        while (i < self.iterations) : (i += 1) {
            _ = self.guarded.incrBy(allocator, self.key, 1) catch |err| {
                std.debug.panic("incrBy failed under real thread contention: {s}", .{@errorName(err)});
            };
        }
    }
};

test "incrBy survives real concurrent OS threads with zero lost updates" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();
    var guarded = GuardedDb.init(&db);

    const threadCount = 6;
    const perThread = 300;
    var threads: [threadCount]std.Thread = undefined;
    for (&threads) |*t| {
        t.* = try std.Thread.spawn(.{}, IncrWorker.run, .{IncrWorker{
            .guarded = &guarded,
            .key = "shared-counter",
            .iterations = perThread,
        }});
    }
    for (threads) |t| t.join();

    const raw = try guarded.get(allocator, "shared-counter");
    defer allocator.free(raw.?);
    const total = std.mem.readInt(i64, raw.?[0..8], .little);
    // A bare unsynchronized `db.get` + `db.put` read-modify-write would
    // routinely lose updates here (two threads read the same stale value,
    // both write back the same next value, one increment vanishes). The
    // mutex inside GuardedDb makes every incrBy fully serialized, so the
    // total below must land exactly on threadCount*perThread every single
    // run - not "usually", not "close to" - or the guard has a real bug.
    try std.testing.expectEqual(@as(i64, threadCount * perThread), total);
}

const TransferWorker = struct {
    guarded: *GuardedDb,
    rounds: usize,

    fn run(self: TransferWorker) void {
        const allocator = std.heap.page_allocator;
        var i: usize = 0;
        while (i < self.rounds) : (i += 1) {
            const g = self.guarded.acquire();
            defer g.release();
            const aRaw = g.db.get(allocator, "balanceA") catch @panic("get a failed");
            defer allocator.free(aRaw.?);
            const bRaw = g.db.get(allocator, "balanceB") catch @panic("get b failed");
            defer allocator.free(bRaw.?);
            const a = std.mem.readInt(i64, aRaw.?[0..8], .little);
            const b = std.mem.readInt(i64, bRaw.?[0..8], .little);
            if (a < 1) continue;
            var aBuf: [8]u8 = undefined;
            var bBuf: [8]u8 = undefined;
            std.mem.writeInt(i64, &aBuf, a - 1, .little);
            std.mem.writeInt(i64, &bBuf, b + 1, .little);
            g.db.put("balanceA", &aBuf) catch @panic("put a failed");
            g.db.put("balanceB", &bBuf) catch @panic("put b failed");
        }
    }
};

const InvariantReader = struct {
    guarded: *GuardedDb,
    checks: usize,
    violations: *std.atomic.Value(usize),

    fn run(self: InvariantReader) void {
        const allocator = std.heap.page_allocator;
        var i: usize = 0;
        while (i < self.checks) : (i += 1) {
            const g = self.guarded.acquire();
            const aRaw = g.db.get(allocator, "balanceA") catch @panic("get a failed");
            const bRaw = g.db.get(allocator, "balanceB") catch @panic("get b failed");
            const a = std.mem.readInt(i64, aRaw.?[0..8], .little);
            const b = std.mem.readInt(i64, bRaw.?[0..8], .little);
            g.release();
            allocator.free(aRaw.?);
            allocator.free(bRaw.?);
            if (a + b != 1000) _ = self.violations.fetchAdd(1, .monotonic);
        }
    }
};

test "acquire lets a writer move value between two keys atomically under a concurrent reader" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();
    var guarded = GuardedDb.init(&db);

    var initA: [8]u8 = undefined;
    var initB: [8]u8 = undefined;
    std.mem.writeInt(i64, &initA, 1000, .little);
    std.mem.writeInt(i64, &initB, 0, .little);
    try guarded.put("balanceA", &initA);
    try guarded.put("balanceB", &initB);

    var violations = std.atomic.Value(usize).init(0);
    const writerThread = try std.Thread.spawn(.{}, TransferWorker.run, .{TransferWorker{
        .guarded = &guarded,
        .rounds = 1000,
    }});
    const readerThread = try std.Thread.spawn(.{}, InvariantReader.run, .{InvariantReader{
        .guarded = &guarded,
        .checks = 1000,
        .violations = &violations,
    }});
    writerThread.join();
    readerThread.join();

    // Every single reader observation is required to see the invariant
    // balanceA + balanceB == 1000 hold, because each transfer's two puts
    // happen inside one `acquire()`/`release()` pair the reader's own
    // acquire() must wait out entirely - a real race would show up here
    // as a nonzero violation count, not as a crash, which is why this
    // asserts the counter is exactly zero instead of merely "did not hang".
    try std.testing.expectEqual(@as(usize, 0), violations.load(.monotonic));

    const finalARaw = try guarded.get(allocator, "balanceA");
    defer allocator.free(finalARaw.?);
    const finalBRaw = try guarded.get(allocator, "balanceB");
    defer allocator.free(finalBRaw.?);
    const finalA = std.mem.readInt(i64, finalARaw.?[0..8], .little);
    const finalB = std.mem.readInt(i64, finalBRaw.?[0..8], .little);
    try std.testing.expectEqual(@as(i64, 1000), finalA + finalB);
    try std.testing.expectEqual(@as(i64, 0), finalA);
    try std.testing.expectEqual(@as(i64, 1000), finalB);
}

test "apply is usable concurrently: disjoint batches from different threads never clobber each other" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();
    var guarded = GuardedDb.init(&db);

    const ApplyWorker = struct {
        guarded: *GuardedDb,
        prefix: u8,
        count: usize,

        fn run(self: @This()) void {
            var batch = WriteBatch.init(std.heap.page_allocator);
            defer batch.deinit();
            var i: usize = 0;
            while (i < self.count) : (i += 1) {
                const key = [_]u8{ self.prefix, @intCast(i) };
                batch.put(&key, "v") catch @panic("batch put failed");
            }
            self.guarded.apply(&batch) catch |err| {
                std.debug.panic("apply failed under contention: {s}", .{@errorName(err)});
            };
        }
    };

    var threads: [4]std.Thread = undefined;
    for (&threads, 0..) |*t, idx| {
        t.* = try std.Thread.spawn(.{}, ApplyWorker.run, .{ApplyWorker{
            .guarded = &guarded,
            .prefix = @intCast('a' + idx),
            .count = 50,
        }});
    }
    for (threads) |t| t.join();

    inline for (.{ 'a', 'b', 'c', 'd' }) |prefix| {
        var iter = try guarded.db.scanPrefix(allocator, &[_]u8{prefix});
        defer iter.deinit();
        var n: usize = 0;
        while (try iter.next()) |_| n += 1;
        try std.testing.expectEqual(@as(usize, 50), n);
    }
}

const TableTransferWorker = struct {
    guarded: *GuardedDb,
    rounds: usize,

    /// Runs a full `Transaction` + `TableTxn` transfer lifecycle (begin,
    /// read both rows, stage both writes, commit) under a single
    /// `acquire()`/`release()` pair, closing the specific gap every prior
    /// session on this project flagged and deferred: "brak testu
    /// wspolbieznosci wielowatkowej (dwa watki komituja TableTxn na tym
    /// samym Db jednoczesnie)". `Db.get`/`Transaction.commit` touch shared
    /// mutable state (the memtable, the OCC read-set bookkeeping) with no
    /// locking of their own, so without the guard two real OS threads
    /// racing through `TableTxn.open`/`get`/`put`/`commit` at once would
    /// be exactly the kind of unguarded concurrent access every other
    /// module in this project takes its own mutex to avoid.
    fn run(self: TableTransferWorker) void {
        const allocator = std.heap.page_allocator;
        var i: usize = 0;
        while (i < self.rounds) : (i += 1) {
            const g = self.guarded.acquire();
            defer g.release();

            var txn = Transaction.begin(g.db);
            defer txn.deinit();
            var t = TableTxn.open(&txn, allocator, "accounts") catch @panic("TableTxn.open failed");
            defer t.deinit();

            var from = (t.get(.{ .int = 1 }) catch @panic("get 1 failed")).?;
            defer from.deinit();
            var to = (t.get(.{ .int = 2 }) catch @panic("get 2 failed")).?;
            defer to.deinit();
            if (from.get("balance").?.int < 1) continue;

            t.put(.{ .int = 1 }, &.{.{ .name = "balance", .value = .{ .int = from.get("balance").?.int - 1 } }}) catch @panic("put 1 failed");
            t.put(.{ .int = 2 }, &.{.{ .name = "balance", .value = .{ .int = to.get("balance").?.int + 1 } }}) catch @panic("put 2 failed");

            // The whole sequence runs under one lock, so no external writer
            // can ever appear mid-transfer here: commit() must never see a
            // real error.Conflict in this test, unlike the deliberately
            // racy scenario in txn_table.zig's own "all-or-nothing" test.
            txn.commit() catch |err| std.debug.panic("unexpected commit failure under serialized access: {s}", .{@errorName(err)});
        }
    }
};

test "TableTxn commits from many real OS threads never lose or duplicate a unit of balance" {
    const allocator = std.testing.allocator;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var db = try openTemp(&tmp, allocator);
    defer db.close();

    const startBalance: i64 = 2000;
    {
        var setup = try querymod.Table.open(&db, allocator, "accounts");
        defer setup.deinit();
        try setup.put(.{ .int = 1 }, &.{.{ .name = "balance", .value = .{ .int = startBalance } }});
        try setup.put(.{ .int = 2 }, &.{.{ .name = "balance", .value = .{ .int = 0 } }});
    }

    var guarded = GuardedDb.init(&db);
    const threadCount = 4;
    const roundsPerThread = 200;
    var threads: [threadCount]std.Thread = undefined;
    for (&threads) |*t| {
        t.* = try std.Thread.spawn(.{}, TableTransferWorker.run, .{TableTransferWorker{
            .guarded = &guarded,
            .rounds = roundsPerThread,
        }});
    }
    for (threads) |t| t.join();

    var verify = try querymod.Table.open(&db, allocator, "accounts");
    defer verify.deinit();
    var row1 = (try verify.get(.{ .int = 1 })).?;
    defer row1.deinit();
    var row2 = (try verify.get(.{ .int = 2 })).?;
    defer row2.deinit();
    const balance1 = row1.get("balance").?.int;
    const balance2 = row2.get("balance").?.int;

    // Every transfer moves exactly 1 unit and every transfer's two writes
    // commit together or not at all, so the total must be exactly
    // preserved and account 2 must have gained exactly what account 1
    // lost - no interleaving of `get`+`put`+`commit` across threads can
    // produce any other pair of numbers.
    try std.testing.expectEqual(startBalance, balance1 + balance2);
    try std.testing.expectEqual(@as(i64, threadCount * roundsPerThread), balance2);
    try std.testing.expectEqual(startBalance - threadCount * roundsPerThread, balance1);
}