neiam /dms

use futures_util::future::BoxFuture;
use ocrs::{ImageSource, OcrEngine, OcrEngineParams};
use rten::Model;
use sea_orm::*;
use std::sync::Arc;
use tokio::sync::OnceCell;
use tracing::{error, info, warn};

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

/// One engine per process, lazily initialised. ocrs's OcrEngine is heavy
/// (loads ML models) so we never want to construct more than one.
static ENGINE: OnceCell<Arc<OcrEngine>> = OnceCell::const_new();

async fn engine() -> Result<Arc<OcrEngine>, String> {
    ENGINE
        .get_or_try_init(|| async {
            let det = std::env::var("OCRS_DETECTION_MODEL")
                .map_err(|_| "OCRS_DETECTION_MODEL not set".to_string())?;
            let rec = std::env::var("OCRS_RECOGNITION_MODEL")
                .map_err(|_| "OCRS_RECOGNITION_MODEL not set".to_string())?;
            let det_model = Model::load_file(&det).map_err(|e| format!("load detection model: {e}"))?;
            let rec_model = Model::load_file(&rec).map_err(|e| format!("load recognition model: {e}"))?;
            let engine = OcrEngine::new(OcrEngineParams {
                detection_model: Some(det_model),
                recognition_model: Some(rec_model),
                ..Default::default()
            })
            .map_err(|e| format!("OcrEngine init: {e}"))?;
            info!(detection = %det, recognition = %rec, "OCR engine ready");
            Ok::<_, String>(Arc::new(engine))
        })
        .await
        .cloned()
}

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

        let engine = match engine().await {
            Ok(e) => e,
            Err(e) => {
                warn!(error = %e, "ocr engine unavailable, skipping job");
                let _ = jobs::mark_failed(ctx.db.as_ref(), msg.upload_id, JobKind::Ocr, &format!("engine unavailable: {e}")).await;
                return JobOutcome::Reject;
            }
        };

        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::Ocr, "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::Ocr, &format!("read: {e}")).await;
                return JobOutcome::Reject;
            }
        };

        let result = tokio::task::spawn_blocking(move || -> Result<String, String> {
            let img = image::load_from_memory(&bytes)
                .map_err(|e| format!("decode image: {e}"))?
                .into_rgb8();
            let (w, h) = img.dimensions();
            let source = ImageSource::from_bytes(img.as_raw(), (w, h))
                .map_err(|e| format!("ImageSource: {e}"))?;
            let prepared = engine
                .prepare_input(source)
                .map_err(|e| format!("prepare_input: {e}"))?;
            engine.get_text(&prepared).map_err(|e| format!("get_text: {e}"))
        })
        .await;

        let text = match result {
            Ok(Ok(t)) => t,
            Ok(Err(e)) => {
                error!(error = %e, "ocr failed");
                let _ = jobs::mark_failed(ctx.db.as_ref(), upload.id, JobKind::Ocr, &e).await;
                return JobOutcome::Reject;
            }
            Err(e) => {
                error!(error = %e, "ocr blocking task panicked");
                return JobOutcome::Retry;
            }
        };

        let active = ocr_results::ActiveModel {
            upload_id: Set(upload.id),
            extracted_text: Set(Some(text.clone())),
            confidence: Set(None),
            processing_engine: Set(Some("ocrs".to_string())),
            ..Default::default()
        };
        if let Err(e) = active.insert(ctx.db.as_ref()).await {
            error!(error = %e, "insert ocr row failed");
            return JobOutcome::Retry;
        }

        // Mark complete + insert pending row for the follow-up extract_tags
        // job BEFORE publishing it. Doing this first prevents the aggregate
        // recompute from briefly seeing "all jobs complete" while extract_tags
        // is in flight.
        let _ = jobs::upsert_pending(ctx.db.as_ref(), upload.id, JobKind::ExtractTags).await;
        let _ = jobs::mark_complete(ctx.db.as_ref(), upload.id, JobKind::Ocr).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, chars = text.len(), "ocr complete");
        JobOutcome::Ack
    })
}