const std = @import("std");
const Init = std.process.Init;
const Writer = Io.Writer;
const Allocator = std.mem.Allocator;
const Io = std.Io;
const Thread = std.Thread;
const Timestamp = std.Io.Timestamp;

const Store = @import("Store.zig");
const jetkv = @import("root.zig");
const build_options = @import("build_options");

const count = 12_500;
const thread_count = 8;

const WorkFn = fn (Allocator, *Store) void;

pub fn main(init: Init) !void {
    const io = init.io;
    const gpa = init.gpa;

    {
        var kv = jetkv.memory(io, gpa);
        defer kv.deinit();
        try benchBackend(io, gpa, &kv.store, "Memory");
    }

    {
        var kv = try jetkv.file(io, .{
            .path = "/tmp/jetkv_bench_file.db",
            .truncate = true,
        });
        defer kv.deinit();
        try benchBackend(io, gpa, &kv.store, "File");
    }

    {
        var kv = try jetkv.sqlite(io, gpa, .{
            .path = "/tmp/jetkv_bench_sqlite.db",
            .truncate = true,
        });
        defer kv.deinit();
        try benchBackend(io, gpa, &kv.store, "SQLite");
    }

    if (!build_options.@"skip-valkey") {
        var kv: jetkv.Valkey(.{ .pool_size = 16 }) = try .init(io, gpa);
        defer kv.deinit();
        try benchBackend(io, gpa, &kv.store, "Valkey");
    }
}

fn benchBackend(io: Io, gpa: Allocator, store: *Store, name: []const u8) !void {
    try runBenchmark(io, gpa, store, name, "get", getWork);
    try runBenchmark(io, gpa, store, name, "put", putWork);
    try runBenchmark(io, gpa, store, name, "append", appendWork);
}

fn runBenchmark(
    io: Io,
    gpa: Allocator,
    store: *Store,
    name: []const u8,
    phase: []const u8,
    comptime workFn: WorkFn,
) !void {
    // Seed a value so read workloads have something to find.
    try store.put("foo", "bar");
    const start_time: Timestamp = .now(io, .awake);

    var threads: [thread_count]Thread = undefined;
    for (0..thread_count) |index| {
        threads[index] = try .spawn(
            .{ .allocator = gpa },
            workFn,
            .{ gpa, store },
        );
    }

    for (threads) |thread| thread.join();

    const duration = start_time.untilNow(io, .awake);
    var buffer: Writer.Allocating = .init(gpa);
    defer buffer.deinit();
    const average_time = @divTrunc(duration.toNanoseconds(), count * thread_count);
    const average_duration: Io.Duration = .fromNanoseconds(average_time);
    try duration.format(&buffer.writer);
    const total = try buffer.toOwnedSlice();
    defer gpa.free(total);
    try average_duration.format(&buffer.writer);
    const average = try buffer.toOwnedSlice();
    defer gpa.free(average);
    const nanoseconds: f32 = @floatFromInt(duration.toNanoseconds());
    const transactions = count * thread_count;
    const tps = (transactions / nanoseconds) * std.time.ns_per_s;

    std.debug.print(
        \\== {[name]s} / {[phase]s} ==
        \\threads: {[threads]d}
        \\total transactions: {[transactions]d}
        \\duration: {[duration]s}
        \\average_duration: {[ave_duration]s}
        \\transactions per second: {[tps]d}
        \\
        \\
    , .{
        .name = name,
        .phase = phase,
        .threads = thread_count,
        .transactions = transactions,
        .duration = total,
        .ave_duration = average,
        .tps = tps,
    });
}

fn getWork(gpa: Allocator, kv: *Store) void {
    var i: usize = 0;
    while (i < count) : (i += 1) {
        var stack_fallback = std.heap.stackFallback(4096, gpa);
        const allocator = stack_fallback.get();
        const value = kv.get(allocator, "foo") catch |err| {
            std.debug.print("{s}\n", .{@errorName(err)});
            continue;
        };
        std.debug.assert(std.mem.eql(u8, value.?, "bar"));
    }
}

fn putWork(gpa: Allocator, kv: *Store) void {
    _ = gpa;
    var i: usize = 0;
    while (i < count) : (i += 1) {
        kv.put("foo", "bar") catch |err| {
            std.debug.print("{s}\n", .{@errorName(err)});
            continue;
        };
    }
}

fn appendWork(gpa: Allocator, kv: *Store) void {
    _ = gpa;
    var i: usize = 0;
    while (i < count) : (i += 1) {
        kv.append("list", "bar") catch |err| {
            std.debug.print("{s}\n", .{@errorName(err)});
            continue;
        };
    }
}