neiam /dms

use async_trait::async_trait;
use futures_util::future::BoxFuture;
use sea_orm::*;
use tracing::{error, info, warn};

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

/// Pluggable backend for image description. Real implementations call out to
/// a VLM (Anthropic, OpenAI, Ollama, etc); the noop variant just records a
/// placeholder so the rest of the pipeline can be exercised in dev.
#[async_trait]
pub trait DescribeBackend: Send + Sync {
    async fn describe(&self, image_bytes: &[u8], content_type: Option<&str>) -> Result<String, String>;
    fn model_name(&self) -> &'static str;
}

pub struct NoopBackend;

#[async_trait]
impl DescribeBackend for NoopBackend {
    async fn describe(&self, _image_bytes: &[u8], _content_type: Option<&str>) -> Result<String, String> {
        Ok(String::new())
    }
    fn model_name(&self) -> &'static str {
        "noop"
    }
}

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::Describe).await, Ok(true)) {
            info!(upload_id = msg.upload_id, "describe already complete, skipping");
            return JobOutcome::Ack;
        }
        let _ = jobs::mark_processing(ctx.db.as_ref(), msg.upload_id, JobKind::Describe).await;

        // For now use the noop backend; swap in a real VLM client once chosen.
        let backend: Box<dyn DescribeBackend> = Box::new(NoopBackend);

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

        let abs = ctx.storage_abs(&upload.item_path);
        let bytes = match tokio::fs::read(&abs).await {
            Ok(b) => b,
            Err(e) => {
                error!(error = %e, "read original failed");
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Describe, &format!("read: {e}")).await;
                return JobOutcome::Reject;
            }
        };

        let description = match backend.describe(&bytes, upload.content_type.as_deref()).await {
            Ok(d) => d,
            Err(e) => {
                warn!(error = %e, "describe backend failed");
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Describe, &e).await;
                return JobOutcome::Reject;
            }
        };

        if !description.trim().is_empty() {
            let active = ai_descriptions::ActiveModel {
                upload_id: Set(upload.id),
                description: Set(description.clone()),
                model_used: Set(Some(backend.model_name().to_string())),
                confidence: Set(None),
                ..Default::default()
            };
            if let Err(e) = active.insert(ctx.db.as_ref()).await {
                error!(error = %e, "insert description failed");
                return JobOutcome::Retry;
            }

            let _ = jobs::upsert_pending(ctx.db.as_ref(), upload.id, JobKind::ExtractTags).await;
            let _ = jobs::mark_complete(ctx.db.as_ref(), upload.id, JobKind::Describe).await;

            let next = JobMessage::new(JobKind::ExtractTags, upload.id, upload.user_id);
            if let Err(e) = publish_job(&ctx.amqp, &next).await {
                warn!(error = %e, "failed to enqueue extract_tags job");
            }
            info!(upload_id = upload.id, "description recorded");
        } else {
            info!(upload_id = upload.id, "noop describe; marking complete with empty result");
            let _ = jobs::mark_complete(ctx.db.as_ref(), upload.id, JobKind::Describe).await;
        }
        JobOutcome::Ack
    })
}