tl/scenes/thumbnails.py

import json
import logging
import math
import shutil
import subprocess
import tempfile
import threading
from pathlib import Path

from django.core.files import File
from django.core.files.storage import default_storage

from .models import Project

logger = logging.getLogger(__name__)

THUMB_INTERVAL_SECONDS = 1.0
THUMB_MAX_HEIGHT = 360
THUMB_JPEG_QUALITY = 9
# ffmpeg's select filter decodes the whole clip on every pass, so each chunk is
# a full re-decode. Keep it one pass (one decode) for all but absurdly long clips
# — thousands of terms evaluate fine once the select expression is balanced (see
# _balanced_sum); the limit is parser depth, not term count.
CLIP_SELECT_CHUNK_SIZE = 4096


def _otio_frames(rational_time):
    return rational_time.get("value", 0) if rational_time else 0


def _clip_sample_frame(start, duration):
    end = start + duration
    first = math.ceil(start)
    last = math.ceil(end) - 1
    if last < first:
        return max(0, math.floor(start))
    midpoint = math.floor(start + (duration / 2))
    return max(0, min(max(midpoint, first), last))


def _clip_tiles(start, duration, fps):
    frames_per_tile = max(1, round(fps * THUMB_INTERVAL_SECONDS))
    tiles = []
    offset = 0
    while offset < duration:
        end = min(duration, offset + frames_per_tile)
        sample = _clip_sample_frame(start + offset, end - offset)
        tiles.append({"start": offset, "end": end, "frame": sample})
        offset = end
    return tiles


def _otio_clip_items(project):
    if not project.otio:
        return []

    # read through the storage backend (works on local disk AND R2/S3)
    with project.otio.open("rb") as fh:
        otio = json.loads(fh.read())

    tracks = otio.get("tracks", {}).get("children", [])
    source_starts = [
        _otio_frames(item.get("source_range", {}).get("start_time"))
        for track in tracks
        for item in track.get("children", [])
        if str(item.get("OTIO_SCHEMA", "")).startswith("Clip")
    ]
    media_offset = min(source_starts) if source_starts else 0

    clips = []
    for track_index, track in enumerate(tracks):
        clip_index = 0
        for item in track.get("children", []):
            if not str(item.get("OTIO_SCHEMA", "")).startswith("Clip"):
                continue
            source_range = item.get("source_range", {})
            start = _otio_frames(source_range.get("start_time")) - media_offset
            duration = _otio_frames(source_range.get("duration"))
            clips.append({
                "id": f"t{track_index}-c{clip_index}",
                "tiles": _clip_tiles(start, duration, project.fps),
            })
            clip_index += 1
    return clips


def _thumb_key(project_id, filename):
    return f"projects/{project_id}/thumbs/v1/{filename}"


def _save_thumb(project_id, filename, local_path):
    """Upload one tile through the storage backend (R2 in prod, disk in dev) and
    return its public URL. Overwrites any prior tile at that key."""
    key = _thumb_key(project_id, filename)
    if default_storage.exists(key):
        default_storage.delete(key)
    with open(local_path, "rb") as fh:
        default_storage.save(key, File(fh))
    return default_storage.url(key)


def _chunks(items, size):
    for i in range(0, len(items), size):
        yield i // size, items[i:i + size]


def _balanced_sum(terms):
    """Join `terms` with `+` as a *balanced* tree, e.g. ((a+b)+(c+d)).

    ffmpeg's expression parser is recursive-descent: a flat "a+b+c+..." chain
    nests right-associatively to depth N and overflows the parser past ~120
    terms ("Error while parsing expression"). Balancing makes the parse depth
    O(log N), so thousands of terms evaluate fine. `+` is associative, so the
    selection is identical either way."""
    nodes = list(terms)
    if not nodes:
        return "0"
    while len(nodes) > 1:
        nodes = [
            f"({nodes[i]}+{nodes[i + 1]})" if i + 1 < len(nodes) else nodes[i]
            for i in range(0, len(nodes), 2)
        ]
    return nodes[0]


def enqueue_thumbnail_generation(project_id):
    thread = threading.Thread(
        target=generate_project_thumbnails,
        args=(project_id,),
        name=f"thumbs-project-{project_id}",
        daemon=True,
    )
    thread.start()


