initial commit, moved libpanto over from the pantograph repo

1 files changed, 358 insertions, 0 deletions

diff --git a/src/image.zig b/src/image.zig
new file mode 100644
index 0000000..26474d1
--- /dev/null
+++ b/src/image.zig
@@ -0,0 +1,358 @@
+//! Native image processing for tool-returned attachments.
+//!
+//! Two responsibilities:
+//!
+//!   1. `detectCodec` — identify an attachment's codec from its leading
+//!      bytes (magic numbers), not its file extension.
+//!   2. `maybeResize` — bound raster images to `max_dim` on each side so a
+//!      single screenshot can't blow out the model's context. PDFs and
+//!      already-small images pass through untouched.
+//!
+//! Raster codecs go through the vendored stb single-header trio
+//! (decode -> Mitchell resize -> re-encode in the *same* codec). WEBP is
+//! decode-only (jebp), so a resized WEBP is re-encoded as JPEG.
+
+const std = @import("std");
+const Allocator = std.mem.Allocator;
+
+const c = @cImport({
+    @cInclude("stb_image.h");
+    @cInclude("stb_image_resize2.h");
+    @cInclude("stb_image_write.h");
+    @cInclude("jebp.h");
+});
+
+/// Longest side (px) allowed before we resize. pi uses 2000x2000.
+pub const max_dim: u32 = 2000;
+/// JPEG quality used when re-encoding (WEBP path, and JPEG inputs).
+const jpeg_quality: c_int = 80;
+
+pub const Codec = enum { png, jpeg, gif, bmp, webp, pdf };
+
+/// The MIME type string for a codec (static; do not free).
+pub fn mediaTypeForCodec(codec: Codec) []const u8 {
+    return switch (codec) {
+        .png => "image/png",
+        .jpeg => "image/jpeg",
+        .gif => "image/gif",
+        .bmp => "image/bmp",
+        .webp => "image/webp",
+        .pdf => "application/pdf",
+    };
+}
+
+/// Detect a supported codec from leading bytes (magic numbers).
+pub fn detectCodec(bytes: []const u8) ?Codec {
+    if (bytes.len >= 8 and std.mem.eql(u8, bytes[0..8], &.{ 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A }))
+        return .png;
+    if (bytes.len >= 3 and bytes[0] == 0xFF and bytes[1] == 0xD8 and bytes[2] == 0xFF)
+        return .jpeg;
+    if (bytes.len >= 6 and (std.mem.eql(u8, bytes[0..6], "GIF87a") or std.mem.eql(u8, bytes[0..6], "GIF89a")))
+        return .gif;
+    if (bytes.len >= 2 and bytes[0] == 'B' and bytes[1] == 'M')
+        return .bmp;
+    // WEBP: "RIFF"????"WEBP"
+    if (bytes.len >= 12 and std.mem.eql(u8, bytes[0..4], "RIFF") and std.mem.eql(u8, bytes[8..12], "WEBP"))
+        return .webp;
+    if (bytes.len >= 5 and std.mem.eql(u8, bytes[0..5], "%PDF-"))
+        return .pdf;
+    return null;
+}
+
+/// Map a media-type string back to a codec (for callers that already have
+/// the MIME string). Returns null for unsupported types.
+pub fn codecForMediaType(media_type: []const u8) ?Codec {
+    if (std.mem.eql(u8, media_type, "image/png")) return .png;
+    if (std.mem.eql(u8, media_type, "image/jpeg")) return .jpeg;
+    if (std.mem.eql(u8, media_type, "image/gif")) return .gif;
+    if (std.mem.eql(u8, media_type, "image/bmp")) return .bmp;
+    if (std.mem.eql(u8, media_type, "image/webp")) return .webp;
+    if (std.mem.eql(u8, media_type, "application/pdf")) return .pdf;
+    return null;
+}
+
+/// The result of `maybeResize`. `media_type` is the MIME type of `data`
+/// (may differ from the input when a WEBP was re-encoded as JPEG).
+/// `data` is always an owned slice the caller must free.
+pub const Processed = struct {
+    media_type: []const u8, // static string, do not free
+    data: []u8, // owned by `allocator`
+};
+
+/// Full attachment pipeline for a tool-returned media part: resolve the
+/// media type (detecting from magic bytes when `hint` is null), then
+/// resize. Returns owned raw bytes + the resolved media type.
+///
+/// Errors `error.UnknownMediaType` when neither the hint nor magic-byte
+/// detection recognizes the bytes — the caller decides how to surface
+/// that (e.g. drop the attachment, or fall back to text).
+pub fn process(allocator: Allocator, bytes: []const u8, hint: ?[]const u8) !Processed {
+    const codec = blk: {
+        if (hint) |h| {
+            if (codecForMediaType(h)) |hinted| break :blk hinted;
+        }
+        break :blk detectCodec(bytes) orelse return error.UnknownMediaType;
+    };
+    return maybeResize(allocator, bytes, codec);
+}
+
+/// Resize `bytes` so neither dimension exceeds `max_dim`, preserving the
+/// input codec where possible. Returns an owned copy of the (possibly
+/// unchanged) bytes plus the resulting media type.
+///
+///   - PDF: returned verbatim (a copy), media type unchanged.
+///   - raster <= max_dim on both sides: returned verbatim (a copy) — we
+///     skip the decode/encode round-trip to avoid quality loss + CPU.
+///   - stb-supported raster larger than max_dim: decode -> resize -> same
+///     codec.
+///   - WEBP larger than max_dim: jebp decode -> resize -> JPEG.
+pub fn maybeResize(allocator: Allocator, bytes: []const u8, codec: Codec) !Processed {
+    const media_type = mediaTypeForCodec(codec);
+
+    if (codec == .pdf)
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+
+    if (codec == .webp) return resizeWebp(allocator, bytes, media_type);
+    return resizeStb(allocator, bytes, media_type, codec);
+}
+
+/// Probe a raster image's dimensions without fully decoding. Returns null
+/// on failure.
+fn probeDims(bytes: []const u8) ?struct { w: u32, h: u32 } {
+    var w: c_int = 0;
+    var h: c_int = 0;
+    var comp: c_int = 0;
+    const ok = c.stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp);
+    if (ok == 0 or w <= 0 or h <= 0) return null;
+    return .{ .w = @intCast(w), .h = @intCast(h) };
+}
+
+/// Compute target dimensions that fit within `max_dim` x `max_dim` while
+/// preserving aspect ratio. Returns null when no resize is needed.
+fn targetDims(w: u32, h: u32) ?struct { w: u32, h: u32 } {
+    if (w <= max_dim and h <= max_dim) return null;
+    const wf: f64 = @floatFromInt(w);
+    const hf: f64 = @floatFromInt(h);
+    const scale = @min(@as(f64, @floatFromInt(max_dim)) / wf, @as(f64, @floatFromInt(max_dim)) / hf);
+    const nw: u32 = @max(1, @as(u32, @intFromFloat(@round(wf * scale))));
+    const nh: u32 = @max(1, @as(u32, @intFromFloat(@round(hf * scale))));
+    return .{ .w = nw, .h = nh };
+}
+
+const StbWriteCtx = struct {
+    list: *std.ArrayList(u8),
+    allocator: Allocator,
+    failed: bool = false,
+};
+
+fn stbWriteCb(ctx_opaque: ?*anyopaque, data: ?*anyopaque, size: c_int) callconv(.c) void {
+    const ctx: *StbWriteCtx = @ptrCast(@alignCast(ctx_opaque.?));
+    if (ctx.failed or size <= 0) return;
+    const bytes: [*]const u8 = @ptrCast(data.?);
+    ctx.list.appendSlice(ctx.allocator, bytes[0..@intCast(size)]) catch {
+        ctx.failed = true;
+    };
+}
+
+fn resizeStb(allocator: Allocator, bytes: []const u8, media_type: []const u8, codec: Codec) !Processed {
+    const dims = probeDims(bytes) orelse
+        // Can't parse it; pass through rather than fail the read.
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+
+    const target = targetDims(dims.w, dims.h) orelse
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+
+    // Decode to RGBA (4 channels) for a uniform resize path.
+    var w: c_int = 0;
+    var h: c_int = 0;
+    var comp: c_int = 0;
+    const pixels = c.stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, 4);
+    if (pixels == null)
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+    defer c.stbi_image_free(pixels);
+
+    const out_pixels = try allocator.alloc(u8, @as(usize, target.w) * @as(usize, target.h) * 4);
+    defer allocator.free(out_pixels);
+
+    const res = c.stbir_resize_uint8_srgb(
+        pixels,
+        w,
+        h,
+        0,
+        out_pixels.ptr,
+        @intCast(target.w),
+        @intCast(target.h),
+        0,
+        c.STBIR_RGBA,
+    );
+    if (res == null)
+        return error.ResizeFailed;
+
+    var out: std.ArrayList(u8) = .empty;
+    errdefer out.deinit(allocator);
+    var ctx = StbWriteCtx{ .list = &out, .allocator = allocator };
+
+    const tw: c_int = @intCast(target.w);
+    const th: c_int = @intCast(target.h);
+    const ok = switch (codec) {
+        .png => c.stbi_write_png_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, tw * 4),
+        .bmp => c.stbi_write_bmp_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr),
+        // stb has no GIF encoder; re-encode resized GIFs as PNG (lossless).
+        .gif => c.stbi_write_png_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, tw * 4),
+        .jpeg => c.stbi_write_jpg_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, jpeg_quality),
+        else => unreachable,
+    };
+    if (ok == 0 or ctx.failed) return error.EncodeFailed;
+
+    const result_media: []const u8 = switch (codec) {
+        .gif => "image/png", // re-encoded
+        else => media_type,
+    };
+    return .{ .media_type = result_media, .data = try out.toOwnedSlice(allocator) };
+}
+
+fn resizeWebp(allocator: Allocator, bytes: []const u8, media_type: []const u8) !Processed {
+    var img: c.jebp_image_t = std.mem.zeroes(c.jebp_image_t);
+    // Peek at the header first to learn dimensions cheaply.
+    if (c.jebp_decode_size(&img, bytes.len, bytes.ptr) != c.JEBP_OK)
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+
+    const w: u32 = @intCast(img.width);
+    const h: u32 = @intCast(img.height);
+    const target = targetDims(w, h) orelse
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+
+    // Full decode to RGBA (jebp_color_t is 4 contiguous bytes per pixel).
+    if (c.jebp_decode(&img, bytes.len, bytes.ptr) != c.JEBP_OK)
+        return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) };
+    defer c.jebp_free_image(&img);
+
+    const src: [*]const u8 = @ptrCast(img.pixels);
+    const out_pixels = try allocator.alloc(u8, @as(usize, target.w) * @as(usize, target.h) * 4);
+    defer allocator.free(out_pixels);
+
+    const res = c.stbir_resize_uint8_srgb(
+        src,
+        @intCast(w),
+        @intCast(h),
+        0,
+        out_pixels.ptr,
+        @intCast(target.w),
+        @intCast(target.h),
+        0,
+        c.STBIR_RGBA,
+    );
+    if (res == null) return error.ResizeFailed;
+
+    // TODO: when the source WEBP has an alpha layer, re-encoding to JPEG
+    // flattens transparency, which can look wrong for screenshots and
+    // diagrams. Consider re-encoding to PNG when alpha is present. For now
+    // we always emit JPEG: there is no small single-header WEBP encoder,
+    // and token size matters more than fidelity for LLM input.
+    var out: std.ArrayList(u8) = .empty;
+    errdefer out.deinit(allocator);
+    var ctx = StbWriteCtx{ .list = &out, .allocator = allocator };
+    const tw: c_int = @intCast(target.w);
+    const th: c_int = @intCast(target.h);
+    const ok = c.stbi_write_jpg_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, jpeg_quality);
+    if (ok == 0 or ctx.failed) return error.EncodeFailed;
+
+    return .{ .media_type = "image/jpeg", .data = try out.toOwnedSlice(allocator) };
+}
+
+// -----------------------------------------------------------------------------
+// Tests
+// -----------------------------------------------------------------------------
+
+const testing = std.testing;
+
+test "detectCodec + mediaTypeForCodec - magic bytes" {
+    try testing.expectEqualStrings("image/png", mediaTypeForCodec(detectCodec(&.{ 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A }).?));
+    try testing.expectEqualStrings("image/jpeg", mediaTypeForCodec(detectCodec(&.{ 0xFF, 0xD8, 0xFF, 0xE0 }).?));
+    try testing.expectEqualStrings("image/gif", mediaTypeForCodec(detectCodec("GIF89a....").?));
+    try testing.expectEqualStrings("image/bmp", mediaTypeForCodec(detectCodec("BM....").?));
+    try testing.expectEqualStrings("application/pdf", mediaTypeForCodec(detectCodec("%PDF-1.7").?));
+    const webp = "RIFF" ++ &[_]u8{ 0, 0, 0, 0 } ++ "WEBP";
+    try testing.expectEqualStrings("image/webp", mediaTypeForCodec(detectCodec(webp).?));
+    try testing.expect(detectCodec("not an image") == null);
+    try testing.expect(detectCodec(&.{0x89}) == null);
+}
+
+test "targetDims - skip when small, scale when large preserving aspect" {
+    try testing.expect(targetDims(100, 100) == null);
+    try testing.expect(targetDims(max_dim, max_dim) == null);
+    const t = targetDims(4000, 2000).?;
+    try testing.expectEqual(@as(u32, 2000), t.w);
+    try testing.expectEqual(@as(u32, 1000), t.h);
+    const t2 = targetDims(1000, 8000).?;
+    try testing.expectEqual(@as(u32, 250), t2.w);
+    try testing.expectEqual(@as(u32, 2000), t2.h);
+}
+
+test "process - detects type from raw bytes when hint absent" {
+    const a = testing.allocator;
+    // A tiny PNG (header + IHDR enough for stbi_info) — but simplest is to
+    // round-trip an stb-encoded small PNG and feed it with no hint.
+    const w: c_int = 4;
+    const h: c_int = 4;
+    var px: [4 * 4 * 4]u8 = undefined;
+    @memset(&px, 0x40);
+    var png: std.ArrayList(u8) = .empty;
+    defer png.deinit(a);
+    var ctx = StbWriteCtx{ .list = &png, .allocator = a };
+    try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &ctx, w, h, 4, &px, w * 4) != 0);
+
+    const out = try process(a, png.items, null);
+    defer a.free(out.data);
+    try testing.expectEqualStrings("image/png", out.media_type);
+
+    // Unknown bytes -> error so the caller can fall back / drop.
+    try testing.expectError(error.UnknownMediaType, process(a, "totally not an image", null));
+}
+
+test "maybeResize - PDF passes through unchanged" {
+    const a = testing.allocator;
+    const pdf = "%PDF-1.7\nfake pdf body";
+    const out = try maybeResize(a, pdf, .pdf);
+    defer a.free(out.data);
+    try testing.expectEqualStrings("application/pdf", out.media_type);
+    try testing.expectEqualStrings(pdf, out.data);
+}
+
+test "maybeResize - small PNG round-trips, large PNG shrinks and stays PNG" {
+    const a = testing.allocator;
+
+    // Build a small (8x8) RGBA PNG via stb and confirm pass-through.
+    const small_w: c_int = 8;
+    const small_h: c_int = 8;
+    var small_px: [8 * 8 * 4]u8 = undefined;
+    for (&small_px, 0..) |*b, i| b.* = @truncate(i);
+    var small_png: std.ArrayList(u8) = .empty;
+    defer small_png.deinit(a);
+    var sctx = StbWriteCtx{ .list = &small_png, .allocator = a };
+    try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &sctx, small_w, small_h, 4, &small_px, small_w * 4) != 0);
+
+    const small_out = try maybeResize(a, small_png.items, .png);
+    defer a.free(small_out.data);
+    try testing.expectEqualStrings("image/png", small_out.media_type);
+    // Small image is returned verbatim (byte-identical copy).
+    try testing.expectEqualSlices(u8, small_png.items, small_out.data);
+
+    // Build a large (2400x100) PNG and confirm it shrinks to <= max_dim.
+    const big_w: c_int = 2400;
+    const big_h: c_int = 100;
+    const big_px = try a.alloc(u8, @as(usize, @intCast(big_w * big_h * 4)));
+    defer a.free(big_px);
+    @memset(big_px, 0x7F);
+    var big_png: std.ArrayList(u8) = .empty;
+    defer big_png.deinit(a);
+    var bctx = StbWriteCtx{ .list = &big_png, .allocator = a };
+    try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &bctx, big_w, big_h, 4, big_px.ptr, big_w * 4) != 0);
+
+    const big_out = try maybeResize(a, big_png.items, .png);
+    defer a.free(big_out.data);
+    try testing.expectEqualStrings("image/png", big_out.media_type);
+    const dims = probeDims(big_out.data).?;
+    try testing.expectEqual(@as(u32, 2000), dims.w);
+    try testing.expect(dims.w <= max_dim and dims.h <= max_dim);
+}