//! GET /api/image?url=&max_bytes= — fetch an image from an allowed host and
//! shrink it under a byte budget. ported from the retired rust backend.
//!
//! the consumer is the bufo bot: bluesky rejects uploads over ~976KB and a
//! handful of bufos are 2048px pngs well past that. same ladder as the rust:
//! pngs are downscaled and re-encoded as png first (75/50/25%), then
//! re-encoded as jpeg at falling quality, then downscaled + jpeg; non-pngs
//! skip straight to the jpeg ladders. if nothing fits, the original goes back.
//!
//! codecs are zigimg (pure zig); the resample is a box filter, which is all a
//! 2–4× downscale of a cartoon needs.

const std = @import("std");
const mem = std.mem;
const Allocator = mem.Allocator;
const zigimg = @import("zigimg");
const log = std.log.scoped(.image);

const allowed_domains = [_][]const u8{
    "all-the.bufo.zone",
    "find-bufo-zig.fly.dev",
    "find-bufo.com",
};

pub const DEFAULT_MAX_BYTES: usize = 900_000;
pub const CACHE_CONTROL = "public, max-age=86400";

pub fn isAllowedUrl(url: []const u8) bool {
    return hasAllowedDomain(url, &allowed_domains);
}

pub fn isAllowedSimilarityUrl(url: []const u8) bool {
    return isAllowedUrl(url);
}

fn hasAllowedDomain(url: []const u8, domains: []const []const u8) bool {
    const prefix = "https://";
    if (!mem.startsWith(u8, url, prefix)) return false;
    const rest = url[prefix.len..];
    const end = mem.indexOfAny(u8, rest, "/?#") orelse rest.len;
    const authority = rest[0..end];
    if (authority.len == 0 or mem.indexOfScalar(u8, authority, '@') != null) return false;
    for (domains) |domain| if (mem.eql(u8, authority, domain)) return true;
    return false;
}

pub const Kind = enum { png, gif, jpeg, webp, other };

/// sniff the container from magic bytes; falls back to the url's extension
pub fn sniff(bytes: []const u8, url: []const u8) Kind {
    if (bytes.len >= 8 and mem.eql(u8, bytes[0..8], "\x89PNG\r\n\x1a\n")) return .png;
    if (bytes.len >= 4 and mem.eql(u8, bytes[0..4], "GIF8")) return .gif;
    if (bytes.len >= 3 and mem.eql(u8, bytes[0..3], "\xff\xd8\xff")) return .jpeg;
    if (bytes.len >= 12 and mem.eql(u8, bytes[0..4], "RIFF") and mem.eql(u8, bytes[8..12], "WEBP")) return .webp;
    if (mem.endsWith(u8, url, ".png")) return .png;
    if (mem.endsWith(u8, url, ".gif")) return .gif;
    if (mem.endsWith(u8, url, ".jpg") or mem.endsWith(u8, url, ".jpeg")) return .jpeg;
    return .other;
}

pub fn contentType(kind: Kind) []const u8 {
    return switch (kind) {
        .png => "image/png",
        .gif => "image/gif",
        .jpeg => "image/jpeg",
        .webp => "image/webp",
        .other => "application/octet-stream",
    };
}

pub fn uploadContentType(bytes: []const u8) ?[]const u8 {
    const kind = sniff(bytes, "");
    return if (kind == .other) null else contentType(kind);
}

pub const Shrunk = struct {
    content_type: []const u8,
    bytes: []const u8,
};

const png_scales = [_]u32{ 75, 50, 25 };
const jpeg_qualities = [_]u8{ 85, 70, 50, 30 };

