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
CLIP_SELECT_CHUNK_SIZE = 30


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 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):
    try:
        project = Project.objects.get(pk=project_id)
    except Project.DoesNotExist:
        return

    if not project.clip:
        Project.objects.filter(pk=project_id).update(
            thumbnail_status="failed",
            thumbnail_error="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.
        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")

        frame_filenames = {}
        for chunk_index, frames in _chunks(unique_clip_frames, CLIP_SELECT_CHUNK_SIZE):
            select_expr = "+".join(f"eq(n\\,{frame})" for frame in frames)
            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"),
            ], check=True, capture_output=True, text=True)
            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
        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="",
        )
    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={},
        )
    finally:
        shutil.rmtree(work, ignore_errors=True)