def generate_project_thumbnails(project_id, log=None):
    """Build the thumbnail manifest for a project. `log` is an optional
    callable(str) for progress lines — the management command passes one so the
    run is watchable on the console instead of only in app logs."""
    emit = log if log is not None else (lambda msg: logger.info("[thumbs %s] %s", project_id, msg))

    try:
        project = Project.objects.get(pk=project_id)
    except Project.DoesNotExist:
        emit("project does not exist")
        return

    if not project.clip:
        Project.objects.filter(pk=project_id).update(
            thumbnail_status="failed",
            thumbnail_error="project has no clip",
        )
        emit("failed: project has no clip")
        return

    Project.objects.filter(pk=project_id).update(
        thumbnail_status="running",
        thumbnail_error="",
        thumbnail_manifest={},
    )

    work = Path(tempfile.mkdtemp(prefix=f"thumbs-{project_id}-"))
    try:
        # pull the clip down to a local temp file so ffmpeg reads it directly,
        # regardless of whether storage is local disk or R2.
        emit(f"downloading clip {project.clip.name}")
        clip_path = work / ("clip" + (Path(project.clip.name).suffix or ".mp4"))
        with project.clip.open("rb") as src, open(clip_path, "wb") as dst:
            shutil.copyfileobj(src, dst, length=1024 * 1024)

        scale = f"scale=w=-2:h={THUMB_MAX_HEIGHT}:force_original_aspect_ratio=decrease"

        # the timeline only renders per-clip tiles, so extract exactly the source
        # frames those tiles need (one select-pass), not a full 1 fps sequence.
        clip_items = _otio_clip_items(project)
        unique_clip_frames = sorted({
            max(0, tile["frame"])
            for clip in clip_items
            for tile in clip["tiles"]
        })
        if not unique_clip_frames:
            raise RuntimeError("no clip frames to sample")

        chunks = list(_chunks(unique_clip_frames, CLIP_SELECT_CHUNK_SIZE))
        emit(f"{len(unique_clip_frames)} frames across {len(clip_items)} clips, "
             f"{len(chunks)} ffmpeg pass(es)")

        frame_filenames = {}
        for chunk_index, frames in chunks:
            emit(f"  ffmpeg pass {chunk_index + 1}/{len(chunks)} ({len(frames)} frames)…")
            select_expr = _balanced_sum(f"eq(n\\,{frame})" for frame in frames)
            proc = subprocess.run([
                "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                "-i", str(clip_path),
                "-vf", f"select='{select_expr}',{scale}",
                "-vsync", "0",
                "-q:v", str(THUMB_JPEG_QUALITY),
                str(work / f"clipframe_{chunk_index:03d}_%06d.jpg"),
            ], capture_output=True, text=True)
            if proc.returncode != 0:
                # surface ffmpeg's stderr — CalledProcessError hides it, which is
                # what turned the last failure into a guessing game.
                raise RuntimeError(
                    f"ffmpeg exited {proc.returncode}: {proc.stderr.strip() or '(no stderr)'}"
                )
            for idx, frame in enumerate(frames, start=1):
                frame_filenames[frame] = f"clipframe_{chunk_index:03d}_{idx:06d}.jpg"

        # upload each needed tile once, then point the manifest at its public URL
        emit("uploading tiles…")
        uploaded = {}
        clip_thumbnails = {}
        for clip in clip_items:
            tiles = []
            for tile in clip["tiles"]:
                fname = frame_filenames.get(tile["frame"])
                if not fname or not (work / fname).exists():
                    continue
                if fname not in uploaded:
                    uploaded[fname] = _save_thumb(project_id, fname, work / fname)
                tiles.append({"start": tile["start"], "end": tile["end"], "url": uploaded[fname]})
            if tiles:
                clip_thumbnails[clip["id"]] = tiles

        if not clip_thumbnails:
            raise RuntimeError("ffmpeg produced no thumbnails")

        manifest = {
            "version": 1,
            "interval": THUMB_INTERVAL_SECONDS,
            "max_height": THUMB_MAX_HEIGHT,
            "quality": THUMB_JPEG_QUALITY,
            "count": len(uploaded),
            "clips": clip_thumbnails,
        }
        Project.objects.filter(pk=project_id).update(
            thumbnail_status="ready",
            thumbnail_manifest=manifest,
            thumbnail_error="",
        )
        emit(f"ready: {len(uploaded)} tiles uploaded for {len(clip_thumbnails)} clips")
    except Exception as exc:
        logger.exception("thumbnail generation failed for project %s", project_id)
        Project.objects.filter(pk=project_id).update(
            thumbnail_status="failed",
            thumbnail_error=str(exc),
            thumbnail_manifest={},
        )
        emit(f"failed: {exc}")
    finally:
        shutil.rmtree(work, ignore_errors=True)