EventSnap/backend/src/services/compression.rs

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use anyhow::{Context, Result};
use sqlx::PgPool;
use tokio::sync::{Semaphore, broadcast};
use uuid::Uuid;

use crate::models::upload::Upload;
use crate::state::SseEvent;

#[derive(Clone)]
pub struct CompressionWorker {
    semaphore: Arc<Semaphore>,
    pool: PgPool,
    media_path: PathBuf,
    sse_tx: broadcast::Sender<SseEvent>,
    /// Bumped whenever the underlying data is reset out from under in-flight work (only the e2e
    /// TRUNCATE does this today). A task captures the value at spawn and abandons itself if it has
    /// changed by the time it runs — see `process`.
    generation: Arc<AtomicU64>,
}

impl CompressionWorker {
    pub fn new(
        pool: PgPool,
        media_path: PathBuf,
        concurrency: usize,
        sse_tx: broadcast::Sender<SseEvent>,
    ) -> Self {
        Self {
            semaphore: Arc::new(Semaphore::new(concurrency)),
            pool,
            media_path,
            sse_tx,
            generation: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Invalidate all in-flight and queued compression work. Called by the e2e TRUNCATE endpoint:
    /// truncating deletes the upload rows and wipes `media/`, so a worker that was queued on the
    /// semaphore when the wipe happened would otherwise wake in the NEXT test, fail to find its
    /// file, and broadcast `upload-error` / `upload-deleted` into that test's live SSE stream —
    /// corrupting any test that asserts on toasts or feed contents. Bumping the generation makes
    /// those stale tasks return silently instead. A no-op in production (never called there).
    pub fn bump_generation(&self) {
        self.generation.fetch_add(1, Ordering::SeqCst);
    }

    /// How many times `do_process` is attempted before an upload is given up on. The
    /// give-up path is user-visible (the photo disappears), so transient infrastructure
    /// errors must not reach it.
    const MAX_PROCESS_ATTEMPTS: u32 = 3;

    /// Revision of the image-derivative pipeline. Bump this whenever a change makes existing
    /// previews/displays wrong, so `backfill_stale_derivatives` regenerates them once on the
    /// next start. Rev 1 = EXIF orientation is applied.
    const DERIVATIVES_REV: i16 = 1;

    /// Spawn a background task to process an uploaded file.
    pub fn process(&self, upload_id: Uuid, original_path: String, mime_type: String) {
        let worker = self.clone();
        let born_at = worker.generation.load(Ordering::SeqCst);
        tokio::spawn(async move {
            let _permit = worker.semaphore.acquire().await;
            // The data this task was queued against may have been reset while it waited for a permit
            // (e2e TRUNCATE). If so, its file and row are gone; doing anything — including
            // broadcasting a failure — would leak into an unrelated test. Abandon quietly.
            if worker.generation.load(Ordering::SeqCst) != born_at {
                return;
            }
            // Retry before giving up. Most failures here are transient and self-clearing —
            // an ENOSPC spike while several guests upload at once, a momentary DB-pool
            // exhaustion, a panic inside the image codec — and the give-up path is
            // user-visible data loss, so it is worth a few seconds to avoid entering it.
            let mut attempt = 1u32;
            let outcome = loop {
                match worker
                    .do_process(upload_id, &original_path, &mime_type)
                    .await
                {
                    Ok(v) => break Ok(v),
                    Err(e) if attempt < Self::MAX_PROCESS_ATTEMPTS => {
                        tracing::warn!(
                            error = ?e, %upload_id, attempt,
                            "compression attempt failed; retrying"
                        );
                        tokio::time::sleep(std::time::Duration::from_secs(2u64.pow(attempt)))
                            .await;
                        attempt += 1;
                        // The data may have been reset while we slept (e2e TRUNCATE).
                        if worker.generation.load(Ordering::SeqCst) != born_at {
                            return;
                        }
                    }
                    Err(e) => break Err(e),
                }
            };

            match outcome {
                Ok(_) => {
                    tracing::info!("compression completed for upload {upload_id}");
                    let _ = worker.sse_tx.send(SseEvent {
                        event_type: "upload-processed".to_string(),
                        data: serde_json::json!({ "upload_id": upload_id }).to_string(),
                    });
                }
                Err(e) => {
                    tracing::error!(
                        "compression failed for upload {upload_id} after {attempt} attempt(s): {e:#}"
                    );
                    // Refund + soft-delete (one tx, so v_feed excludes it) so a failed
                    // transcode doesn't leave a permanently broken feed card or silently
                    // charge the uploader's quota. Then tell the uploader (upload-error
                    // toast) and evict the card everywhere (upload-deleted).
                    //
                    // The ORIGINAL IS DELIBERATELY KEPT. This path used to `remove_file` it
                    // unconditionally, which meant any transient error — a disk-full blip
                    // while saving a derivative, a pool hiccup, a panic in the image codec —
                    // irreversibly destroyed the guest's only copy of a photo they can never
                    // retake. The row is only soft-deleted, so keeping the bytes makes the
                    // upload fully recoverable; the file is orphaned rather than lost, and
                    // the path is logged so it can be found. `backfill_stale_derivatives`
                    // already refuses to destroy data on error for exactly this reason.
                    let _ = Upload::set_compression_status(&worker.pool, upload_id, "failed").await;
                    if let Err(del) = Upload::soft_delete(&worker.pool, upload_id).await {
                        tracing::warn!(error = ?del, %upload_id, "failed to soft-delete after compression failure");
                    }
                    tracing::warn!(
                        %upload_id,
                        path = %worker.media_path.join(&original_path).display(),
                        "original retained for recovery after compression failure"
                    );
                    let _ = worker.sse_tx.send(SseEvent {
                        event_type: "upload-error".to_string(),
                        data: serde_json::json!({ "upload_id": upload_id, "error": e.to_string() })
                            .to_string(),
                    });
                    let _ = worker.sse_tx.send(SseEvent {
                        event_type: "upload-deleted".to_string(),
                        data: serde_json::json!({ "upload_id": upload_id }).to_string(),
                    });
                }
            }
        });
    }

    async fn do_process(
        &self,
        upload_id: Uuid,
        original_path: &str,
        mime_type: &str,
    ) -> Result<()> {
        Upload::set_compression_status(&self.pool, upload_id, "processing").await?;

        let original = self.media_path.join(original_path);

        if mime_type.starts_with("image/") {
            let (preview_rel, display_rel) = self
                .generate_image_derivatives(upload_id, &original, mime_type)
                .await?;
            Upload::set_preview_path(&self.pool, upload_id, &preview_rel).await?;
            Upload::set_display_path(&self.pool, upload_id, &display_rel).await?;
            Upload::set_derivatives_rev(&self.pool, upload_id, Self::DERIVATIVES_REV).await?;
            tracing::info!("preview + display generated for upload {upload_id}");
        } else if mime_type.starts_with("video/") {
            let thumb_rel = self.generate_video_thumbnail(upload_id, &original).await?;
            Upload::set_thumbnail_path(&self.pool, upload_id, &thumb_rel).await?;
            tracing::info!("thumbnail generated for upload {upload_id}");
        }

        Upload::set_compression_status(&self.pool, upload_id, "done").await?;
        Ok(())
    }

    /// Longest edge of the big-screen "display" derivative used by the diashow. Sized to be
    /// sharp on 1080p/4K while staying bounded (a ~2048px JPEG decodes to ~16 MB — trivial
    /// for any kiosk, unlike a raw multi-thousand-pixel original).
    const DISPLAY_MAX_EDGE: u32 = 2048;
    /// Longest edge of the phone-feed "preview" (data-saver default).
    const PREVIEW_MAX_EDGE: u32 = 800;

    /// Decode the image ONCE and emit both derivatives — the 800px `preview` (phone feed)
    /// and the 2048px `display` (diashow). Returns `(preview_rel, display_rel)`.
    async fn generate_image_derivatives(
        &self,
        upload_id: Uuid,
        original: &Path,
        mime_type: &str,
    ) -> Result<(String, String)> {
        let previews_dir = self.media_path.join("previews");
        let displays_dir = self.media_path.join("displays");
        tokio::fs::create_dir_all(&previews_dir).await?;
        tokio::fs::create_dir_all(&displays_dir).await?;

        let filename = format!("{upload_id}.jpg");
        let preview_path = previews_dir.join(&filename);
        let display_path = displays_dir.join(&filename);
        let original = original.to_path_buf();
        let mime_owned = mime_type.to_string();
        let preview_max = Self::PREVIEW_MAX_EDGE;
        let display_max = Self::DISPLAY_MAX_EDGE;

        // Run blocking image operations in a spawn_blocking task
        tokio::task::spawn_blocking(move || -> Result<()> {
            // Decompression-bomb limits + EXIF orientation, both in one place — see
            // services::imaging for why neither may be skipped.
            let img = crate::services::imaging::decode_oriented(&original)?;

            // Preview: max 800px, preserving aspect ratio (data-saver feed).
            img.resize(
                preview_max,
                preview_max,
                image::imageops::FilterType::Lanczos3,
            )
            .save_with_format(&preview_path, image::ImageFormat::Jpeg)
            .context("failed to save preview")?;

            // Display: max 2048px for the diashow. Only DOWNSCALE — never upscale a smaller
            // original (that adds bytes with no quality gain); re-encode it as JPEG as-is.
            let display = if img.width() > display_max || img.height() > display_max {
                img.resize(
                    display_max,
                    display_max,
                    image::imageops::FilterType::Lanczos3,
                )
            } else {
                img
            };
            display
                .save_with_format(&display_path, image::ImageFormat::Jpeg)
                .context("failed to save display")?;

            // If the original is PNG, try lossless compression in-place
            if mime_owned == "image/png" {
                let opts = oxipng::Options::from_preset(2);
                let _ = oxipng::optimize(
                    &oxipng::InFile::Path(original),
                    &oxipng::OutFile::Path {
                        path: None,
                        preserve_attrs: true,
                    },
                    &opts,
                );
            }

            Ok(())
        })
        .await??;

        Ok((
            format!("previews/{filename}"),
            format!("displays/{filename}"),
        ))
    }

    /// Regenerate image derivatives that an older pipeline produced. Fire-and-forget from
    /// startup; picks up two cases, both of which leave the ORIGINAL untouched:
    ///
    /// - uploads processed before the `display` derivative existed (preview but no
    ///   `display_path`), and
    /// - uploads whose derivatives predate `DERIVATIVES_REV` — currently rev 1, which applies
    ///   the EXIF orientation. Everything generated before it is stored sideways for any
    ///   portrait phone photo.
    ///
    /// Unlike the failure path in `process`, a backfill error is logged and skipped — it must
    /// NEVER destroy or soft-delete an upload that already has a working preview.
    pub async fn backfill_stale_derivatives(&self) {
        let rows = sqlx::query_as::<_, (Uuid, String, String)>(
            "SELECT id, original_path, mime_type FROM upload
             WHERE deleted_at IS NULL AND mime_type LIKE 'image/%'
               AND original_path IS NOT NULL
               AND (
                     (display_path IS NULL AND preview_path IS NOT NULL)
                  OR derivatives_rev < $1
               )",
        )
        .bind(Self::DERIVATIVES_REV)
        .fetch_all(&self.pool)
        .await;
        let rows = match rows {
            Ok(r) => r,
            Err(e) => {
                tracing::warn!(error = ?e, "derivative backfill query failed");
                return;
            }
        };
        if rows.is_empty() {
            return;
        }
        tracing::info!("regenerating derivatives for {} upload(s)", rows.len());
        for (id, original_path, mime_type) in rows {
            let worker = self.clone();
            tokio::spawn(async move {
                let _permit = worker.semaphore.acquire().await;
                let original = worker.media_path.join(&original_path);
                match worker
                    .generate_image_derivatives(id, &original, &mime_type)
                    .await
                {
                    Ok((preview_rel, display_rel)) => {
                        let _ = Upload::set_preview_path(&worker.pool, id, &preview_rel).await;
                        let _ = Upload::set_display_path(&worker.pool, id, &display_rel).await;
                        let _ = Upload::set_derivatives_rev(
                            &worker.pool,
                            id,
                            Self::DERIVATIVES_REV,
                        )
                        .await;
                        tracing::info!("derivatives regenerated for upload {id}");
                    }
                    Err(e) => {
                        // Leave the existing derivatives and the original intact; this row is
                        // simply retried on the next start. The rev stays behind, which is the
                        // marker that it still needs doing.
                        tracing::warn!(error = ?e, %id, "derivative backfill failed; leaving as-is");
                    }
                }
            });
        }
    }

    async fn generate_video_thumbnail(&self, upload_id: Uuid, original: &Path) -> Result<String> {
        let thumbs_dir = self.media_path.join("thumbnails");
        tokio::fs::create_dir_all(&thumbs_dir).await?;

        let thumb_filename = format!("{upload_id}.jpg");
        let thumb_path = thumbs_dir.join(&thumb_filename);

        // Hard timeout — a malformed video can hang `ffmpeg` indefinitely. Without a
        // cap, the held compression-worker semaphore permit is never released and the
        // pool eventually deadlocks (no further uploads ever processed). 120s is well
        // above the time to extract one frame from any sane input.
        let mut child = tokio::process::Command::new("ffmpeg")
            .args([
                "-i",
                original.to_str().unwrap_or_default(),
                "-vframes",
                "1",
                "-ss",
                "00:00:01",
                "-vf",
                "scale=800:-1",
                "-y",
                thumb_path.to_str().unwrap_or_default(),
            ])
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true)
            .spawn()
            .context("failed to spawn ffmpeg")?;

        let status =
            match tokio::time::timeout(std::time::Duration::from_secs(120), child.wait()).await {
                Ok(res) => res.context("ffmpeg wait failed")?,
                Err(_) => {
                    let _ = child.kill().await;
                    anyhow::bail!("ffmpeg timeout after 120s");
                }
            };

        if !status.success() {
            // Best-effort: drain stderr for the log.
            let mut stderr = Vec::new();
            if let Some(mut handle) = child.stderr.take() {
                use tokio::io::AsyncReadExt;
                let _ = handle.read_to_end(&mut stderr).await;
            }
            anyhow::bail!("ffmpeg failed: {}", String::from_utf8_lossy(&stderr));
        }

        Ok(format!("thumbnails/{thumb_filename}"))
    }
}