AI Image Generator — Zig source

//! AI Image Generator — prompt validation, gate verdicts, and ?prompt= URL
//! state for FLUX.1-schnell (Zig 0.13, stdlib only). Port of src/lib/image-gen.ts.

const std = @import("std");

pub const model = "@cf/black-forest-labs/flux-1-schnell"; // + 4 steps = 57.6 n/img
pub const steps = 4;
pub const prompt_max_chars = 600; // UX clamp, below the model's 2048 limit
pub const ip_hour_cap = 10;
pub const daily_cap = 100;
pub const PromptInvalid = error{ Empty, TooLong, ControlChars };

/// Daily cap wins; retry is minutes (UTC-midnight reset for daily, 60 for IP).
pub const GateVerdict = union(enum) { ok, rate_limited_ip, rate_limited_daily: u32 };

fn hexval(c: u8) u8 { return if (c <= '9') c - '0' else (c | 32) - 'a' + 10; }

/// Normalize CRLF/CR to \n and trim in one pass; length is bytes (stricter
/// than the TS char count), then reject control chars except \n (C1 is
/// exactly lead 0xC2 + continuation 0x80..0x9F). Returns the length in `out`.
pub fn validate_prompt(raw: []const u8, out: []u8) PromptInvalid!usize {
    var n: usize = 0;
    var started = false;
    for (raw, 0..) |b, i| {
        var c = b;
        if (c == '\r') { if (i + 1 < raw.len and raw[i + 1] == '\n') continue; c = '\n'; }
        if (!started and std.ascii.isWhitespace(c)) continue;
        started = true;
        out[n] = c; n += 1;
    }
    while (n > 0 and std.ascii.isWhitespace(out[n - 1])) n -= 1;
    if (n == 0) return error.Empty;
    if (n > prompt_max_chars) return error.TooLong;
    for (out[0..n], 0..) |c, j| {
        const cont = (c & 0xC0) == 0x80; // UTF-8 continuation byte
        const c1 = c == 0xC2 and j + 1 < n and (out[j + 1] & 0xE0) == 0x80; // C1 block
        if (!cont and (c1 or (c < 0x20 and c != '\n') or c == 0x7F)) return error.ControlChars;
    }
    return n;
}

/// Daily cap wins: retry = minutes (>= 1) to the next UTC midnight via pure
/// 86,400-second day arithmetic; the IP window is a fixed 60 minutes.
pub fn check_gate(ip_count: u32, daily_count: u32, now_s: i64) GateVerdict {
    if (daily_count >= daily_cap) {
        const left: i64 = 86_400 - @mod(now_s, 86_400);
        return .{ .rate_limited_daily = @intCast(@divTrunc(left + 59, 60)) };
    }
    if (ip_count >= ip_hour_cap) return .rate_limited_ip;
    return .ok;
}

/// Encode the shareable state into `out` ("?prompt=…", %XX-escaped; empty
/// when there is nothing to share). Needs 8 + 3 * prompt.len bytes.
pub fn encode_url_state(prompt: []const u8, out: []u8) []const u8 {
    if (prompt.len == 0) return out[0..0];
    const hex = "0123456789ABCDEF";
    @memcpy(out[0..8], "?prompt=");
    var n: usize = 8;
    for (prompt) |c| {
        const unreserved = std.ascii.isAlphanumeric(c) or c == '-' or c == '.' or c == '_' or c == '~';
        if (unreserved) { out[n] = c; n += 1; }
        else { out[n] = '%'; out[n + 1] = hex[c >> 4]; out[n + 2] = hex[c & 15]; n += 3; }
    }
    return out[0..n];
}

/// Inverse of encode_url_state — this tool only encodes ?prompt=… as the
/// leading pair, so decode just that (well-formed %XX assumed).
pub fn decode_url_state(search: []const u8, out: []u8) []const u8 {
    if (!std.mem.startsWith(u8, search, "?prompt=")) return out[0..0];
    var n: usize = 0;
    var i: usize = 8;
    while (i < search.len and search[i] != '&') : (i += 1) {
        if (search[i] == '%' and i + 2 < search.len) {
            out[n] = (hexval(search[i + 1]) << 4) | hexval(search[i + 2]); n += 1; i += 2;
        } else { out[n] = search[i]; n += 1; }
    }
    return out[0..n];
}