neiam /dms

//! Poster-frame extraction for video uploads.
//!
//! Thumbnail, OCR and describe all work on a decoded still, and none of them
//! can open an `.mp4`. A video upload therefore runs this job first: shell out
//! to ffmpeg for one representative frame, store it beside the original, then
//! chain the usual three against that frame. From their point of view a video
//! becomes an image, which is the whole trick.

use futures_util::future::BoxFuture;
use sea_orm::*;
use std::process::Stdio;
use std::time::Duration;
use tracing::{error, info, warn};

use crate::entities::uploads;
use crate::jobs;
use crate::queue::{publish_job, JobKind, JobMessage, JobOutcome};
use crate::storage::Storage;
use crate::workers::WorkerContext;

/// Where the extracted frame lands, as `<stem>.poster.jpg` beside the original.
/// [`crate::workers::source_image_relative`] resolves the same path when a
/// categorizer needs to read it back.
pub const POSTER_SUFFIX: &str = "poster";
pub const POSTER_EXT: &str = "jpg";

/// The jobs that run against the frame once it exists.
const FOLLOW_UPS: [JobKind; 3] = [JobKind::Thumbnail, JobKind::Ocr, JobKind::Describe];

/// Why a frame didn't come out — the two cases want opposite outcomes.
enum PosterError {
    /// ffmpeg could not be run at all. Environmental, not the file's fault, so
    /// the job is worth retrying once the image is fixed.
    Spawn(String),
    /// ffmpeg ran and rejected the input. Retrying re-reads the same bytes and
    /// fails the same way, so this one is terminal.
    Failed(String),
}

pub fn run(ctx: WorkerContext, msg: JobMessage) -> BoxFuture<'static, JobOutcome> {
    Box::pin(async move {
        if matches!(
            jobs::is_complete(ctx.db.as_ref(), msg.upload_id, JobKind::Poster).await,
            Ok(true)
        ) {
            info!(upload_id = %msg.upload_id, "poster already complete, skipping");
            return JobOutcome::Ack;
        }
        let _ = jobs::mark_processing(ctx.db.as_ref(), msg.upload_id, JobKind::Poster).await;

        let upload = match uploads::Entity::find_by_id(msg.upload_id)
            .one(ctx.db.as_ref())
            .await
        {
            Ok(Some(u)) => u,
            Ok(None) => {
                warn!(upload_id = %msg.upload_id, "upload not found, dropping");
                let _ = jobs::mark_failed(ctx.db.as_ref(), msg.upload_id, JobKind::Poster, "upload not found").await;
                return JobOutcome::Reject;
            }
            Err(e) => {
                error!(error = %e, "db error loading upload");
                return JobOutcome::Retry;
            }
        };

        let abs = ctx.storage_abs(&upload.item_path);
        let frame = match extract_frame(&ctx, &abs).await {
            Ok(bytes) => bytes,
            Err(PosterError::Spawn(e)) => {
                error!(upload_id = %upload.id, error = %e, "ffmpeg unavailable");
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Poster, &e).await;
                return JobOutcome::Retry;
            }
            Err(PosterError::Failed(e)) => {
                error!(upload_id = %upload.id, error = %e, "poster extraction failed");
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Poster, &e).await;
                return JobOutcome::Reject;
            }
        };

        let stored = match ctx
            .storage
            .put_derived_as(&upload.item_path, POSTER_SUFFIX, POSTER_EXT, &frame)
            .await
        {
            Ok(rel) => rel,
            Err(e) => {
                error!(error = %e, "write poster failed");
                return JobOutcome::Retry;
            }
        };

        // Create the follow-up rows BEFORE completing this one. The recompute
        // that runs on every terminal transition looks at the whole set, so
        // completing the poster while it is the only row would briefly flip the
        // upload to "complete" before the real work has even been queued.
        for kind in FOLLOW_UPS {
            if let Err(e) = jobs::upsert_pending(ctx.db.as_ref(), upload.id, kind).await {
                error!(error = %e, ?kind, "upsert follow-up job row failed");
            }
        }
        let _ = jobs::mark_complete(ctx.db.as_ref(), upload.id, JobKind::Poster).await;

        for kind in FOLLOW_UPS {
            let next = JobMessage::new(kind, upload.id, upload.user_id);
            if let Err(e) = publish_job(&ctx.amqp, &next).await {
                error!(error = %e, ?kind, "publish follow-up failed");
            }
        }

        info!(upload_id = %upload.id, poster = %stored, "poster frame extracted");
        JobOutcome::Ack
    })
}

/// One representative frame, as JPEG bytes.
async fn extract_frame(ctx: &WorkerContext, path: &std::path::Path) -> Result<Vec<u8>, PosterError> {
    let bin = &ctx.config.ffmpeg_bin;
    let mut cmd = tokio::process::Command::new(bin);
    cmd.arg("-v")
        .arg("error")
        .arg("-i")
        .arg(path)
        // `thumbnail` scores frames across a window and picks the most
        // representative, rather than taking frame 0 — which on a GIF-turned-mp4
        // is routinely a blank fade-in.
        .arg("-vf")
        .arg("thumbnail")
        .arg("-frames:v")
        .arg("1")
        // Straight to stdout: a single JPEG needs no temp file.
        .arg("-f")
        .arg("image2pipe")
        .arg("-vcodec")
        .arg("mjpeg")
        .arg("-")
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());

    let timeout = Duration::from_secs(ctx.config.poster_timeout_secs);
    let output = match tokio::time::timeout(timeout, cmd.output()).await {
        Ok(Ok(o)) => o,
        Ok(Err(e)) => {
            return Err(PosterError::Spawn(format!("could not run {bin}: {e}")));
        }
        Err(_) => {
            // A hung decode is the file's problem, not the environment's.
            return Err(PosterError::Failed(format!(
                "ffmpeg timed out after {}s",
                ctx.config.poster_timeout_secs
            )));
        }
    };

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(PosterError::Failed(format!(
            "ffmpeg exited {}: {}",
            output.status,
            stderr.trim().chars().take(400).collect::<String>()
        )));
    }
    if output.stdout.is_empty() {
        return Err(PosterError::Failed(
            "ffmpeg produced no frame — no decodable video stream?".to_string(),
        ));
    }
    Ok(output.stdout)
}