FabledCurator/agent/fc_agent/media.py

"""Image + video handling. Stills load directly; videos are sampled into frames
(ffmpeg) at the cadence FC sends — so a video becomes a bag of per-frame
instances, each with a timestamp."""
import io
import os
import subprocess
import tempfile
import time

from PIL import Image, ImageFile

# Load slightly-truncated images (a few missing trailing bytes) instead of
# raising — matches the server embedder. These are common in scraped libraries
# and would otherwise fail the job 3× then error (operator-flagged 2026-06-30).
ImageFile.LOAD_TRUNCATED_IMAGES = True

# Disable PIL's decompression-bomb guard: this is a TRUSTED local library, not an
# untrusted upload surface, so a legitimately huge image (high-res scans/prints,
# 90M+ pixels) must load. The default 89M-pixel limit only WARNS, but PIL raises
# DecompressionBombError at 2× (~179M px) — which would fail those jobs outright
# (operator-flagged 2026-06-30, images of 90–95M px).
Image.MAX_IMAGE_PIXELS = None


def is_video(mime: str) -> bool:
    return bool(mime) and (mime.startswith("video/") or mime in {"image/gif"})


def _dhash(img: Image.Image, size: int = 8) -> int:
    """Difference hash: compare adjacent pixels of a (size+1 × size) grayscale
    thumbnail → a `size*size`-bit fingerprint. Cheap (64 comparisons on a 72-px
    thumbnail) and robust to scaling/compression noise — near-identical frames
    hash within a few bits, a real scene change moves many."""
    small = img.convert("L").resize((size + 1, size))
    px = list(small.getdata())
    bits = 0
    for row in range(size):
        base = row * (size + 1)
        for col in range(size):
            bits = (bits << 1) | int(px[base + col] > px[base + col + 1])
    return bits


def dedupe_frames(
    frames: list[tuple[float, Image.Image]], min_distance: int
) -> list[tuple[float, Image.Image]]:
    """Drop visually near-duplicate frames. A near-static video sampled into many
    frames re-runs the WHOLE detect→CCIP→SigLIP chain on ~identical frames — the
    dominant video load. Greedy perceptual-hash dedup: keep a frame only if its
    dHash differs from every already-kept frame by >= min_distance bits (Hamming),
    so a static run collapses to one frame while genuinely distinct scenes all
    survive. Order + timestamps preserved. CPU-only (64-bit int XORs), so it runs
    in the decode stage and spares the GPU the skipped frames entirely.

    min_distance is the coarseness dial: higher keeps more frames (safer for brief
    localized changes an 8×8 hash can miss), 0 disables. The first frame is always
    kept (nothing to compare against)."""
    if min_distance <= 0 or len(frames) <= 1:
        return frames
    kept: list[tuple[float, Image.Image]] = []
    hashes: list[int] = []
    for t, frame in frames:
        h = _dhash(frame)
        if all(bin(h ^ k).count("1") >= min_distance for k in hashes):
            hashes.append(h)
            kept.append((t, frame))
    return kept


def to_rgb(img: Image.Image) -> Image.Image:
    """RGB, flattening any transparency onto white first. A naive convert('RGB')
    on a palette-with-transparency image (common for character PNGs on a clear
    background) lets PIL guess the transparent pixels — usually black artifacts
    that bleed into the crop + the embedding (and the "should be converted to
    RGBA" warning). Compositing over white gives a clean, consistent background."""
    if img.mode in ("RGBA", "LA", "PA") or (
        img.mode == "P" and "transparency" in img.info
    ):
        img = img.convert("RGBA")
        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
        return Image.alpha_composite(bg, img).convert("RGB")
    return img.convert("RGB")


def load_image(data: bytes) -> Image.Image:
    return to_rgb(Image.open(io.BytesIO(data)))


# ffmpeg reconnect flags — resume a dropped HTTP transfer (a slow/contended media
# store can cut a long stream) instead of failing the whole job. Relies only on
# HTTP + Range, which every FC deployment serves → environment-agnostic.
_RECONNECT = [
    "-reconnect", "1", "-reconnect_streamed", "1",
    "-reconnect_on_network_error", "1", "-reconnect_delay_max", "5",
]


def _collect_frames(
    tmp: str, interval: float, cap: int
) -> list[tuple[float, Image.Image]]:
    out: list[tuple[float, Image.Image]] = []
    names = sorted(n for n in os.listdir(tmp) if n.startswith("f_"))
    for i, name in enumerate(names[:cap]):
        with Image.open(os.path.join(tmp, name)) as im:
            out.append((round(i * interval, 2), to_rgb(im)))
    return out


def _terminate(proc: subprocess.Popen) -> None:
    """Stop an ffmpeg cleanly, then hard-kill if it ignores SIGTERM."""
    proc.terminate()
    try:
        proc.wait(timeout=2)
    except subprocess.TimeoutExpired:
        proc.kill()
        try:
            proc.wait(timeout=2)
        except subprocess.TimeoutExpired:
            pass


def sample_frames_from_url(
    url: str, interval_seconds: float, max_frames: int,
    *, headers: str = "", timeout: float = 1200.0, should_stop=None,
) -> list[tuple[float, Image.Image]]:
    """Sample frames by pointing ffmpeg STRAIGHT at the media URL — it Range-reads
    only the video index + up to max_frames worth of content, so the agent never
    downloads the whole file (VR/4K originals run 800MB+ and would buffer ~1GB in
    RAM and get cut off mid-download). Reconnect flags resume a dropped transfer;
    the timeout is the per-video ceiling (a slow/reconnecting stream can otherwise
    run for minutes). `should_stop` is polled while ffmpeg runs so a Stop KILLS the
    subprocess at once — otherwise a downloader stuck in a long decode keeps the
    agent "working" long after Stop. Empty on failure / stop / timeout."""
    interval = max(0.5, float(interval_seconds or 4.0))
    cap = max(1, int(max_frames or 64))
    hdr = ["-headers", headers] if headers else []
    with tempfile.TemporaryDirectory() as tmp:
        pattern = os.path.join(tmp, "f_%05d.jpg")
        cmd = ["ffmpeg", "-nostdin", "-loglevel", "error", *_RECONNECT, *hdr,
               "-i", url, "-vf", f"fps=1/{interval}", "-frames:v", str(cap),
               "-q:v", "3", pattern]
        try:
            proc = subprocess.Popen(
                cmd, stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            )
        except (OSError, ValueError):
            return []
        # Poll rather than block, so a Stop (or the per-video timeout) can kill a
        # slow/wedged ffmpeg promptly instead of waiting it out.
        start = time.monotonic()
        while True:
            try:
                proc.wait(timeout=0.5)
                break
            except subprocess.TimeoutExpired:
                if (should_stop and should_stop()) or (time.monotonic() - start > timeout):
                    _terminate(proc)
                    return []
        if proc.returncode != 0:
            return []
        return _collect_frames(tmp, interval, cap)