//! Headless MuPDF raster scoreboard.
//!
//!   zig build pdf-bench
//!   zig build pdf-bench -- --json
//!   zig build pdf-bench -- --reps 25 --warmup 4
//!   zig build pdf-bench -- --path other.pdf --page 2 --pages 6
//!
//! The benchmark links the real src/pdf.zig wrapper and C exception bridge in
//! ReleaseFast. Timed regions include MuPDF rasterization directly into the
//! allocator-owned RGBA output; tint scenarios additionally include the exact
//! in-place PDF transform used by the core. Checksum/free/reporting stay
//! outside. "cold" means a fresh Document and MuPDF page cache (the OS
//! filesystem cache is deliberately uncontrolled), while warm scenarios
//! retain one Document.
const std = @import("std");
const bench_util = @import("bench_util.zig");
const nowNs = bench_util.nowNs;
const fatal = bench_util.Fatal("pdf-bench").f;
const pdf = @import("mupdf");

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

const gpa = std.heap.page_allocator;
const default_reps: usize = 25;
const default_warmup: usize = 3;
const default_multi_pages: usize = 8;
const checksum_seed: u64 = 0x7061_7264_6573_5044;

const Config = struct {
    path: []const u8 = "docs/design.pdf",
    page: usize = 0,
    reps: usize = default_reps,
    warmup: usize = default_warmup,
    multi_pages: usize = default_multi_pages,
    json: bool = false,
};

const Profile = struct {
    id: []const u8,
    description: []const u8,
    viewport_w: u32,
    viewport_h: u32,
    request: pdf.RenderRequest,
};

const profiles = [_]Profile{
    .{
        .id = "kitty-96",
        .description = "Kitty 96 DPI, bandwidth-capped raster",
        .viewport_w = 0,
        .viewport_h = 0,
        .request = .{ .dpi = 96, .max_dimension = 1200 },
    },
    .{
        .id = "sdl-192-1600x900",
        .description = "SDL 192 DPI, physical 1600x900 viewport minimum",
        .viewport_w = 1600,
        .viewport_h = 900,
        .request = .{
            .dpi = 192,
            .minimum_width = 1600,
            .minimum_height = 900,
            .max_dimension = 4096,
        },
    },
};

/// Kept orthogonal so a later benchmark can add another pixel treatment
/// without changing lifecycle/cold-cache accounting.
const Decoration = enum { plain, highlighted };
const Lifecycle = enum { cold_document_page, warm_page, warm_document_pass };

const Scenario = struct {
    id: []const u8,
    lifecycle: Lifecycle,
    decoration: Decoration,
    tint: pdf.TintMode = .disabled,
};

// Fixed benchmark colors, independent of user config and host defaults. They
// are Pardes's helix page background/foreground and make run-to-run tint
// checksums a stable optimization target.
const tint_colors: pdf.TintColors = .{
    .background = .{ 0x08, 0x08, 0x08 },
    .foreground = .{ 0xbd, 0xbd, 0xbd },
};

// Append only: baseline readers and the cross-scenario checks below rely on
// the original four scenarios retaining their indices and meaning.
const scenarios = [_]Scenario{
    .{ .id = "cold-document-plain", .lifecycle = .cold_document_page, .decoration = .plain },
    .{ .id = "warm-page-plain", .lifecycle = .warm_page, .decoration = .plain },
    .{ .id = "warm-page-highlighted", .lifecycle = .warm_page, .decoration = .highlighted },
    .{ .id = "warm-document-pass", .lifecycle = .warm_document_pass, .decoration = .plain },
    // The checked-in design page is grayscale. Reuse the original deterministic
    // highlight decoration here so filtered/full exercise source chroma and
    // must retain distinct correctness checksums.
    .{ .id = "warm-page-filtered", .lifecycle = .warm_page, .decoration = .highlighted, .tint = .filtered },
    .{ .id = "warm-page-full", .lifecycle = .warm_page, .decoration = .highlighted, .tint = .full },
};

const highlights = [_]pdf.Highlight{
    pdf.Highlight.init(.{
        .ul = .{ .x = 0.10, .y = 0.12 },
        .ur = .{ .x = 0.90, .y = 0.12 },
        .ll = .{ .x = 0.10, .y = 0.20 },
        .lr = .{ .x = 0.90, .y = 0.20 },
    }, .{ 0xff, 0xd5, 0x4f, 0x70 }, .search),
    pdf.Highlight.init(.{
        .ul = .{ .x = 0.18, .y = 0.46 },
        .ur = .{ .x = 0.82, .y = 0.46 },
        .ll = .{ .x = 0.18, .y = 0.54 },
        .lr = .{ .x = 0.82, .y = 0.54 },
    }, .{ 0x55, 0xa7, 0xff, 0x78 }, .selection),
};

