#!/usr/bin/env python3

# This script creates and caches thumbnails of image/video files. 
# It sweeps the user's home folder only (excluding hidden files), checking the
# status of image and video files for cached thumbnails, if needed it creates 
# new thumbnails and adds them to the thumbnail cache in ~/.cache.
# Intended use is for PcmanFM, to assist in generating a proper thumbnail cache. 
# It can be set to run once at boot/login, and/or used with a menu entry to 
# create an updated cache file at any time.
 
# PcmanFM will continue to generate the cached thumbnails for images larger than 
# 4.2 mb on-the-fly per the default setting in Vuu-do, at 4.2 mb or higher the 
# thumbnail setting in PcmanFM works (mostly), any lower and you start losing 
# thumbnails of arbitrary sizes randomly for some reason, this script catches the 
# smaller stuff and anything else, albeit not on-the-fly which would be optimal, 
# but even just running this once at boot helps a LOT.

# Copyleft: greenjeans 2025, use as you see fit. https://sourceforge.net/projects/vuu-do/
# Version: 1.02 - adds support for .jxl format.
# This is free software with no warranty, use at your own risk.
# Depends: python3, python3-pil, the gdk-pixbuf, heif, and ffmpeg thumbnailers,
# libjxl and libjxl-gdk-pixbuf version 0.11.1 or higher.

import os
import subprocess
import hashlib
import urllib.parse
import mimetypes
import time
from pathlib import Path
import logging
from PIL import Image
import PIL.PngImagePlugin
from multiprocessing import Pool, cpu_count

# Set up logging. Logs/errors in ~/.xsession-errors.
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Configuration. Svg, webp, and jxl files are handled now by gdk-pixbuf-thumbnailer. 
# Image/heic must be defined here in addition to image/heif.
HOME = str(Path.home())
CACHE_DIR = os.path.join(HOME, '.cache', 'thumbnails', 'normal')
MIN_FILE_SIZE = 100 * 1024  # 100 KB in bytes, everything above this size gets thumbnailed/cached. (edit to change)
THUMBNAIL_SIZE = 128  # Normal size per XDG spec. (edit to change)
THUMBNAILERS = {
    'image/png': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'image/jpeg': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'image/gif': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'image/jxl': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'image/heif': ['heif-thumbnailer', '{input}', '{output}'],
    'image/heic': ['heif-thumbnailer', '{input}', '{output}'],
    'image/avif': ['heif-thumbnailer', '{input}', '{output}'],
    'image/webp': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'image/svg+xml': ['gdk-pixbuf-thumbnailer', '-s', str(THUMBNAIL_SIZE), '{input}', '{output}'],
    'video/mp4': ['ffmpegthumbnailer', '-i', '{input}', '-o', '{output}', '-s', str(THUMBNAIL_SIZE), '-t', '10'],
    'video/mpeg': ['ffmpegthumbnailer', '-i', '{input}', '-o', '{output}', '-s', str(THUMBNAIL_SIZE), '-t', '10'],
    'video/x-matroska': ['ffmpegthumbnailer', '-i', '{input}', '-o', '{output}', '-s', str(THUMBNAIL_SIZE), '-t', '10'],
}

# EXTENSIONS must be defined below in lowercase for case-insensitive matching.
EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.heic', '.heif', '.avif', '.webp', '.svg', '.mp4', '.mpeg', '.mpg', '.mkv', '.jxl'}

def ensure_cache_dir():
    """Ensure the thumbnail cache directory exists."""
    os.makedirs(CACHE_DIR, exist_ok=True)
    logger.info(f"Thumbnail cache directory: {CACHE_DIR}")

def get_thumbnail_path(file_path):
    """Generate the thumbnail filename per XDG spec (MD5 of file URI)."""
    file_uri = f"file://{urllib.parse.quote(file_path)}"
    thumb_hash = hashlib.md5(file_uri.encode('utf-8')).hexdigest()
    return os.path.join(CACHE_DIR, f"{thumb_hash}.png")

