ImageApi/src/thumbnails.rs

//! Thumbnail generation + the media-count Prometheus gauges.
//!
//! Startup and per-tick scans walk each library and produce a 200×200
//! thumbnail under `THUMBNAILS/<library_id>/<rel_path>`, falling through
//! a fast path (`image` crate), a RAW-preview path (`exif::extract_embedded_jpeg_preview`),
//! and ffmpeg for video / HEIF / NEF / ARW. Files that fail every
//! decoder get a sibling `.unsupported` sentinel so subsequent scans
//! skip them silently.

use std::path::{Path, PathBuf};
use std::process::Command;

use image::GenericImageView;
use image::codecs::jpeg::JpegEncoder;
use lazy_static::lazy_static;
use log::{debug, error, info, warn};
use opentelemetry::{
    KeyValue,
    trace::{Span, TraceContextExt, Tracer},
};
use prometheus::IntGauge;
use rayon::prelude::*;
use walkdir::DirEntry;

use crate::content_hash;
use crate::exif;
use crate::file_types;
use crate::libraries;
use crate::otel::global_tracer;
use crate::video::actors::{generate_image_thumbnail_ffmpeg, generate_video_thumbnail};

/// Maximum long-edge size (px) for the large preview tier. Tuned to look
/// crisp full-screen on a 3× phone (≈1290×2796 native) and to hold up
/// through a few stops of pinch-zoom before the original streams in.
/// Bigger doesn't help: callers that need true full resolution request
/// `size=full` and the handler streams the original bytes.
pub const LARGE_PREVIEW_MAX_DIM: u32 = 2048;

/// JPEG quality for the large and xlarge preview tiers. 85 is the
/// conventional "indistinguishable from source at viewing size" point —
/// well above the `image` crate's default ~75, but well below quality-90+
/// territory where file size doubles for no perceptible win.
const LARGE_PREVIEW_JPEG_QUALITY: u8 = 85;

/// Maximum long-edge size (px) for the xlarge preview tier. Bridges the
/// gap between `large` (2048px, ~16MB decoded) and the original bytes
/// (potentially 48+ MP / ~192MB decoded). At 4096px the decoded bitmap is
/// ~64MB — enough for 2-3× pinch-zoom on any phone before the viewer
/// needs to stream the true original.
pub const XLARGE_PREVIEW_MAX_DIM: u32 = 4096;

lazy_static! {
    pub static ref IMAGE_GAUGE: IntGauge = IntGauge::new(
        "imageserver_image_total",
        "Count of the images on the server"
    )
    .unwrap();
    pub static ref VIDEO_GAUGE: IntGauge = IntGauge::new(
        "imageserver_video_total",
        "Count of the videos on the server"
    )
    .unwrap();
}

/// Sentinel path written next to a would-be thumbnail when a file cannot be
/// decoded by either the `image` crate or ffmpeg. Its presence causes future
/// scans to skip the file instead of re-logging the failure.
pub fn unsupported_thumbnail_sentinel(thumb_path: &Path) -> PathBuf {
    let mut s = thumb_path.as_os_str().to_owned();
    s.push(".unsupported");
    PathBuf::from(s)
}

pub fn generate_image_thumbnail(src: &Path, thumb_path: &Path) -> std::io::Result<()> {
    // The `image` crate doesn't auto-apply EXIF Orientation on load, and
    // saving back out as JPEG drops EXIF entirely — so without baking the
    // rotation into the pixels here, browsers see the raw landscape buffer
    // of a portrait phone shot and render it sideways. Read once up front
    // and apply to whichever decode branch we end up taking.
    let orientation = exif::read_orientation(src).unwrap_or(1);

    // RAW formats (ARW/NEF/CR2/etc): try the file's embedded JPEG preview
    // first. Avoids ffmpeg choking on proprietary RAW compression (Sony ARW
    // in particular), and is faster than decoding RAW pixels anyway.
    if let Some(preview) = exif::extract_embedded_jpeg_preview(src) {
        let img = image::load_from_memory(&preview).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("decode embedded preview {:?}: {}", src, e),
            )
        })?;
        let img = exif::apply_orientation(img, orientation);
        let scaled = img.thumbnail(200, u32::MAX);
        scaled
            .save_with_format(thumb_path, image::ImageFormat::Jpeg)
            .map_err(|e| std::io::Error::other(format!("save {:?}: {}", thumb_path, e)))?;
        return Ok(());
    }

    if file_types::needs_ffmpeg_thumbnail(src) {
        return generate_image_thumbnail_ffmpeg(src, thumb_path, orientation);
    }

    let img = image::open(src).map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{:?}: {}", src, e))
    })?;
    let img = exif::apply_orientation(img, orientation);
    let scaled = img.thumbnail(200, u32::MAX);
    scaled
        .save(thumb_path)
        .map_err(|e| std::io::Error::other(format!("save {:?}: {}", thumb_path, e)))?;
    Ok(())
}

