const std = @import("std");
const syntax = @import("syntax");
const gpa = std.heap.c_allocator;

const Case = struct {
    name: []const u8,
    path: []const u8,
    source: []const u8,
    locations: bool = false,
    start: usize = 0,
};

const cases = [_]Case{
    .{ .name = "zig", .path = "a.zig", .source = "const answer: u32 = 42;\npub fn main() void { // note\n    _ = \"hello\";\n}\n" },
    .{ .name = "markdown", .path = "a.md", .source = "# Heading *slant*\n\n**bold** and `code`\n\n```zig\nconst n = 42;\n```\n" },
    .{ .name = "markdown-nested", .path = "a.md", .source = "> # Heading *slant*\n>\n> Prose with **bold** and [a link](https://example.org).\n\n- `code` with <https://example.org>\n- plain text\n\n```zig\nfn main() void { return; }\n```\n" },
    .{ .name = "c-comment-window", .path = "a.c", .source = "/* opening\nstill a comment 42\n*/\nint n = 7;\n", .start = 11 },
    .{ .name = "markdown-indent", .path = "a.md", .source = "    # code, not a heading\n" },
    .{ .name = "result-markdown", .path = "+Grep", .source = "a.md:1 # Heading *slant*\na.md:3 **bold** and `code`\n", .locations = true },
    .{ .name = "result-indent", .path = "+Grep", .source = "a.md:1     # code, not a heading\n", .locations = true },
    .{ .name = "result-zig", .path = "+Grep", .source = "a.zig:1:1 const answer: u32 = 42;\na.zig:2:1 pub fn main() void { // note\na.zig:3:1     _ = \"hello\";\na.zig:4:1 }\n", .locations = true },
    .{ .name = "prose", .path = "+Help", .source = "A document without code locations.\nRead a.md for help.\n", .locations = true },
};

fn highlight(case: Case) ![]u8 {
    if (comptime @hasDecl(syntax, "highlightLocations")) {
        if (case.locations) return syntax.highlightLocations(gpa, case.source, case.start, case.source.len);
    }
    return syntax.highlightFileRange(gpa, case.path, case.source, case.start, case.source.len);
}

fn out(comptime format: []const u8, args: anytype) !void {
    var buffer: [4096]u8 = undefined;
    const text = try std.fmt.bufPrint(&buffer, format, args);
    var offset: usize = 0;
    while (offset < text.len) {
        const n = std.c.write(1, text.ptr + offset, text.len - offset);
        if (n <= 0) return error.WriteFailed;
        offset += @intCast(n);
    }
}

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

const Counter = struct {
    calls: usize = 0,
    bytes: usize = 0,
    live: usize = 0,
    peak: usize = 0,

    fn allocator(self: *Counter) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = alloc, .resize = resize, .remap = remap, .free = free } };
    }

    fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = gpa.rawAlloc(len, alignment, ret) orelse return null;
        self.calls += 1;
        self.bytes += len;
        self.live += len;
        self.peak = @max(self.peak, self.live);
        return ptr;
    }

    fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, len: usize, ret: usize) bool {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        if (!gpa.rawResize(memory, alignment, len, ret)) return false;
        self.live = self.live - memory.len + len;
        self.peak = @max(self.peak, self.live);
        return true;
    }

    fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, len: usize, ret: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = gpa.rawRemap(memory, alignment, len, ret) orelse return null;
        self.calls += 1;
        self.bytes += len;
        self.live = self.live - memory.len + len;
        self.peak = @max(self.peak, self.live);
        return ptr;
    }

    fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret: usize) void {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        self.live -= memory.len;
        gpa.rawFree(memory, alignment, ret);
    }
};