def add_png_metadata(thumb_path, file_path, mtime):
    """Add XDG-required metadata to the PNG thumbnail."""
    try:
        img = Image.open(thumb_path)
        metadata = PIL.PngImagePlugin.PngInfo()
        file_uri = f"file://{urllib.parse.quote(file_path)}"
        metadata.add_text("Thumb::URI", file_uri)
        metadata.add_text("Thumb::MTime", str(int(mtime)))
        img.save(thumb_path, "PNG", pnginfo=metadata)
        logger.debug(f"Added metadata to {thumb_path}")
    except Exception as e:
        logger.error(f"Failed to add metadata to {thumb_path}: {e}")

def generate_thumbnail(file_path, thumb_path):
    """Generate a thumbnail using the appropriate thumbnailer."""
    mime_type, _ = mimetypes.guess_type(file_path)
    if not mime_type or mime_type not in THUMBNAILERS:
        logger.warning(f"No thumbnailer for MIME type {mime_type or 'unknown'} ({file_path})")
        return False
    
    cmd = [arg.format(input=file_path, output=thumb_path) for arg in THUMBNAILERS[mime_type]]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        logger.debug(f"Generated thumbnail for {file_path}")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Thumbnailer failed for {file_path}: {e}")
        return False

def needs_update(file_path, thumb_path):
    """Check if the thumbnail needs to be updated."""
    if not os.path.exists(thumb_path):
        return True
    file_mtime = os.path.getmtime(file_path)
    try:
        img = Image.open(thumb_path)
        metadata_mtime = img.info.get('Thumb::MTime')
        if metadata_mtime and int(metadata_mtime) == int(file_mtime):
            return False
        return True
    except Exception as e:
        logger.warning(f"Invalid thumbnail {thumb_path}: {e}")
        return True

def generate_thumbnail_worker(args):
    """Worker function for multiprocessing to generate a thumbnail for a single file."""
    file_path, thumb_path = args
    try:
        if generate_thumbnail(file_path, thumb_path):
            add_png_metadata(thumb_path, file_path, os.path.getmtime(file_path))
            logger.info(f"Generated thumbnail for {file_path}")
            return True
        else:
            logger.warning(f"Failed to generate thumbnail for {file_path}")
            return False
    except Exception as e:
        logger.error(f"Error processing {file_path}: {e}")
        return False

def scan_and_cache():
    """Scan home directory and generate/update thumbnails using multiprocessing."""
    ensure_cache_dir()
    file_count = 0
    thumb_count = 0
    tasks = []

    # Collect all files that need thumbnailing, excluding hidden files, symlinks, and unreadables.
    for root, dirs, files in os.walk(HOME, topdown=True, followlinks=False):
        dirs[:] = [d for d in dirs if not d.startswith('.')]
        for fname in files:
            if not any(fname.lower().endswith(ext) for ext in EXTENSIONS):
                continue
            file_path = os.path.join(root, fname)
            try:
                if os.path.getsize(file_path) < MIN_FILE_SIZE:
                    logger.debug(f"Skipping {file_path} (size < {MIN_FILE_SIZE} bytes)")
                    continue
                thumb_path = get_thumbnail_path(file_path)
                file_count += 1
                if needs_update(file_path, thumb_path):
                    tasks.append((file_path, thumb_path))
                else:
                    logger.debug(f"Thumbnail up-to-date for {file_path}")
            except OSError as e:
                logger.warning(f"Skipping {file_path}: {e}")
                continue

    # Process tasks in parallel, default setting is to use up to 4 processes, this can be changed if needed.
    if tasks:
        num_processes = min(cpu_count(), 4)
        logger.info(f"Generating {len(tasks)} thumbnails using {num_processes} processes")
        pool = Pool(processes=num_processes)
        try:
             results = pool.map(generate_thumbnail_worker, tasks)
        finally:
             pool.close()
             pool.join()
        thumb_count = sum(1 for result in results if result)

    logger.info(f"Processed {file_count} files, generated/updated {thumb_count} thumbnails")

def main():
    """Main function to run the thumbnail caching process."""
    try:
        scan_and_cache()
    except Exception as e:
        logger.error(f"Script failed: {e}")
        raise

if __name__ == "__main__":
    main()