/// returns bytes that fit `max_bytes` when possible. every returned slice is
/// owned by `alloc` (the original is returned as-is when it already fits or
/// nothing else works).
pub fn shrink(alloc: Allocator, original: []const u8, kind: Kind, max_bytes: usize) Allocator.Error!Shrunk {
    const as_is: Shrunk = .{ .content_type = contentType(kind), .bytes = original };
    if (original.len <= max_bytes) return as_is;

    var img = zigimg.Image.fromMemory(alloc, original) catch |err| {
        log.err("failed to decode image: {t}", .{err});
        return as_is;
    };
    defer img.deinit(alloc);
    img.convert(alloc, .rgba32) catch |err| {
        log.err("failed to convert image to rgba32: {t}", .{err});
        return as_is;
    };
    log.info("image {d} bytes exceeds {d} limit, resizing ({d}x{d} {t})", .{ original.len, max_bytes, img.width, img.height, kind });

    // encoders write into a caller-owned buffer; nothing we produce should
    // exceed the original, but leave headroom for png overhead on small images
    const scratch = try alloc.alloc(u8, original.len + (1 << 20));

    if (kind == .png) {
        for (png_scales) |scale| {
            var small = downscale(alloc, img, scale) catch continue;
            defer small.deinit(alloc);
            const out = small.writeToMemory(alloc, scratch, .{ .png = .{} }) catch continue;
            if (out.len <= max_bytes) {
                log.info("resized png to {d}x{d} ({d}%), {d} bytes", .{ small.width, small.height, scale, out.len });
                return .{ .content_type = "image/png", .bytes = try alloc.dupe(u8, out) };
            }
        }
    }

    var rgb = compositeOnWhite(alloc, img) catch return as_is;
    defer rgb.deinit(alloc);
    for (jpeg_qualities) |q| {
        const out = rgb.writeToMemory(alloc, scratch, .{ .jpeg = .{ .quality = q } }) catch continue;
        if (out.len <= max_bytes) {
            log.info("re-encoded as jpeg q={d}, {d} bytes", .{ q, out.len });
            return .{ .content_type = "image/jpeg", .bytes = try alloc.dupe(u8, out) };
        }
    }
    for (png_scales) |scale| {
        var small = downscale(alloc, rgb, scale) catch continue;
        defer small.deinit(alloc);
        const out = small.writeToMemory(alloc, scratch, .{ .jpeg = .{ .quality = 50 } }) catch continue;
        if (out.len <= max_bytes) {
            log.info("resized to {d}x{d} ({d}%) + jpeg q=50, {d} bytes", .{ small.width, small.height, scale, out.len });
            return .{ .content_type = "image/jpeg", .bytes = try alloc.dupe(u8, out) };
        }
    }

    log.warn("could not shrink image under {d} bytes, returning original", .{max_bytes});
    return as_is;
}

/// box-filter downscale to `percent` of each dimension. works for rgba32 and
/// rgb24 sources and preserves the pixel format.
fn downscale(alloc: Allocator, src: zigimg.Image, percent: u32) !zigimg.Image {
    const dw: usize = @max(1, src.width * percent / 100);
    const dh: usize = @max(1, src.height * percent / 100);
    switch (src.pixels) {
        .rgba32 => |px| {
            var dst = try zigimg.Image.create(alloc, dw, dh, .rgba32);
            errdefer dst.deinit(alloc);
            boxFilter(zigimg.color.Rgba32, px, src.width, src.height, dst.pixels.rgba32, dw, dh, true);
            return dst;
        },
        .rgb24 => |px| {
            var dst = try zigimg.Image.create(alloc, dw, dh, .rgb24);
            errdefer dst.deinit(alloc);
            boxFilter(zigimg.color.Rgb24, px, src.width, src.height, dst.pixels.rgb24, dw, dh, false);
            return dst;
        },
        else => return error.UnsupportedPixelFormat,
    }
}

fn boxFilter(
    comptime Px: type,
    src: []const Px,
    sw: usize,
    sh: usize,
    dst: []Px,
    dw: usize,
    dh: usize,
    comptime has_alpha: bool,
) void {
    for (0..dh) |dy| {
        const y0 = dy * sh / dh;
        const y1 = @max(y0 + 1, (dy + 1) * sh / dh);
        for (0..dw) |dx| {
            const x0 = dx * sw / dw;
            const x1 = @max(x0 + 1, (dx + 1) * sw / dw);
            var r: u64 = 0;
            var g: u64 = 0;
            var b: u64 = 0;
            var a: u64 = 0;
            var n: u64 = 0;
            for (y0..y1) |y| for (x0..x1) |x| {
                const p = src[y * sw + x];
                r += p.r;
                g += p.g;
                b += p.b;
                if (has_alpha) a += p.a;
                n += 1;
            };
            var out: Px = undefined;
            out.r = @intCast(r / n);
            out.g = @intCast(g / n);
            out.b = @intCast(b / n);
            if (has_alpha) out.a = @intCast(a / n);
            dst[dy * dw + dx] = out;
        }
    }
}

/// jpeg has no alpha; blend transparent pixels onto white like a browser would
fn compositeOnWhite(alloc: Allocator, src: zigimg.Image) !zigimg.Image {
    var dst = try zigimg.Image.create(alloc, src.width, src.height, .rgb24);
    errdefer dst.deinit(alloc);
    const px = src.pixels.rgba32;
    for (px, 0..) |p, i| {
        const a: u32 = p.a;
        dst.pixels.rgb24[i] = .{
            .r = @intCast((@as(u32, p.r) * a + 255 * (255 - a)) / 255),
            .g = @intCast((@as(u32, p.g) * a + 255 * (255 - a)) / 255),
            .b = @intCast((@as(u32, p.b) * a + 255 * (255 - a)) / 255),
        };
    }
    return dst;
}