const Identity = struct {
    pixels: u64 = 0,
    rgba_bytes: u64 = 0,
    checksum: u64 = checksum_seed,

    fn eql(a: Identity, b: Identity) bool {
        return a.pixels == b.pixels and a.rgba_bytes == b.rgba_bytes and
            a.checksum == b.checksum;
    }
};

const Timed = struct { ns: u64, identity: Identity };

const Result = struct {
    profile: usize,
    scenario: usize,
    identity: Identity,
    min_ns: u64,
    median_ns: u64,
    p90_ns: u64,
    max_ns: u64,
    sampled_ns: u64,
};

pub fn main(init: std.process.Init) !void {
    const args = try init.minimal.args.toSlice(init.arena.allocator());
    const config = parseArgs(args);

    var probe = pdf.Document.open(config.path) catch |err|
        fatal("cannot open {s}: {t}", .{ config.path, err });
    const document_pages = probe.pages;
    probe.deinit();
    if (config.page >= document_pages)
        fatal("--page {d} is outside the document's {d} pages", .{ config.page + 1, document_pages });
    const multi_pages = @min(config.multi_pages, document_pages);

    var results: [profiles.len * scenarios.len]Result = undefined;
    for (profiles, 0..) |profile, pi| {
        for (scenarios, 0..) |scenario, si| {
            results[pi * scenarios.len + si] = try measure(
                config,
                pi,
                profile,
                si,
                scenario,
                multi_pages,
            );
        }
        checkProfileResults(results[pi * scenarios.len ..][0..scenarios.len], profile);
    }

    if (config.json)
        reportJson(init.io, config, document_pages, multi_pages, &results)
    else
        reportText(init.io, config, document_pages, multi_pages, &results);
}

fn parseArgs(args: []const []const u8) Config {
    var config: Config = .{};
    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        const arg = args[i];
        if (std.mem.eql(u8, arg, "--json")) {
            config.json = true;
        } else if (std.mem.eql(u8, arg, "--reps")) {
            i += 1;
            if (i >= args.len) usage();
            config.reps = parsePositive("--reps", args[i], 1000);
        } else if (std.mem.eql(u8, arg, "--warmup")) {
            i += 1;
            if (i >= args.len) usage();
            config.warmup = std.fmt.parseUnsigned(usize, args[i], 10) catch
                fatal("--warmup expects an integer", .{});
            if (config.warmup > 100) fatal("--warmup must be at most 100", .{});
        } else if (std.mem.eql(u8, arg, "--path")) {
            i += 1;
            if (i >= args.len or args[i].len == 0) usage();
            config.path = args[i];
        } else if (std.mem.eql(u8, arg, "--page")) {
            i += 1;
            if (i >= args.len) usage();
            config.page = parsePositive("--page", args[i], std.math.maxInt(usize)) - 1;
        } else if (std.mem.eql(u8, arg, "--pages")) {
            i += 1;
            if (i >= args.len) usage();
            config.multi_pages = parsePositive("--pages", args[i], 10_000);
        } else if (std.mem.eql(u8, arg, "--help") or std.mem.eql(u8, arg, "-h")) {
            printUsage();
            std.process.exit(0);
        } else usage();
    }
    return config;
}

fn usage() noreturn {
    printUsage();
    std.process.exit(2);
}

fn printUsage() void {
    std.debug.print(
        "usage: pardes-pdf-bench [--json] [--reps N] [--warmup N] " ++
            "[--path FILE] [--page N] [--pages N]\n",
        .{},
    );
}

