AI Image Generator — Python source

"""AI Image Generator — prompt validation, usage-gate verdicts, and ?prompt=
URL state for FLUX.1-schnell image generation.

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the AI Image Generator tool,
          from src/lib/image-gen.ts (the canonical TypeScript implementation).
"""

from __future__ import annotations

import math
import re
from datetime import datetime, timedelta, timezone
from urllib.parse import parse_qs, quote, urlencode

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

# Control chars except \n: C0 minus U+000A, DEL, and C1.
_CONTROL_CHARS = re.compile(r"[\x00-\x09\x0B-\x1F\x7F-\x9F]")


def validate_prompt(raw: str) -> tuple[str | None, str | None]:
    """Return (prompt, None) on success or (None, reason) where reason is
    'empty' | 'tooLong' | 'controlChars'."""
    prompt = re.sub(r"\r\n?", "\n", raw).strip()  # normalize line endings first
    if not prompt:
        return None, "empty"
    if len(prompt) > PROMPT_MAX_CHARS:
        return None, "tooLong"
    if _CONTROL_CHARS.search(prompt):
        return None, "controlChars"
    return prompt, None


def minutes_to_utc_midnight(now: datetime) -> int:
    """Minutes (>= 1) until the next UTC midnight — the daily window reset."""
    utc = now.astimezone(timezone.utc)
    midnight = datetime(utc.year, utc.month, utc.day, tzinfo=timezone.utc) + timedelta(days=1)
    return max(1, math.ceil((midnight - utc).total_seconds() / 60))


def check_gate(
    ip_count: int, daily_count: int, now: datetime | None = None
) -> dict:
    """Daily cap wins: the site-global budget resets at UTC midnight."""
    if daily_count >= DAILY_CAP:
        return {
            "ok": False,
            "error": "RATE_LIMITED_DAILY",
            "retryAfterMinutes": minutes_to_utc_midnight(now or datetime.now(timezone.utc)),
        }
    if ip_count >= IP_HOUR_CAP:
        return {"ok": False, "error": "RATE_LIMITED_IP", "retryAfterMinutes": 60}
    return {"ok": True}


def encode_url_state(prompt: str) -> str:
    """The shareable state: '' when there is nothing to share. quote_via=quote
    gives encodeURIComponent-style %20 (not form-style +) for spaces."""
    if not prompt:
        return ""
    return "?prompt=" + urlencode({"prompt": prompt}, quote_via=quote)


def decode_url_state(search: str) -> str:
    """Inverse of encode_url_state — a missing param decodes to ''."""
    return parse_qs(search.lstrip("?")).get("prompt", [""])[0]