/// Generate the on-demand large-preview tier (≈2048 long edge JPEG).
///
/// Mirrors [`generate_image_thumbnail`]'s decode waterfall — embedded RAW
/// preview, then ffmpeg for HEIC/HEIF, then the `image` crate — but
/// resizes to [`LARGE_PREVIEW_MAX_DIM`] instead of 200 and encodes at
/// quality 85 rather than the crate default. Caller is expected to have
/// already created the destination's parent dir.
///
/// Does not upscale: if the source's long edge is already below the cap,
/// the file is encoded at its native size (still re-saved as JPEG so the
/// served bytes match for callers that key off `Content-Length`).
pub fn generate_large_preview(src: &Path, dest: &Path) -> std::io::Result<()> {
    let orientation = exif::read_orientation(src).unwrap_or(1);

    // RAW: prefer the in-file embedded JPEG preview over raw-sensor decode.
    // The preview is typically already 1–2 MP and avoids RAW codec quirks.
    if let Some(preview) = exif::extract_embedded_jpeg_preview(src) {
        let img = image::load_from_memory(&preview).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("decode embedded preview {:?}: {}", src, e),
            )
        })?;
        let img = exif::apply_orientation(img, orientation);
        return encode_large_jpeg(img, dest);
    }

    if file_types::needs_ffmpeg_thumbnail(src) {
        return generate_large_preview_ffmpeg(src, dest, orientation);
    }

    let img = image::open(src).map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{:?}: {}", src, e))
    })?;
    let img = exif::apply_orientation(img, orientation);
    encode_large_jpeg(img, dest)
}

/// Resize-if-needed + JPEG-encode at q85. Used by both the embedded-preview
/// and image-crate-decode branches of `generate_large_preview`.
fn encode_large_jpeg(img: image::DynamicImage, dest: &Path) -> std::io::Result<()> {
    let (w, h) = img.dimensions();
    let max_dim = w.max(h);
    // Avoid upscaling tiny sources — pointless work and adds nothing for
    // the viewer. `thumbnail` would scale up freely; explicit guard.
    let scaled = if max_dim > LARGE_PREVIEW_MAX_DIM {
        img.thumbnail(LARGE_PREVIEW_MAX_DIM, LARGE_PREVIEW_MAX_DIM)
    } else {
        img
    };
    let file = std::fs::File::create(dest)
        .map_err(|e| std::io::Error::other(format!("create {:?}: {}", dest, e)))?;
    let mut writer = std::io::BufWriter::new(file);
    let mut encoder = JpegEncoder::new_with_quality(&mut writer, LARGE_PREVIEW_JPEG_QUALITY);
    encoder
        .encode_image(&scaled)
        .map_err(|e| std::io::Error::other(format!("encode {:?}: {}", dest, e)))?;
    Ok(())
}