fn measure(
    config: Config,
    profile_index: usize,
    profile: Profile,
    scenario_index: usize,
    scenario: Scenario,
    multi_pages: usize,
) !Result {
    const samples = try gpa.alloc(u64, config.reps);
    defer gpa.free(samples);
    var sample_at: usize = 0;
    var expected: ?Identity = null;
    var sampled_ns: u64 = 0;
    const rounds = config.warmup + config.reps;

    if (scenario.lifecycle == .cold_document_page) {
        for (0..rounds) |round| {
            const timed = try timeCold(
                config.path,
                config.page,
                profile.request,
                scenario.decoration,
                scenario.tint,
            );
            try verifyIdentity(&expected, timed.identity, profile, scenario, round);
            if (round >= config.warmup) {
                samples[sample_at] = timed.ns;
                sample_at += 1;
                sampled_ns +|= timed.ns;
            }
        }
    } else {
        var document = try pdf.Document.open(config.path);
        defer document.deinit();
        for (0..rounds) |round| {
            const timed = switch (scenario.lifecycle) {
                .warm_page => try timePage(
                    &document,
                    config.page,
                    profile.request,
                    scenario.decoration,
                    scenario.tint,
                ),
                .warm_document_pass => try timeDocumentPass(
                    &document,
                    multi_pages,
                    profile.request,
                    scenario.decoration,
                    scenario.tint,
                ),
                .cold_document_page => unreachable,
            };
            try verifyIdentity(&expected, timed.identity, profile, scenario, round);
            if (round >= config.warmup) {
                samples[sample_at] = timed.ns;
                sample_at += 1;
                sampled_ns +|= timed.ns;
            }
        }
    }

    std.debug.assert(sample_at == samples.len);
    std.mem.sort(u64, samples, {}, std.sort.asc(u64));
    const p90_index = @min(samples.len - 1, (samples.len * 90 + 99) / 100 - 1);
    return .{
        .profile = profile_index,
        .scenario = scenario_index,
        .identity = expected.?,
        .min_ns = samples[0],
        .median_ns = samples[samples.len / 2],
        .p90_ns = samples[p90_index],
        .max_ns = samples[samples.len - 1],
        .sampled_ns = sampled_ns,
    };
}

fn timeCold(
    path: []const u8,
    page: usize,
    request: pdf.RenderRequest,
    decoration: Decoration,
    tint: pdf.TintMode,
) !Timed {
    const started = nowNs();
    var document = try pdf.Document.open(path);
    errdefer document.deinit();
    const rendered = try renderPage(&document, page, request, decoration, tint);
    const elapsed = @max(@as(u64, 1), nowNs() -| started);
    const identity = try consumeRender(rendered, page, .{});
    document.deinit();
    return .{ .ns = elapsed, .identity = identity };
}

fn timePage(
    document: *pdf.Document,
    page: usize,
    request: pdf.RenderRequest,
    decoration: Decoration,
    tint: pdf.TintMode,
) !Timed {
    const started = nowNs();
    const rendered = try renderPage(document, page, request, decoration, tint);
    const elapsed = @max(@as(u64, 1), nowNs() -| started);
    return .{
        .ns = elapsed,
        .identity = try consumeRender(rendered, page, .{}),
    };
}

fn timeDocumentPass(
    document: *pdf.Document,
    page_count: usize,
    request: pdf.RenderRequest,
    decoration: Decoration,
    tint: pdf.TintMode,
) !Timed {
    const rendered = try gpa.alloc(pdf.Render, page_count);
    defer gpa.free(rendered);
    var elapsed: u64 = 0;
    {
        var initialized: usize = 0;
        errdefer for (rendered[0..initialized]) |item| gpa.free(item.rgba);
        const started = nowNs();
        for (rendered, 0..) |*item, page| {
            item.* = try renderPage(document, page, request, decoration, tint);
            initialized += 1;
        }
        elapsed = @max(@as(u64, 1), nowNs() -| started);
    }
    var identity: Identity = .{};
    for (rendered, 0..) |item, page|
        identity = try consumeRender(item, page, identity);
    return .{ .ns = elapsed, .identity = identity };
}

fn renderPage(
    document: *pdf.Document,
    page: usize,
    request: pdf.RenderRequest,
    decoration: Decoration,
    tint: pdf.TintMode,
) !pdf.Render {
    const rendered = switch (decoration) {
        .plain => document.renderAt(gpa, page, request),
        .highlighted => document.renderWithHighlightsAt(gpa, page, request, &highlights),
    };
    const fresh = try rendered;
    errdefer gpa.free(fresh.rgba);
    // Keep the original scenarios instruction-for-instruction identical.
    // Only the appended tint scenarios enter the linear pixel pass.
    if (tint != .disabled) try pdf.tintRgba(fresh.rgba, tint, tint_colors);
    return fresh;
}