fn benchmark(case: Case, reps: usize) !void {
    var counter: Counter = .{};
    syntax.start(counter.allocator());
    const cold_start = now();
    gpa.free(try highlight(case));
    const cold = now() - cold_start;
    const retained = counter.live;
    const calls = counter.calls;
    const bytes = counter.bytes;
    const samples = try gpa.alloc(u64, reps);
    defer gpa.free(samples);
    for (samples) |*sample| {
        const before = now();
        gpa.free(try highlight(case));
        sample.* = now() - before;
    }
    std.mem.sort(u64, samples, {}, std.sort.asc(u64));
    const leaked = counter.live - retained;
    try out("{{\"case\":\"{s}\",\"bytes\":{d},\"cold_ns\":{d},\"median_ns\":{d},\"p95_ns\":{d},\"allocations\":{d},\"allocated_bytes\":{d},\"retained_growth\":{d}}}\n", .{
        case.name,                      case.source.len,                cold,   samples[reps / 2], samples[@min(reps - 1, reps * 95 / 100)],
        (counter.calls - calls) / reps, (counter.bytes - bytes) / reps, leaked,
    });
    syntax.stop();
    if (counter.live != 0) return error.LeakedSyntaxMemory;
}

pub fn main(init: std.process.Init) !void {
    const args = try init.minimal.args.toSlice(init.arena.allocator());
    var bench = false;
    var reps: usize = 21;
    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "--bench")) {
            bench = true;
        } else if (std.mem.eql(u8, args[i], "--reps") and i + 1 < args.len) {
            i += 1;
            reps = try std.fmt.parseInt(usize, args[i], 10);
            if (reps == 0) return error.InvalidRepetitions;
        } else return error.InvalidArgument;
    }
    if (!bench) {
        syntax.start(gpa);
        defer syntax.stop();
        for (cases) |case| {
            try out("{s}\n", .{case.name});
            const styles = try highlight(case);
            defer gpa.free(styles);
            for (case.source[case.start..], 0..) |byte, at| {
                if (byte == '\n') {
                    try out("\n", .{});
                } else {
                    const style = if (styles.len == 0) 0 else styles[at];
                    try out("{c}", .{".ksnc"[style]});
                }
            }
        }
        return;
    }
    for (cases) |case| try benchmark(case, reps);
    var results: std.ArrayList(u8) = .empty;
    defer results.deinit(gpa);
    for (0..512) |_| try results.appendSlice(gpa, "src/a.zig:12:1 const value: u32 = 42; // result\n");
    try benchmark(.{ .name = "results-512", .path = "+Grep", .source = results.items, .locations = true }, reps);
    results.clearRetainingCapacity();
    for (0..128) |_| try results.appendSlice(gpa, "a.zig:1 const value = 42;\nb.c:1 int value = 42;\nc.md:1 **bold** and `code`\nd.rs:1 let value = 42;\n");
    try benchmark(.{ .name = "results-mixed-512", .path = "+Grep", .source = results.items, .locations = true }, reps);
    results.clearRetainingCapacity();
    var row_buffer: [512]u8 = undefined;
    for (0..128) |line| {
        const rows = try std.fmt.bufPrint(&row_buffer, "a.zig:{d} const value_{d} = {d};\nb.c:{d} int value_{d} = {d};\nc.md:{d} **item {d}** and `code {d}`\nd.rs:{d} let value_{d} = {d};\n", .{ line + 1, line, line, line + 1, line, line, line + 1, line, line, line + 1, line, line });
        try results.appendSlice(gpa, rows);
    }
    try benchmark(.{ .name = "results-varied-512", .path = "+Grep", .source = results.items, .locations = true }, reps);
    results.clearRetainingCapacity();
    for (0..512) |line| {
        const row = try std.fmt.bufPrint(&row_buffer, "notes.md:{d} Ordinary prose for item {d}, without inline markup.\n", .{ line + 1, line });
        try results.appendSlice(gpa, row);
    }
    try benchmark(.{ .name = "results-prose-512", .path = "+Grep", .source = results.items, .locations = true }, reps);
    results.clearRetainingCapacity();
    for (0..128) |line| {
        const block = try std.fmt.bufPrint(&row_buffer, "# Heading {d}\n\nPlain paragraph {d}.\n\n**Bold {d}** and `code`.\n\n```zig\nconst value_{d} = {d};\n```\n\n", .{ line, line, line, line, line });
        try results.appendSlice(gpa, block);
    }
    try benchmark(.{ .name = "markdown-blocks-128", .path = "notes.md", .source = results.items }, reps);
}