/// ffmpeg path for HEIC/HEIF (image crate can't decode these). Mirrors
/// [`crate::video::actors::generate_image_thumbnail_ffmpeg`] but scales
/// to the large-preview cap instead of 200.
fn generate_large_preview_ffmpeg(src: &Path, dest: &Path, orientation: i32) -> std::io::Result<()> {
    // Rotation + scale + colorspace. HEIC sources use Display P3; without
    // colorspace=bt709 the mjpeg encoder treats P3 values as sRGB, producing
    // warm/oversaturated output. The min(iw,cap) trick caps the long edge
    // regardless of orientation, mirroring image::thumbnail.
    let rotation = match orientation {
        2 => "hflip",
        3 => "transpose=2",
        4 => "vflip",
        5 => "transpose=0,hflip",
        6 => "transpose=0",
        7 => "transpose=1,hflip",
        8 => "transpose=1",
        _ => "",
    };
    let scale_expr = format!(
        "scale='if(gt(iw,ih),min(iw,{cap}),-1)':'if(gt(iw,ih),-1,min(ih,{cap}))'",
        cap = LARGE_PREVIEW_MAX_DIM
    );
    let vf = if rotation.is_empty() {
        format!("{},colorspace=bt709", scale_expr)
    } else {
        format!("{},{},colorspace=bt709", rotation, scale_expr)
    };
    let output = Command::new("ffmpeg")
        .arg("-y")
        .arg("-i")
        .arg(src)
        .arg("-vframes")
        .arg("1")
        .arg("-vf")
        .arg(&vf)
        .arg("-q:v")
        // ffmpeg's mjpeg qscale: 2 ≈ ~q95, 5 ≈ ~q85, 10 ≈ ~q70. We pick
        // 5 to match the non-ffmpeg branch's q85 target.
        .arg("5")
        .arg("-f")
        .arg("image2")
        .arg("-c:v")
        .arg("mjpeg")
        .arg(dest)
        .output()?;

    if !output.status.success() {
        return Err(std::io::Error::other(format!(
            "ffmpeg failed ({}): {}",
            output.status,
            String::from_utf8_lossy(&output.stderr).trim()
        )));
    }
    Ok(())
}

/// Generate the on-demand xlarge-preview tier (≈4096 long edge JPEG).
///
/// Same waterfall as [`generate_large_preview`] but targeting
/// [`XLARGE_PREVIEW_MAX_DIM`]. Sources whose long edge is already below
/// the cap are encoded at native size (no upscale).
pub fn generate_xlarge_preview(src: &Path, dest: &Path) -> std::io::Result<()> {
    let orientation = exif::read_orientation(src).unwrap_or(1);

    if let Some(preview) = exif::extract_embedded_jpeg_preview(src) {
        let img = image::load_from_memory(&preview).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("decode embedded preview {:?}: {}", src, e),
            )
        })?;
        let img = exif::apply_orientation(img, orientation);
        return encode_xlarge_jpeg(img, dest);
    }

    if file_types::needs_ffmpeg_thumbnail(src) {
        return generate_xlarge_preview_ffmpeg(src, dest, orientation);
    }

    let img = image::open(src).map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{:?}: {}", src, e))
    })?;
    let img = exif::apply_orientation(img, orientation);
    encode_xlarge_jpeg(img, dest)
}

fn encode_xlarge_jpeg(img: image::DynamicImage, dest: &Path) -> std::io::Result<()> {
    let (w, h) = img.dimensions();
    let max_dim = w.max(h);
    let scaled = if max_dim > XLARGE_PREVIEW_MAX_DIM {
        img.thumbnail(XLARGE_PREVIEW_MAX_DIM, XLARGE_PREVIEW_MAX_DIM)
    } else {
        img
    };
    let file = std::fs::File::create(dest)
        .map_err(|e| std::io::Error::other(format!("create {:?}: {}", dest, e)))?;
    let mut writer = std::io::BufWriter::new(file);
    let mut encoder = JpegEncoder::new_with_quality(&mut writer, LARGE_PREVIEW_JPEG_QUALITY);
    encoder
        .encode_image(&scaled)
        .map_err(|e| std::io::Error::other(format!("encode {:?}: {}", dest, e)))?;
    Ok(())
}