// --- tests ---

const t = std.testing;

test "allowed domains" {
    try t.expect(isAllowedUrl("https://all-the.bufo.zone/bufo-hugs-moo-deng.png"));
    try t.expect(isAllowedUrl("https://find-bufo.com/e/bufo.png"));
    try t.expect(!isAllowedSimilarityUrl("https://cdn.bsky.app/img/feed_fullsize/plain/did:plc:test/bafkrei@jpeg"));
    try t.expect(!isAllowedUrl("https://evil.example/x.png"));
    try t.expect(!isAllowedUrl("https://all-the.bufo.zone.evil.example/x.png"));
    try t.expect(!isAllowedUrl("https://evil.example/all-the.bufo.zone/x.png"));
    try t.expect(!isAllowedUrl("https://all-the.bufo.zone@evil.example/x.png"));
    try t.expect(!isAllowedUrl("http://all-the.bufo.zone/x.png"));
}

test "sniff by magic bytes, then extension" {
    try t.expectEqual(Kind.png, sniff("\x89PNG\r\n\x1a\nxxxx", "whatever"));
    try t.expectEqual(Kind.gif, sniff("GIF89a...", "x"));
    try t.expectEqual(Kind.jpeg, sniff("\xff\xd8\xff\xe0", "x"));
    try t.expectEqual(Kind.webp, sniff("RIFF\x04\x00\x00\x00WEBPrest", "x"));
    try t.expectEqual(Kind.png, sniff("", "https://x/a.png"));
    try t.expectEqual(Kind.other, sniff("", "https://x/a"));
}

test "upload content type comes from image bytes, not a filename" {
    try t.expectEqualStrings("image/png", uploadContentType("\x89PNG\r\n\x1a\nrest").?);
    try t.expectEqualStrings("image/webp", uploadContentType("RIFF\x04\x00\x00\x00WEBPrest").?);
    try t.expect(uploadContentType("<svg></svg>") == null);
    try t.expect(uploadContentType("") == null);
}

test "small images pass through untouched" {
    var arena = std.heap.ArenaAllocator.init(t.allocator);
    defer arena.deinit();
    const bytes = "\x89PNG\r\n\x1a\nnot-really";
    const out = try shrink(arena.allocator(), bytes, .png, 1000);
    try t.expectEqualStrings("image/png", out.content_type);
    try t.expect(out.bytes.ptr == bytes.ptr);
}

/// incompressible noise, so the png is big and the shrink ladder is exercised
fn noisePng(alloc: Allocator, side: usize) ![]u8 {
    var img = try zigimg.Image.create(alloc, side, side, .rgba32);
    defer img.deinit(alloc);
    var prng = std.Random.DefaultPrng.init(42);
    const rand = prng.random();
    for (img.pixels.rgba32) |*p| {
        p.* = .{ .r = rand.int(u8), .g = rand.int(u8), .b = rand.int(u8), .a = 255 };
    }
    const buf = try alloc.alloc(u8, side * side * 4 + (1 << 16));
    return try img.writeToMemory(alloc, buf, .{ .png = .{} });
}

test "oversized png is shrunk under the budget and still decodes" {
    var arena = std.heap.ArenaAllocator.init(t.allocator);
    defer arena.deinit();
    const alloc = arena.allocator();

    const original = try noisePng(alloc, 128);
    try t.expect(original.len > 40_000);
    const budget = original.len / 4;

    const out = try shrink(alloc, original, .png, budget);
    try t.expect(out.bytes.len <= budget);
    try t.expect(out.bytes.ptr != original.ptr);

    var decoded = try zigimg.Image.fromMemory(alloc, out.bytes);
    defer decoded.deinit(alloc);
    try t.expect(decoded.width < 128);
    try t.expect(mem.eql(u8, out.content_type, "image/png") or mem.eql(u8, out.content_type, "image/jpeg"));
}

test "box filter averages blocks" {
    const Px = zigimg.color.Rgba32;
    const src = [_]Px{
        .{ .r = 0, .g = 0, .b = 0, .a = 255 },   .{ .r = 200, .g = 0, .b = 0, .a = 255 },
        .{ .r = 100, .g = 0, .b = 0, .a = 255 }, .{ .r = 100, .g = 0, .b = 0, .a = 255 },
    };
    var dst = [_]Px{undefined};
    boxFilter(Px, &src, 2, 2, &dst, 1, 1, true);
    try t.expectEqual(@as(u8, 100), dst[0].r);
    try t.expectEqual(@as(u8, 255), dst[0].a);
}