backend/sanitize.py

"""Guardrails on everything that reaches the model.

We never trust the client. The browser already shrinks images, but a direct
POST to /api/interpret could carry anything — so we decode, validate, and
re-encode every image server-side to a known-safe size and format, and we cap
and clean the fragment. Bad input raises InputError, which the API turns into a
clean 4xx instead of letting it reach (or crash) the model.
"""

from __future__ import annotations

import base64
import binascii
import io
import re

from PIL import Image, UnidentifiedImageError

MAX_IMAGES = 4
# Reject a single source image larger than this (decoded bytes) before we touch it.
MAX_SOURCE_BYTES = 12 * 1024 * 1024  # 12 MB
# What we hand the model: longest edge and JPEG quality after re-encoding.
TARGET_MAX_DIM = 896
JPEG_QUALITY = 72
# A fragment is a line or two, not an essay.
MAX_FRAGMENT_CHARS = 600
# Cap the "already seen" list so a long session can't bloat the prompt.
MAX_EXCLUDE = 24
MAX_TITLE_CHARS = 120

_DATA_URL_RE = re.compile(r"^data:(?P<mime>[\w.+-]+/[\w.+-]+)?;base64,(?P<data>.+)$", re.DOTALL)
# Control chars except tab/newline/carriage-return.
_CONTROL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")


class InputError(Exception):
    """User-supplied input that we refuse to forward to the model."""


def sanitize_fragment(text: str) -> str:
    """Trim, strip control characters, collapse runaway whitespace, and cap length."""
    if not text:
        return ""
    if not isinstance(text, str):
        raise InputError("Fragment must be text.")
    cleaned = _CONTROL_RE.sub("", text).strip()
    # Collapse 3+ blank lines / long whitespace runs that bloat the prompt.
    cleaned = re.sub(r"[ \t]{2,}", " ", cleaned)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    if len(cleaned) > MAX_FRAGMENT_CHARS:
        cleaned = cleaned[:MAX_FRAGMENT_CHARS].rstrip() + "…"
    return cleaned


def sanitize_titles(titles: list[str]) -> list[str]:
    """Clean and de-duplicate the 'already seen' titles, capped in count and length."""
    if not titles:
        return []
    if not isinstance(titles, list):
        raise InputError("Exclude must be a list of titles.")
    out: list[str] = []
    seen: set[str] = set()
    for t in titles:
        if not isinstance(t, str):
            continue
        cleaned = _CONTROL_RE.sub("", t).strip()[:MAX_TITLE_CHARS]
        key = cleaned.lower()
        if cleaned and key not in seen:
            seen.add(key)
            out.append(cleaned)
        if len(out) >= MAX_EXCLUDE:
            break
    return out


def _decode(one: str) -> bytes:
    """Accept a data URL or raw base64 and return the decoded bytes."""
    if not isinstance(one, str) or not one.strip():
        raise InputError("Empty image entry.")
    m = _DATA_URL_RE.match(one.strip())
    payload = m.group("data") if m else one.strip()
    try:
        raw = base64.b64decode(payload, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise InputError("An image wasn't valid base64.") from exc
    if not raw:
        raise InputError("An image decoded to nothing.")
    if len(raw) > MAX_SOURCE_BYTES:
        raise InputError("An image is too large (max 12 MB).")
    return raw


def _reencode(raw: bytes) -> str:
    """Open, normalize, downscale, and re-encode to a JPEG data URL we control."""
    try:
        img = Image.open(io.BytesIO(raw))
        img.load()
    except (UnidentifiedImageError, OSError) as exc:
        raise InputError("An image couldn't be read.") from exc

    # Flatten transparency/odd modes onto black, matching the cinematic ground.
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGBA") if "A" in img.mode else img.convert("RGB")
        if img.mode == "RGBA":
            bg = Image.new("RGB", img.size, (12, 9, 7))
            bg.paste(img, mask=img.split()[-1])
            img = bg
    if img.mode != "RGB":
        img = img.convert("RGB")

    w, h = img.size
    if max(w, h) > TARGET_MAX_DIM:
        if w >= h:
            img = img.resize((TARGET_MAX_DIM, round(h * TARGET_MAX_DIM / w)))
        else:
            img = img.resize((round(w * TARGET_MAX_DIM / h), TARGET_MAX_DIM))

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=JPEG_QUALITY)
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/jpeg;base64,{b64}"


def sanitize_images(images: list[str]) -> list[str]:
    """Validate, downscale, and re-encode each image. Returns safe JPEG data URLs."""
    if not isinstance(images, list) or not images:
        raise InputError("Bring at least one image.")
    if len(images) > MAX_IMAGES:
        raise InputError(f"Up to {MAX_IMAGES} images, please.")
    return [_reencode(_decode(one)) for one in images]