fn generate_xlarge_preview_ffmpeg(
    src: &Path,
    dest: &Path,
    orientation: i32,
) -> std::io::Result<()> {
    let rotation = match orientation {
        2 => "hflip",
        3 => "transpose=2",
        4 => "vflip",
        5 => "transpose=0,hflip",
        6 => "transpose=0",
        7 => "transpose=1,hflip",
        8 => "transpose=1",
        _ => "",
    };
    let scale_expr = format!(
        "scale='if(gt(iw,ih),min(iw,{cap}),-1)':'if(gt(iw,ih),-1,min(ih,{cap}))'",
        cap = XLARGE_PREVIEW_MAX_DIM
    );
    let vf = if rotation.is_empty() {
        format!("{},colorspace=bt709", scale_expr)
    } else {
        format!("{},{},colorspace=bt709", rotation, scale_expr)
    };
    let output = Command::new("ffmpeg")
        .arg("-y")
        .arg("-i")
        .arg(src)
        .arg("-vframes")
        .arg("1")
        .arg("-vf")
        .arg(&vf)
        .arg("-q:v")
        .arg("5")
        .arg("-f")
        .arg("image2")
        .arg("-c:v")
        .arg("mjpeg")
        .arg(dest)
        .output()?;

    if !output.status.success() {
        return Err(std::io::Error::other(format!(
            "ffmpeg failed ({}): {}",
            output.status,
            String::from_utf8_lossy(&output.stderr).trim()
        )));
    }
    Ok(())
}

pub fn create_thumbnails(libs: &[libraries::Library], excluded_dirs: &[String]) {
    let tracer = global_tracer();
    let span = tracer.start("creating thumbnails");

    let thumbs = &dotenv::var("THUMBNAILS").expect("THUMBNAILS not defined");
    let thumbnail_directory: &Path = Path::new(thumbs);

    for lib in libs {
        info!(
            "Scanning thumbnails for library '{}' at {}",
            lib.name, lib.root_path
        );
        let images = PathBuf::from(&lib.root_path);
        // Effective excludes = global env-var excludes ∪ library row's
        // excluded_dirs. Lets a parent-library mount skip the subtree
        // already covered by a child library.
        let effective_excludes = lib.effective_excluded_dirs(excluded_dirs);

        // Prune EXCLUDED_DIRS so we don't generate thumbnails-of-thumbnails
        // for Synology @eaDir trees. file_scan handles filter_entry pruning.
        crate::file_scan::walk_library_files(&images, &effective_excludes)
            .into_par_iter()
            .for_each(|entry| {
                let src = entry.path();
                let Ok(relative_path) = src.strip_prefix(&images) else {
                    return;
                };
                // Library-scoped legacy path: prevents two libraries with
                // the same rel_path from clobbering each other's thumbs.
                // Hash-keyed promotion happens lazily on first hash-aware
                // request — keeping this loop ExifDao-free preserves the
                // current "cargo build && go" startup story.
                let thumb_path = content_hash::library_scoped_legacy_path(
                    thumbnail_directory,
                    lib.id,
                    relative_path,
                );
                let bare_legacy = thumbnail_directory.join(relative_path);

                // Backwards-compat check: if a single-library install has a
                // bare-legacy thumb here already, accept it as present.
                // Same for the sentinel. Means we don't redo work after
                // upgrade and we don't leave stale duplicates around.
                if thumb_path.exists()
                    || bare_legacy.exists()
                    || unsupported_thumbnail_sentinel(&thumb_path).exists()
                    || unsupported_thumbnail_sentinel(&bare_legacy).exists()
                {
                    return;
                }

                let Some(parent) = thumb_path.parent() else {
                    return;
                };
                if let Err(e) = std::fs::create_dir_all(parent) {
                    error!("Failed to create thumbnail dir {:?}: {}", parent, e);
                    return;
                }

                if is_video(&entry) {
                    let mut video_span = tracer.start_with_context(
                        "generate_video_thumbnail",
                        &opentelemetry::Context::new()
                            .with_remote_span_context(span.span_context().clone()),
                    );
                    video_span.set_attributes(vec![
                        KeyValue::new("type", "video"),
                        KeyValue::new("file-name", thumb_path.display().to_string()),
                        KeyValue::new("library", lib.name.clone()),
                    ]);

                    debug!("Generating video thumbnail: {:?}", thumb_path);
                    if let Err(e) = generate_video_thumbnail(src, &thumb_path) {
                        let sentinel = unsupported_thumbnail_sentinel(&thumb_path);
                        error!(
                            "Unable to thumbnail video {:?}: {}. Writing sentinel {:?}",
                            src, e, sentinel
                        );
                        if let Err(se) = std::fs::write(&sentinel, b"") {
                            warn!("Failed to write sentinel {:?}: {}", sentinel, se);
                        }
                    }
                    video_span.end();
                } else if is_image(&entry) {
                    match generate_image_thumbnail(src, &thumb_path) {
                        Ok(_) => info!("Saved thumbnail: {:?}", thumb_path),
                        Err(e) => {
                            let sentinel = unsupported_thumbnail_sentinel(&thumb_path);
                            error!(
                                "Unable to thumbnail {:?}: {}. Writing sentinel {:?}",
                                src, e, sentinel
                            );
                            if let Err(se) = std::fs::write(&sentinel, b"") {
                                warn!("Failed to write sentinel {:?}: {}", sentinel, se);
                            }
                        }
                    }
                }
            });
    }

    debug!("Finished making thumbnails");

    for lib in libs {
        let effective_excludes = lib.effective_excluded_dirs(excluded_dirs);
        update_media_counts(Path::new(&lib.root_path), &effective_excludes);
    }
}

