const std = @import("std");
const pardes = @import("pardes");
const filesystem = pardes.filesystem;
const ctlfs = pardes.ctlfs;

pub const std_options: std.Options = .{ .log_level = .err };

const backing = std.heap.page_allocator;
const harness_id = blk: {
    @setEvalBranchQuota(100_000);
    break :blk std.fmt.comptimePrint("{x}", .{std.hash.Wyhash.hash(0, @embedFile("fs_bench.zig"))});
};

fn nowNs() u64 {
    var ts: std.c.timespec = undefined;
    _ = std.c.clock_gettime(.MONOTONIC, &ts);
    return @as(u64, @intCast(ts.sec)) *| 1_000_000_000 +| @as(u64, @intCast(ts.nsec));
}

const Row = struct {
    name: []const u8,
    ns: u64,
    reps: usize,
    allocs: usize,
    note: []const u8 = "",
};

var rows: std.ArrayList(Row) = .empty;

fn record(name: []const u8, total_ns: u64, reps: usize, counting: *std.testing.FailingAllocator, note: []const u8) !void {
    try rows.append(backing, .{
        .name = name,
        .ns = total_ns / reps,
        .reps = reps,
        .allocs = counting.allocations,
        .note = note,
    });
}

const Session = struct {
    core: *pardes.Pardes,
    counting: *std.testing.FailingAllocator,
    file_id: usize,
    file_serial: u32,
    body_bytes: usize,
    tag: []u8,

    fn init(counting: *std.testing.FailingAllocator, body_bytes: usize) !Session {
        const gpa = counting.allocator();
        const core = try pardes.Pardes.init(gpa, .{ .tty_only = true, .cols = 120, .rows = 40 });
        errdefer core.deinit();
        while (core.nextEffect()) |_| {}

        const line = "the quick brown fox jumps over the lazy dog\n";
        var content: std.ArrayList(u8) = .empty;
        defer content.deinit(backing);
        while (content.items.len < body_bytes)
            try content.appendSlice(backing, line[0..@min(line.len, body_bytes - content.items.len)]);
        const pane = try core.setTestFile(content.items);
        const id = core.active;
        while (core.nextEffect()) |_| {}
        return .{
            .core = core,
            .counting = counting,
            .file_id = id,
            .file_serial = pane.serial,
            .body_bytes = pane.file.?.content.len,
            .tag = try core.tagText(backing, pane),
        };
    }

    fn deinit(s: *Session) void {
        backing.free(s.tag);
        s.core.deinit();
    }

    fn node(s: *const Session, file: ctlfs.PaneFile) u64 {
        return ctlfs.Node.of(s.file_serial, file);
    }
};

fn checkNumbers(bytes: []const u8, expected: []const usize) ![]const u8 {
    var rest = bytes;
    for (expected) |value| {
        rest = std.mem.trimStart(u8, rest, " ");
        const end = std.mem.indexOfScalar(u8, rest, ' ') orelse return error.IncorrectReply;
        const got = std.fmt.parseInt(usize, rest[0..end], 10) catch return error.IncorrectReply;
        if (got != value) return error.IncorrectReply;
        rest = rest[end + 1 ..];
    }
    return std.mem.trimStart(u8, rest, " ");
}

