neiam /dms

use async_trait::async_trait;
use futures_util::future::BoxFuture;
use sea_orm::*;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::OnceCell;
use tracing::{error, info, warn};

use crate::config::{Config, DescribeBackendKind};
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, rknn-llm); 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>,
        filename: Option<&str>,
    ) -> Result<String, String>;
    fn model_name(&self) -> String;
}

pub struct NoopBackend;

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

/// HTTP client for the rknn-llm multimodal server (see
/// `tools/rknn-describe-server/` for the reference implementation).
///
/// Wire protocol — POST `endpoint`, multipart/form-data:
///   - `image`     binary file part
///   - `prompt`    optional text part (overrides server default)
///   - `model`     optional text part (selects loaded model on the server)
///   - `max_new_tokens`  optional integer text part
///
/// Response — JSON: `{ "description": "...", "model": "qwen2-vl-2b", "tokens": 123 }`
pub struct RknnBackend {
    client: reqwest::Client,
    endpoint: String,
    model: String,
    prompt: String,
    max_new_tokens: u32,
}

impl RknnBackend {
    pub fn new(cfg: &crate::config::DescribeConfig) -> Result<Self, String> {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(cfg.timeout_secs))
            .build()
            .map_err(|e| format!("build http client: {e}"))?;
        Ok(RknnBackend {
            client,
            endpoint: cfg.endpoint.clone(),
            model: cfg.model.clone(),
            prompt: cfg.prompt.clone(),
            max_new_tokens: cfg.max_new_tokens,
        })
    }
}

#[derive(Serialize, Deserialize)]
struct DescribeResponse {
    description: String,
    #[serde(default)]
    model: Option<String>,
}

#[async_trait]
impl DescribeBackend for RknnBackend {
    async fn describe(
        &self,
        image_bytes: &[u8],
        content_type: Option<&str>,
        filename: Option<&str>,
    ) -> Result<String, String> {
        let mut part = reqwest::multipart::Part::bytes(image_bytes.to_vec())
            .file_name(filename.unwrap_or("upload.bin").to_string());
        if let Some(ct) = content_type {
            part = part.mime_str(ct).map_err(|e| format!("invalid content-type: {e}"))?;
        }
        let form = reqwest::multipart::Form::new()
            .part("image", part)
            .text("prompt", self.prompt.clone())
            .text("model", self.model.clone())
            .text("max_new_tokens", self.max_new_tokens.to_string());

        let resp = self
            .client
            .post(&self.endpoint)
            .multipart(form)
            .send()
            .await
            .map_err(|e| format!("rknn-llm POST {}: {e}", self.endpoint))?;
        let status = resp.status();
        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(format!("rknn-llm http {status}: {body}"));
        }
        let parsed: DescribeResponse = resp
            .json()
            .await
            .map_err(|e| format!("decode rknn-llm response: {e}"))?;
        Ok(parsed.description)
    }
    fn model_name(&self) -> String {
        format!("rknn-llm:{}", self.model)
    }
}

static BACKEND: OnceCell<Arc<dyn DescribeBackend>> = OnceCell::const_new();

async fn backend_for(cfg: &Config) -> Arc<dyn DescribeBackend> {
    BACKEND
        .get_or_init(|| async {
            match cfg.describe.backend {
                DescribeBackendKind::Rknn => match RknnBackend::new(&cfg.describe) {
                    Ok(b) => {
                        info!(
                            endpoint = %cfg.describe.endpoint,
                            model = %cfg.describe.model,
                            "describe backend: rknn-llm"
                        );
                        Arc::new(b) as Arc<dyn DescribeBackend>
                    }
                    Err(e) => {
                        error!(error = %e, "rknn backend init failed; falling back to noop");
                        Arc::new(NoopBackend) as Arc<dyn DescribeBackend>
                    }
                },
                DescribeBackendKind::Noop => {
                    info!("describe backend: noop");
                    Arc::new(NoopBackend) as Arc<dyn DescribeBackend>
                }
            }
        })
        .await
        .clone()
}

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;

        let backend = backend_for(ctx.config.as_ref()).await;

        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 filename = std::path::Path::new(&upload.item_path)
            .file_name()
            .and_then(|s| s.to_str());

        let description = match backend
            .describe(&bytes, upload.content_type.as_deref(), filename)
            .await
        {
            Ok(d) => d,
            Err(e) => {
                warn!(error = %e, "describe backend failed");
                // Network/transport errors are retryable; persistent ones get
                // surfaced via the dead-letter queue when retries exhaust.
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Describe, &e).await;
                return JobOutcome::Retry;
            }
        };

        if !description.trim().is_empty() {
            let active = ai_descriptions::ActiveModel {
                id: Set(uuid::Uuid::now_v7()),
                upload_id: Set(upload.id),
                description: Set(description.clone()),
                model_used: Set(Some(backend.model_name())),
                confidence: Set(None),
                processed_at: Set(chrono::Utc::now().into()),
            };
            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, "empty describe result; marking complete");
            let _ = jobs::mark_complete(ctx.db.as_ref(), upload.id, JobKind::Describe).await;
        }
        JobOutcome::Ack
    })
}