neiam /dms

pub mod poster;
pub mod preview;
pub mod thumbnail;
pub mod ocr;
pub mod describe;
pub mod extract_tags;

use async_trait::async_trait;
use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter};
use std::sync::Arc;
use tracing::{error, warn};

use crate::entities::{upload_jobs, uploads};
use crate::jobs;
use crate::queue::{JobHandler, JobKind, JobMessage, JobOutcome};
use crate::storage::LocalStorage;
use crate::Config;

/// Shared worker context passed to every handler.
#[derive(Clone)]
pub struct WorkerContext {
    pub db: Arc<DatabaseConnection>,
    pub storage: Arc<LocalStorage>,
    pub config: Arc<Config>,
    pub amqp: Arc<lapin::Connection>,
}

impl WorkerContext {
    pub fn storage_abs(&self, relative: &str) -> std::path::PathBuf {
        use crate::storage::Storage;
        self.storage.absolute_path(relative)
    }
}

pub fn is_video(upload: &uploads::Model) -> bool {
    upload
        .content_type
        .as_deref()
        .map(|ct| ct.starts_with("video/"))
        .unwrap_or(false)
}

/// What a still-image categorizer should actually read for this upload: the
/// original for an image, the frame the `poster` job extracted for a video.
/// Videos only ever reach thumbnail/ocr/describe after that job has run, so
/// the frame is on disk by the time this path is used.
pub fn source_image_relative(upload: &uploads::Model) -> String {
    if is_video(upload) {
        crate::storage::derived_relative(&upload.item_path, poster::POSTER_SUFFIX, poster::POSTER_EXT)
    } else {
        upload.item_path.clone()
    }
}

/// Adapter so per-kind worker functions can plug into the generic JobHandler.
pub struct DispatchHandler {
    pub ctx: WorkerContext,
    pub run: fn(WorkerContext, JobMessage) -> futures_util::future::BoxFuture<'static, JobOutcome>,
}

#[async_trait]
impl JobHandler for DispatchHandler {
    async fn handle(&self, msg: &JobMessage) -> JobOutcome {
        // Hard cap on retries. The DB-tracked `upload_jobs.attempts` column
        // ticks up every time `mark_processing` runs, so a worker stuck in a
        // requeue loop can build up arbitrarily many redeliveries; without a
        // cap, the same poisoned message starves the queue forever.
        let max = self.ctx.config.worker_max_attempts;
        if max > 0 {
            let kind = msg.kind;
            match read_attempts(&self.ctx.db, msg.upload_id, kind).await {
                Ok(Some(prior)) if prior >= max => {
                    let err_msg = format!(
                        "exceeded WORKER_MAX_ATTEMPTS ({max}); giving up after {prior} attempts"
                    );
                    warn!(
                        upload_id = %msg.upload_id,
                        kind = ?kind,
                        attempts = prior,
                        max,
                        "max attempts exceeded, sending to DLQ"
                    );
                    if let Err(e) = jobs::mark_failed(self.ctx.db.as_ref(), msg.upload_id, kind, &err_msg).await {
                        error!(error = %e, "mark_failed during attempts-cap reject");
                    }
                    return JobOutcome::Reject;
                }
                Ok(_) => {}
                Err(e) => {
                    error!(error = %e, "read_attempts failed; allowing job through");
                }
            }
        }
        (self.run)(self.ctx.clone(), msg.clone()).await
    }
}

async fn read_attempts(
    db: &DatabaseConnection,
    upload_id: uuid::Uuid,
    kind: JobKind,
) -> Result<Option<i32>, sea_orm::DbErr> {
    let row = upload_jobs::Entity::find()
        .filter(upload_jobs::Column::UploadId.eq(upload_id))
        .filter(upload_jobs::Column::Kind.eq(kind.as_str()))
        .one(db)
        .await?;
    Ok(row.map(|r| r.attempts))
}