fn checkReply(s: *Session, req: ctlfs.Req, reply: ctlfs.Reply) !void {
    const event_read = req.op == .read and req.node == s.node(.event);
    const expected_status: ctlfs.Status = if (event_read) .again else .ok;
    if (reply.tag != req.tag or reply.errno != 0 or reply.status != expected_status)
        return error.IncorrectReply;
    const body = s.core.panes[s.file_id].?.file.?.content;
    switch (req.op) {
        .getattr => if (reply.attr.node != req.node or reply.attr.size != body.len or reply.attr.dir) return error.IncorrectReply,
        .lookup => if (reply.attr.node != s.node(.ctl) or reply.attr.dir) return error.IncorrectReply,
        .read => {
            if (event_read) {
                if (reply.payload != .none) return error.IncorrectReply;
            } else if (req.node == s.node(.body)) {
                if (reply.payload != .region) return error.IncorrectReply;
                const region = reply.payload.region;
                if (region.pane != s.file_id or region.serial != s.file_serial or region.off != req.off or region.len != req.size)
                    return error.IncorrectReply;
            } else {
                const bytes = s.core.fs.out.items;
                if (reply.payload != .staged or reply.payload.staged != bytes.len or bytes.len > req.size)
                    return error.IncorrectReply;
                const pane = s.core.panes[s.file_id].?;
                if (req.node == @intFromEnum(ctlfs.TopFile.index)) {
                    var expected_buf: [4200]u8 = undefined;
                    const expected = try std.fmt.bufPrint(&expected_buf, "{d} text {d} {s}\n", .{ s.file_serial, @intFromBool(pane.file.?.revision != pane.file.?.saved_revision), pane.file.?.path });
                    if (!std.mem.eql(u8, expected, bytes)) return error.IncorrectReply;
                    return;
                }
                const rest = try checkNumbers(bytes, &.{ s.file_serial, s.tag.len, body.len, 0, @intFromBool(pane.file.?.revision != pane.file.?.saved_revision) });
                if (req.node == s.node(.ctl)) {
                    const tail = try checkNumbers(rest, &.{pane.cols});
                    if (!std.mem.startsWith(u8, tail, "default ")) return error.IncorrectReply;
                    if ((try checkNumbers(tail[8..], &.{pardes.config.tab_width})).len != 0)
                        return error.IncorrectReply;
                } else return error.UnmeasuredOperation;
            }
        },
        .readdir => {
            const bytes = s.core.fs.out.items;
            if (reply.payload != .staged or reply.payload.staged != bytes.len or bytes.len > req.size)
                return error.IncorrectReply;
            const names = [_][]const u8{ "README", "index", "status", "look", "exec", "log", "screen", "listeners", "pane", "os" };
            const nodes = [_]u64{ @intFromEnum(ctlfs.TopFile.README), @intFromEnum(ctlfs.TopFile.index), @intFromEnum(ctlfs.TopFile.status), @intFromEnum(ctlfs.TopFile.look), @intFromEnum(ctlfs.TopFile.exec), @intFromEnum(ctlfs.TopFile.log), @intFromEnum(ctlfs.TopFile.screen), @intFromEnum(ctlfs.TopFile.listeners), @intFromEnum(ctlfs.TopFile.pane), filesystem.os_root };
            var off: usize = 0;
            var entry: usize = 0;
            while (off < bytes.len) : (entry += 1) {
                if (bytes.len - off < 10) return error.IncorrectReply;
                const node = std.mem.readInt(u64, bytes[off..][0..8], .little);
                const dir = bytes[off + 8];
                const size: usize = bytes[off + 9];
                off += 10;
                if (dir > 1 or size == 0 or size > bytes.len - off) return error.IncorrectReply;
                if (entry < names.len and (node != nodes[entry] or dir != @intFromBool(entry >= 7) or
                    !std.mem.eql(u8, names[entry], bytes[off..][0..size]))) return error.IncorrectReply;
                off += size;
            }
            if (entry < names.len) return error.IncorrectReply;
        },
        .write => if (reply.written != req.data.len or !std.mem.endsWith(u8, body, req.data)) return error.IncorrectReply,
        else => return error.UnmeasuredOperation,
    }
}

fn bench(s: *Session, name: []const u8, reps: usize, req: ctlfs.Req, note: []const u8) !void {
    std.debug.assert(reps > 0);
    const initial_bytes = s.core.panes[s.file_id].?.file.?.content.len;
    const initial_hash = if (req.op == .write) std.hash.Wyhash.hash(0, s.core.panes[s.file_id].?.file.?.content) else 0;
    for (0..64) |_| {
        try checkReply(s, req, ctlfs.handle(s.core, req));
        _ = s.core.scratch.reset(.retain_capacity);
    }
    s.counting.allocations = 0;
    const expected_status: ctlfs.Status = if (req.op == .read and req.node == s.node(.event)) .again else .ok;
    var valid = true;
    var reply: ctlfs.Reply = undefined;
    const start = nowNs();
    for (0..reps) |_| {
        reply = ctlfs.handle(s.core, req);
        valid = valid and reply.tag == req.tag and reply.status == expected_status and reply.errno == 0;
        std.mem.doNotOptimizeAway(reply);
        _ = s.core.scratch.reset(.retain_capacity);
    }
    const elapsed = nowNs() - start;
    if (!valid) return error.IncorrectReply;
    try checkReply(s, req, reply);
    if (req.op == .write) {
        const body = s.core.panes[s.file_id].?.file.?.content;
        if (body.len != initial_bytes + (64 + reps) * req.data.len or
            std.hash.Wyhash.hash(0, body[0..initial_bytes]) != initial_hash or
            !std.mem.allEqual(u8, body[initial_bytes..], 'x')) return error.MissingWrite;
    }
    try record(name, elapsed, reps, s.counting, note);
}

