EventSnap/backend/src/services/compression.rs

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

use anyhow::{Context, Result};
use sqlx::PgPool;
use tokio::sync::{broadcast, Semaphore};
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>,
}

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,
        }
    }

    /// 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();
        tokio::spawn(async move {
            let _permit = worker.semaphore.acquire().await;
            match worker.do_process(upload_id, &original_path, &mime_type).await {
                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}: {e:#}");
                    // Auto-cleanup: a failed transcode would otherwise leave a
                    // permanently broken feed card, silently charge the uploader's
                    // quota, and orphan the original on disk. Refund + soft-delete
                    // (one tx, so v_feed excludes it), remove the orphan file, then
                    // tell the uploader (upload-error toast) and evict the card
                    // everywhere (upload-deleted, already handled by the feed).
                    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");
                    }
                    let orphan = worker.media_path.join(&original_path);
                    if let Err(rm) = tokio::fs::remove_file(&orphan).await {
                        tracing::warn!(error = ?rm, path = %orphan.display(), "failed to remove orphaned original");
                    }
                    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 = self.generate_image_preview(upload_id, &original, mime_type).await?;
            Upload::set_preview_path(&self.pool, upload_id, &preview_rel).await?;
            tracing::info!("preview 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(())
    }

    async fn generate_image_preview(
        &self,
        upload_id: Uuid,
        original: &Path,
        mime_type: &str,
    ) -> Result<String> {
        let previews_dir = self.media_path.join("previews");
        tokio::fs::create_dir_all(&previews_dir).await?;

        let preview_filename = format!("{upload_id}.jpg");
        let preview_path = previews_dir.join(&preview_filename);
        let original = original.to_path_buf();
        let preview_path_clone = preview_path.clone();
        let mime_owned = mime_type.to_string();

        // Run blocking image operations in a spawn_blocking task
        tokio::task::spawn_blocking(move || -> Result<()> {
            // Reject decompression bombs *before* fully decoding: the upload body
            // cap bounds the file size on disk, but a small file can still decode to
            // enormous dimensions (e.g. a ~1 MB image expanding to 50k×50k px →
            // gigabytes), OOM-ing the box during decode/resize. 12000×12000 covers
            // any real phone photo; max_alloc hard-caps the decode allocation.
            let mut reader = image::ImageReader::open(&original)
                .context("failed to open image")?
                .with_guessed_format()
                .context("failed to read image header")?;
            let mut limits = image::Limits::default();
            limits.max_image_width = Some(12_000);
            limits.max_image_height = Some(12_000);
            limits.max_alloc = Some(256 * 1024 * 1024);
            reader.limits(limits);
            let img = reader.decode().context("failed to decode image")?;

            // Resize to max 800px wide, preserving aspect ratio
            let preview = img.resize(800, 800, image::imageops::FilterType::Lanczos3);
            preview.save_with_format(&preview_path_clone, image::ImageFormat::Jpeg)
                .context("failed to save preview")?;

            // 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/{preview_filename}"))
    }

    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}"))
    }
}