warmup_thumbs.py

#!/usr/bin/env python3
"""Pre-generate thumbnail cache for all media files.

Walks the database for image and video media, generates thumbnails
for any that don't already have a cached version on disk.

Usage:
    python warmup_thumbs.py --workers 8 --type image
    python warmup_thumbs.py --workers 4 --type video
    python warmup_thumbs.py --workers 8 --type all --limit 1000
"""

import argparse
import hashlib
import logging
import os
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import psycopg2
from PIL import Image

THUMB_CACHE_DIR = Path(__file__).parent / ".thumb_cache"
THUMB_MAX_WIDTH = 400
THUMB_QUALITY = 75
DB_URL = os.environ.get("DB_URL", "postgresql://coomer:coomer_dev@localhost:5433/coomer")

_OLD_GVFS_PREFIX = "/run/user/1000/gvfs/smb-share:server=unas-pro-8.local,share=drive_2"
NAS_BASE = os.environ.get("NAS_BASE", "/mnt/nas/drive_2")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger(__name__)


def _resolve_path(local_path: str) -> str:
    """Resolve media path (mirrors main.py logic)."""
    if local_path.startswith(_OLD_GVFS_PREFIX):
        local_path = NAS_BASE + local_path[len(_OLD_GVFS_PREFIX):]
    return local_path


def generate_thumb(
    media_id: int,
    local_path: str,
    cache_dir: Path = THUMB_CACHE_DIR,
    width: int = THUMB_MAX_WIDTH,
) -> bool:
    """Generate a single image thumbnail. Returns True if generated, False if skipped."""
    cache_key = hashlib.md5(f"{media_id}:{width}:jpg".encode()).hexdigest()
    cache_path = cache_dir / f"{cache_key}.jpg"

    if cache_path.exists():
        return False  # already cached

    resolved = _resolve_path(local_path)
    if not os.path.isfile(resolved):
        return False  # source file missing

    try:
        with Image.open(resolved) as img:
            img = img.convert("RGB")
            ratio = width / img.width
            new_h = int(img.height * ratio)
            img = img.resize((width, new_h), Image.LANCZOS)
            img.save(str(cache_path), "JPEG", quality=THUMB_QUALITY, optimize=True)
        return True
    except Exception as e:
        log.warning("Failed to generate thumb for media %d: %s", media_id, e)
        # Clean up partial file
        if cache_path.exists():
            cache_path.unlink()
        return False


def generate_video_thumb(
    media_id: int,
    local_path: str,
    cache_dir: Path = THUMB_CACHE_DIR,
    width: int = THUMB_MAX_WIDTH,
) -> bool:
    """Generate a single video thumbnail via ffmpeg. Returns True if generated."""
    cache_path = cache_dir / f"video_{media_id}.jpg"

    if cache_path.exists():
        return False  # already cached

    resolved = _resolve_path(local_path)
    if not os.path.isfile(resolved):
        return False  # source file missing

    try:
        result = subprocess.run(
            [
                "ffmpeg", "-y",
                "-i", resolved,
                "-vframes", "1",
                "-vf", f"scale={width}:-1",
                "-f", "image2",
                str(cache_path),
            ],
            capture_output=True,
            timeout=30,
        )
        if result.returncode != 0 or not cache_path.exists():
            return False
        return True
    except Exception as e:
        log.warning("Failed to generate video thumb for media %d: %s", media_id, e)
        if cache_path.exists():
            cache_path.unlink()
        return False


def main():
    """CLI entry point for batch thumbnail warmup."""
    parser = argparse.ArgumentParser(description="Pre-generate thumbnail cache")
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 4,
                        help="Number of parallel workers (default: CPU count)")
    parser.add_argument("--limit", type=int, default=0,
                        help="Max number of media to process (0 = all)")
    parser.add_argument("--type", dest="media_type", choices=["image", "video", "all"],
                        default="all", help="Which media type to process")
    parser.add_argument("--db-url", default=DB_URL,
                        help="Database connection URL")
    parser.add_argument("--cache-dir", type=Path, default=THUMB_CACHE_DIR,
                        help="Thumbnail cache directory")
    parser.add_argument("--width", type=int, default=THUMB_MAX_WIDTH,
                        help="Thumbnail max width")
    args = parser.parse_args()

    args.cache_dir.mkdir(exist_ok=True)

    # Build query based on media type filter
    if args.media_type == "all":
        type_filter = "file_type IN ('image', 'video')"
    else:
        type_filter = f"file_type = '{args.media_type}'"

    limit_clause = f"LIMIT {args.limit}" if args.limit > 0 else ""
    query = f"SELECT id, local_path, file_type FROM media_files WHERE {type_filter} ORDER BY id {limit_clause}"

    log.info("Connecting to database...")
    conn = psycopg2.connect(args.db_url)
    cur = conn.cursor()
    cur.execute(query)
    rows = cur.fetchall()
    cur.close()
    conn.close()

    total = len(rows)
    log.info("Found %d media files to process with %d workers", total, args.workers)

    generated = 0
    skipped = 0
    failed = 0

    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {}
        for media_id, local_path, file_type in rows:
            if file_type == "video":
                fut = executor.submit(
                    generate_video_thumb, media_id, local_path,
                    cache_dir=args.cache_dir, width=args.width,
                )
            else:
                fut = executor.submit(
                    generate_thumb, media_id, local_path,
                    cache_dir=args.cache_dir, width=args.width,
                )
            futures[fut] = (media_id, file_type)

        for i, fut in enumerate(as_completed(futures), 1):
            media_id, file_type = futures[fut]
            try:
                result = fut.result()
                if result:
                    generated += 1
                else:
                    skipped += 1
            except Exception as e:
                failed += 1
                log.error("Error processing media %d: %s", media_id, e)

            if i % 100 == 0 or i == total:
                log.info(
                    "Progress: %d/%d (generated=%d, skipped=%d, failed=%d)",
                    i, total, generated, skipped, failed,
                )

    log.info(
        "Done! Total=%d, Generated=%d, Skipped=%d, Failed=%d",
        total, generated, skipped, failed,
    )


if __name__ == "__main__":
    main()