const Options = struct {
    reps: usize = 100_000,
    json: bool = false,

    fn parse(args: []const []const u8) !Options {
        var result: Options = .{};
        var seen_reps = false;
        var i: usize = 0;
        while (i < args.len) : (i += 1) {
            if (std.mem.eql(u8, args[i], "--json") and !result.json) {
                result.json = true;
            } else if (std.mem.eql(u8, args[i], "--reps") and !seen_reps) {
                i += 1;
                if (i == args.len) return error.MissingRepetitions;
                result.reps = std.fmt.parseInt(usize, args[i], 10) catch return error.InvalidRepetitions;
                if (result.reps == 0) return error.InvalidRepetitions;
                seen_reps = true;
            } else return error.InvalidArgument;
        }
        return result;
    }
};

pub fn main(init: std.process.Init) !void {
    const args = try init.minimal.args.toSlice(init.arena.allocator());
    const options = try Options.parse(args[1..]);
    const reps = options.reps;
    defer rows.deinit(backing);

    var counting: std.testing.FailingAllocator = .init(backing, .{});
    var s = try Session.init(&counting, 1 << 20);
    defer s.deinit();

    try bench(&s, "getattr body", reps, .{
        .tag = 1,
        .op = .getattr,
        .node = s.node(.body),
    }, "stat of a 1 MiB body");
    try bench(&s, "lookup ctl", reps, .{
        .tag = 2,
        .op = .lookup,
        .node = ctlfs.Node.of(s.file_serial, .dir),
        .data = "ctl",
    }, "name -> node");
    try bench(&s, "read body 4K", reps, .{
        .tag = 3,
        .op = .read,
        .node = s.node(.body),
        .off = 4096,
        .size = 4096,
    }, "must be zero-copy");
    try bench(&s, "read body 1M", @max(1, reps / 10), .{
        .tag = 4,
        .op = .read,
        .node = s.node(.body),
        .off = 0,
        .size = 1 << 20,
    }, "same cost as 4K if truly zero-copy");
    try bench(&s, "read ctl", reps, .{
        .tag = 5,
        .op = .read,
        .node = s.node(.ctl),
        .size = 256,
    }, "formatted into the staging buffer");
    try bench(&s, "read index", reps, .{
        .tag = 6,
        .op = .read,
        .node = @intFromEnum(ctlfs.TopFile.index),
        .size = 4096,
    }, "one line per pane");
    try bench(&s, "readdir root", reps, .{
        .tag = 7,
        .op = .readdir,
        .node = ctlfs.root,
        .size = 4096,
    }, "staged dirents");
    try bench(&s, "read event (empty)", reps, .{
        .tag = 8,
        .op = .read,
        .node = s.node(.event),
        .size = 256,
    }, "Status.again — the blocking primitive");
    try bench(&s, "write body 1K", @min(reps, 200), .{
        .tag = 9,
        .op = .write,
        .node = s.node(.body),
        .data = "x" ** 1024,
        .size = 1024,
    }, "append: whole-body swap + undo snapshot");

    const key_reps = @min(reps, 2000);
    var key_body_bytes: usize = 0;
    for ([_]bool{ false, true }) |scripted| {
        var keys = try Session.init(&counting, 32 * 1024);
        defer keys.deinit();
        const pane = keys.core.panes[keys.file_id].?;
        pane.mode = .insert;
        const initial_bytes = pane.file.?.content.len;
        key_body_bytes = initial_bytes;
        if (scripted) {
            keys.core.fs.panes[keys.file_id].readers = 1;
            keys.core.fs.listeners = 1;
        }
        counting.allocations = 0;
        const start = nowNs();
        for (0..key_reps) |_| {
            keys.core.update(.{ .key = .{ .cp = 'x', .text = "x" } });
            while (keys.core.nextEffect()) |_| {}
        }
        const elapsed = nowNs() - start;
        if (pane.file.?.content.len != initial_bytes + key_reps or pane.cur_col != key_reps or
            !std.mem.allEqual(u8, pane.file.?.content[0..key_reps], 'x')) return error.MissingKeystroke;
        const events = &keys.core.fs.panes[keys.file_id].events;
        if (scripted) for (0..key_reps) |i| {
            var buf: [64]u8 = undefined;
            const expected = try std.fmt.bufPrint(&buf, "KI{d} {d} 0 1 x\n", .{ i, i + 1 });
            if (!std.mem.eql(u8, expected, events.peek() orelse return error.MissingEvent)) return error.IncorrectEvent;
            events.pop();
        };
        if (!events.empty()) return error.IncorrectEvent;
        try record(
            if (scripted) "keystroke, 1 listener" else "keystroke, no listener",
            elapsed,
            key_reps,
            &counting,
            if (scripted) "diff + record" else "one branch",
        );
    }

    var out: std.Io.Writer.Allocating = .init(backing);
    defer out.deinit();
    const w = &out.writer;
    if (options.json) {
        try w.print("{{\"harness\":\"{s}\",\"body_bytes\":{d},\"key_body_bytes\":{d},\"rows\":[", .{ harness_id, s.body_bytes, key_body_bytes });
        for (rows.items, 0..) |r, n| {
            if (n > 0) try w.writeAll(",");
            try w.print("{{\"name\":\"{s}\",\"ns\":{d},\"reps\":{d},\"allocs\":{d}}}", .{ r.name, r.ns, r.reps, r.allocs });
        }
        try w.writeAll("]}\n");
    } else {
        try w.print("{s:<26} {s:>9} {s:>8}  {s}\n", .{ "request", "ns/op", "allocs", "note" });
        try w.print("{s:<26} {s:>9} {s:>8}  {s}\n", .{ "-" ** 26, "-" ** 9, "-" ** 8, "-" ** 20 });
        for (rows.items) |r|
            try w.print("{s:<26} {d:>9} {d:>8}  {s}\n", .{ r.name, r.ns, r.allocs, r.note });
    }
    try std.Io.File.stdout().writeStreamingAll(init.io, out.written());
}