/// Hash every output byte after the timer stops. Reporting and cross-round
/// equality make the pixels observable, while keeping checksum cost out of the
/// renderer latency. Page/dimension mixing makes a multi-page pass ordered.
fn consumeRender(rendered: pdf.Render, page: usize, prior: Identity) !Identity {
    defer gpa.free(rendered.rgba);
    const pixels = std.math.mul(usize, rendered.width, rendered.height) catch
        return error.BadRasterDimensions;
    const expected_bytes = std.math.mul(usize, pixels, 4) catch
        return error.BadRasterDimensions;
    if (rendered.rgba.len != expected_bytes) return error.BadRasterLength;
    var checksum = std.hash.Wyhash.hash(
        prior.checksum ^ (@as(u64, page) *% 0x9e37_79b9_7f4a_7c15),
        rendered.rgba,
    );
    checksum ^= @as(u64, rendered.width) *% 0xbf58_476d_1ce4_e5b9;
    checksum ^= @as(u64, rendered.height) *% 0x94d0_49bb_1331_11eb;
    return .{
        .pixels = std.math.add(u64, prior.pixels, @intCast(pixels)) catch
            return error.CounterOverflow,
        .rgba_bytes = std.math.add(u64, prior.rgba_bytes, @intCast(expected_bytes)) catch
            return error.CounterOverflow,
        .checksum = checksum,
    };
}

fn verifyIdentity(
    expected: *?Identity,
    got: Identity,
    profile: Profile,
    scenario: Scenario,
    round: usize,
) !void {
    if (expected.*) |want| {
        if (!want.eql(got)) {
            std.debug.print(
                "pdf-bench: unstable pixels in {s}/{s} round {d}: " ++
                    "expected {x:0>16}, got {x:0>16}\n",
                .{ profile.id, scenario.id, round, want.checksum, got.checksum },
            );
            return error.UnstableRaster;
        }
    } else expected.* = got;
}

fn checkProfileResults(results: []const Result, profile: Profile) void {
    std.debug.assert(results.len == scenarios.len);
    const cold = results[0].identity;
    const warm = results[1].identity;
    const marked = results[2].identity;
    const filtered = results[4].identity;
    const full = results[5].identity;
    if (!cold.eql(warm))
        fatal("{s}: cold and warm plain renders produced different pixels", .{profile.id});
    if (warm.pixels != marked.pixels or warm.rgba_bytes != marked.rgba_bytes)
        fatal("{s}: highlighted render changed raster dimensions", .{profile.id});
    if (warm.checksum == marked.checksum)
        fatal("{s}: highlighted render did not change pixels", .{profile.id});
    if (marked.pixels != filtered.pixels or marked.rgba_bytes != filtered.rgba_bytes or
        marked.pixels != full.pixels or marked.rgba_bytes != full.rgba_bytes)
        fatal("{s}: tint changed raster dimensions", .{profile.id});
    if (marked.checksum == filtered.checksum or marked.checksum == full.checksum or
        filtered.checksum == full.checksum)
        fatal("{s}: tint did not change pixels", .{profile.id});
}

fn medianMpixPerSecond(result: Result) f64 {
    return @as(f64, @floatFromInt(result.identity.pixels)) * 1000.0 /
        @as(f64, @floatFromInt(result.median_ns));
}

fn medianGibPerSecond(result: Result) f64 {
    return @as(f64, @floatFromInt(result.identity.rgba_bytes)) * 1_000_000_000.0 /
        (@as(f64, @floatFromInt(result.median_ns)) * 1024.0 * 1024.0 * 1024.0);
}

fn ms(ns: u64) f64 {
    return @as(f64, @floatFromInt(ns)) / 1_000_000.0;
}

fn mib(bytes: u64) f64 {
    return @as(f64, @floatFromInt(bytes)) / (1024.0 * 1024.0);
}

const report_max_bytes = 64 * 1024;

