AI Image Generator — Rust source

//! AI Image Generator — prompt validation, usage-gate verdicts, and ?prompt=
//! URL state for FLUX.1-schnell image generation.
//! Language: Rust (edition 2021, standard library only). Port of src/lib/image-gen.ts.

pub const MODEL: &str = "@cf/black-forest-labs/flux-1-schnell";
pub const STEPS: u32 = 4; // flux-1-schnell's design point -> 57.6 neurons/image
pub const PROMPT_MAX_CHARS: usize = 600; // UX clamp, below the model's 2048 limit
pub const IP_HOUR_CAP: u32 = 10;
pub const DAILY_CAP: u32 = 100;

/// Why a prompt was refused.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PromptInvalid { Empty, TooLong, ControlChars }

/// Gate verdict — daily cap wins; the retry hint is in minutes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GateVerdict { Ok, RateLimitedIp, RateLimitedDaily { retry_after_minutes: u32 } }

/// Normalize CRLF/CR to `\n`, trim, enforce 1..=600 chars, and reject control
/// characters except `\n` (C0 minus U+000A, DEL, C1). Counted in `char`s.
pub fn validate_prompt(raw: &str) -> Result<String, PromptInvalid> {
    let prompt = raw.replace("\r\n", "\n").replace('\r', "\n").trim().to_string();
    if prompt.is_empty() { return Err(PromptInvalid::Empty); }
    if prompt.chars().count() > PROMPT_MAX_CHARS { return Err(PromptInvalid::TooLong); }
    let is_ctrl = |c: char| { let v = c as u32; (v < 0x20 && v != 0x0A) || (0x7F..=0x9F).contains(&v) };
    if prompt.chars().any(is_ctrl) { return Err(PromptInvalid::ControlChars); }
    Ok(prompt)
}

/// Minutes (>= 1) until the next UTC midnight, from epoch milliseconds —
/// pure 86,400,000-ms day arithmetic, no calendar machinery needed.
pub fn minutes_to_utc_midnight(now_ms: u64) -> u32 {
    let ms_left = (now_ms / 86_400_000 + 1) * 86_400_000 - now_ms;
    (ms_left + 59_999) as u32 / 60_000
}

pub fn check_gate(ip_count: u32, daily_count: u32, now_ms: u64) -> GateVerdict {
    if daily_count >= DAILY_CAP {
        return GateVerdict::RateLimitedDaily { retry_after_minutes: minutes_to_utc_midnight(now_ms) };
    }
    if ip_count >= IP_HOUR_CAP { return GateVerdict::RateLimitedIp; }
    GateVerdict::Ok
}

/// Percent-encode a query value: unreserved chars pass, the rest become %XX.
fn percent_encode(value: &str) -> String {
    let mut out = String::with_capacity(value.len());
    for b in value.bytes() {
        let unreserved = b.is_ascii_alphanumeric() || matches!(b, b'-' | b'.' | b'_' | b'~');
        if unreserved { out.push(b as char); } else { out.push_str(&format!("%{b:02X}")); }
    }
    out
}

fn percent_decode(value: &str) -> String {
    let b = value.as_bytes();
    let mut out = Vec::with_capacity(b.len());
    let mut i = 0;
    while i < b.len() {
        if b[i] == b'%' && i + 2 < b.len() {
            if let Ok(byte) = u8::from_str_radix(&value[i + 1..i + 3], 16) { out.push(byte); i += 3; continue; }
        }
        out.push(b[i]);
        i += 1;
    }
    String::from_utf8_lossy(&out).into_owned()
}

/// Encode the shareable state: empty when there is nothing to share.
pub fn encode_url_state(prompt: &str) -> String {
    if prompt.is_empty() { String::new() } else { format!("?prompt={}", percent_encode(prompt)) }
}

/// Inverse of [`encode_url_state`] — a missing param decodes to "".
pub fn decode_url_state(search: &str) -> String {
    search.strip_prefix('?').unwrap_or(search).split('&')
        .find_map(|pair| pair.strip_prefix("prompt=")).map(percent_decode).unwrap_or_default()
}