test "fs benchmark rejects empty workloads and ambiguous options" {
    try std.testing.expectEqual(@as(usize, 100_000), (try Options.parse(&.{})).reps);
    try std.testing.expectEqualDeep(Options{ .reps = 1, .json = true }, try Options.parse(&.{ "--json", "--reps", "1" }));
    try std.testing.expectError(error.MissingRepetitions, Options.parse(&.{"--reps"}));
    for ([_][]const u8{ "0", "-1", "no", "99999999999999999999999999" }) |value|
        try std.testing.expectError(error.InvalidRepetitions, Options.parse(&.{ "--reps", value }));
    for ([_][]const []const u8{ &.{"--typo"}, &.{ "--json", "--json" }, &.{ "--reps", "1", "--reps", "2" } }) |args|
        try std.testing.expectError(error.InvalidArgument, Options.parse(args));
}

test "fs benchmark rejects wrong or incomplete staged metadata" {
    try std.testing.expectEqualStrings("/test.txt\n", try checkNumbers("          2          9       1024          0          0 /test.txt\n", &.{ 2, 9, 1024, 0, 0 }));
    try std.testing.expectError(error.IncorrectReply, checkNumbers("2 9 1023 0 0 /test.txt\n", &.{ 2, 9, 1024, 0, 0 }));
    try std.testing.expectError(error.IncorrectReply, checkNumbers("3 9 1024 0 0 /test.txt\n", &.{ 2, 9, 1024, 0, 0 }));
    try std.testing.expectError(error.IncorrectReply, checkNumbers("2 9 1024 ", &.{ 2, 9, 1024, 0, 0 }));
}