fn reportText(
    io: std.Io,
    config: Config,
    document_pages: usize,
    multi_pages: usize,
    results: []const Result,
) void {
    var storage: [report_max_bytes]u8 = undefined;
    var out: std.Io.Writer = .fixed(&storage);
    out.print(
        "pardes PDF render benchmark (ReleaseFast)\n" ++
            "path: {s} ({d} pages), selected page: {d}, multi-page pass: {d} pages\n" ++
            "warmup: {d}, sampled reps: {d}; cold = fresh Document (OS cache uncontrolled)\n" ++
            "tint colors: background=#{x:0>2}{x:0>2}{x:0>2}, foreground=#{x:0>2}{x:0>2}{x:0>2}\n\n",
        .{
            config.path,
            document_pages,
            config.page + 1,
            multi_pages,
            config.warmup,
            config.reps,
            tint_colors.background[0],
            tint_colors.background[1],
            tint_colors.background[2],
            tint_colors.foreground[0],
            tint_colors.foreground[1],
            tint_colors.foreground[2],
        },
    ) catch return;
    out.writeAll(
        "profile              scenario                   min ms   median       p90    max ms  MPix/s  GiB/s  pixels/op  RGBA MiB  checksum\n" ++
            "-----------------------------------------------------------------------------------------------------------------------------------\n",
    ) catch return;
    for (results) |result| {
        const profile = profiles[result.profile];
        const scenario = scenarios[result.scenario];
        out.print(
            "{s:<20} {s:<26} {d:>8.2} {d:>8.2} {d:>8.2} {d:>8.2} {d:>7.1} {d:>6.2} {d:>10} {d:>9.2}  {x:0>16}\n",
            .{
                profile.id,
                scenario.id,
                ms(result.min_ns),
                ms(result.median_ns),
                ms(result.p90_ns),
                ms(result.max_ns),
                medianMpixPerSecond(result),
                medianGibPerSecond(result),
                result.identity.pixels,
                mib(result.identity.rgba_bytes),
                result.identity.checksum,
            },
        ) catch return;
    }
    out.writeAll("\nrequests:\n") catch return;
    for (profiles) |profile| {
        out.print(
            "  {s}: dpi={d}, viewport-min={d}x{d}, max-dimension={d} — {s}\n",
            .{
                profile.id,
                profile.request.dpi,
                profile.viewport_w,
                profile.viewport_h,
                profile.request.max_dimension,
                profile.description,
            },
        ) catch return;
    }
    std.Io.File.stdout().writeStreamingAll(io, out.buffered()) catch {};
}

fn reportJson(
    io: std.Io,
    config: Config,
    document_pages: usize,
    multi_pages: usize,
    results: []const Result,
) void {
    var storage: [report_max_bytes]u8 = undefined;
    var out: std.Io.Writer = .fixed(&storage);
    const path_json = std.json.Stringify.valueAlloc(gpa, config.path, .{}) catch return;
    defer gpa.free(path_json);
    out.print(
        "{{\"benchmark\":\"pardes-pdf-render\",\"build\":\"ReleaseFast\"," ++
            "\"path\":{s},\"document_pages\":{d},\"selected_page\":{d}," ++
            "\"multi_pages\":{d},\"warmup\":{d},\"reps\":{d}," ++
            "\"tint_colors\":{{\"background\":\"{x:0>2}{x:0>2}{x:0>2}\"," ++
            "\"foreground\":\"{x:0>2}{x:0>2}{x:0>2}\"}},\"profiles\":[",
        .{
            path_json,
            document_pages,
            config.page + 1,
            multi_pages,
            config.warmup,
            config.reps,
            tint_colors.background[0],
            tint_colors.background[1],
            tint_colors.background[2],
            tint_colors.foreground[0],
            tint_colors.foreground[1],
            tint_colors.foreground[2],
        },
    ) catch return;
    for (profiles, 0..) |profile, i| {
        out.print(
            "{s}{{\"id\":\"{s}\",\"dpi\":{d},\"minimum_width\":{d}," ++
                "\"minimum_height\":{d},\"max_dimension\":{d}}}",
            .{
                if (i == 0) "" else ",",
                profile.id,
                profile.request.dpi,
                profile.request.minimum_width,
                profile.request.minimum_height,
                profile.request.max_dimension,
            },
        ) catch return;
    }
    out.writeAll("],\"results\":[") catch return;
    for (results, 0..) |result, i| {
        out.print(
            "{s}{{\"profile\":\"{s}\",\"scenario\":\"{s}\"," ++
                "\"min_ns\":{d},\"median_ns\":{d},\"p90_ns\":{d}," ++
                "\"max_ns\":{d},\"sampled_ns\":{d},\"pixels\":{d}," ++
                "\"rgba_bytes\":{d},\"median_mpix_s\":{d:.3}," ++
                "\"median_gib_s\":{d:.3},\"checksum\":\"{x:0>16}\"}}",
            .{
                if (i == 0) "" else ",",
                profiles[result.profile].id,
                scenarios[result.scenario].id,
                result.min_ns,
                result.median_ns,
                result.p90_ns,
                result.max_ns,
                result.sampled_ns,
                result.identity.pixels,
                result.identity.rgba_bytes,
                medianMpixPerSecond(result),
                medianGibPerSecond(result),
                result.identity.checksum,
            },
        ) catch return;
    }
    out.writeAll("]}\n") catch return;
    std.Io.File.stdout().writeStreamingAll(io, out.buffered()) catch {};
}

fn parsePositive(name: []const u8, text: []const u8, maximum: usize) usize {
    return bench_util.parsePositive(fatal, name, text, maximum);
}