pub fn update_media_counts(media_dir: &Path, excluded_dirs: &[String]) {
    let mut image_count = 0;
    let mut video_count = 0;
    for entry in crate::file_scan::walk_library_files(media_dir, excluded_dirs) {
        if is_image(&entry) {
            image_count += 1;
        } else if is_video(&entry) {
            video_count += 1;
        }
    }

    IMAGE_GAUGE.set(image_count);
    VIDEO_GAUGE.set(video_count);
}

pub fn is_image(entry: &DirEntry) -> bool {
    file_types::direntry_is_image(entry)
}

pub fn is_video(entry: &DirEntry) -> bool {
    file_types::direntry_is_video(entry)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    #[test]
    fn unsupported_thumbnail_sentinel_appends_suffix() {
        let p = Path::new("/thumbs/lib1/photo.jpg");
        let s = unsupported_thumbnail_sentinel(p);
        assert_eq!(s, PathBuf::from("/thumbs/lib1/photo.jpg.unsupported"));
    }

    #[test]
    fn unsupported_thumbnail_sentinel_preserves_extension_so_existing_thumb_is_distinct() {
        // A future scan checks both `thumb.exists()` and
        // `sentinel.exists()` — they must be distinct paths.
        let p = Path::new("foo.jpeg");
        let s = unsupported_thumbnail_sentinel(p);
        assert_ne!(s, PathBuf::from("foo.jpeg"));
        assert!(s.to_string_lossy().ends_with(".unsupported"));
    }

    #[test]
    fn unsupported_thumbnail_sentinel_handles_paths_without_extension() {
        let p = Path::new("/thumbs/notes");
        let s = unsupported_thumbnail_sentinel(p);
        assert_eq!(s, PathBuf::from("/thumbs/notes.unsupported"));
    }

    /// Smoke-test update_media_counts: build a tempdir with two images
    /// and one video, run the walker, and assert the gauges line up.
    /// Exercises the is_image / is_video classifier on real DirEntry
    /// values without needing a Prometheus registry.
    #[test]
    fn update_media_counts_counts_images_and_videos_in_tempdir() {
        let tmp = TempDir::new().expect("tempdir");
        fs::write(tmp.path().join("a.jpg"), b"").unwrap();
        fs::write(tmp.path().join("b.png"), b"").unwrap();
        fs::write(tmp.path().join("c.mp4"), b"").unwrap();
        fs::write(tmp.path().join("notes.txt"), b"").unwrap();
        // Reset gauges first in case another test mutated them — the
        // gauges are process-global statics.
        IMAGE_GAUGE.set(0);
        VIDEO_GAUGE.set(0);

        update_media_counts(tmp.path(), &[]);

        assert_eq!(IMAGE_GAUGE.get(), 2, "jpg + png");
        assert_eq!(VIDEO_GAUGE.get(), 1, "mp4");
    }
}