neiam /dms

use clap::Parser;
use std::sync::Arc;
use tracing::info;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

use dms::queue::{run_worker, JobKind};
use dms::storage::LocalStorage;
use dms::workers::{describe, extract_tags, ocr, poster, thumbnail, DispatchHandler, WorkerContext};
use dms::{config::Config, create_db_connection};

#[derive(Parser, Debug)]
#[command(name = "queue-worker", about = "DMS background categorizer worker")]
struct Cli {
    /// Which queue to consume: poster, thumbnail, ocr, describe, extract_tags.
    /// Also sourceable via the `WORKER_KIND` env var so compose can drive
    /// every worker container off the same image.
    #[arg(long, env = "WORKER_KIND")]
    kind: String,

    /// Override prefetch count. Falls back to `WORKER_PREFETCH` env / 4.
    #[arg(long, env = "WORKER_PREFETCH")]
    prefetch: Option<u16>,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "dms=debug,queue_worker=debug".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let cli = Cli::parse();
    let kind = JobKind::from_str(&cli.kind)
        .ok_or_else(|| anyhow::anyhow!("unknown --kind: {}", cli.kind))?;

    let config = Arc::new(Config::from_env()?);
    let db = Arc::new(create_db_connection(&config.database_url).await?);
    let storage = Arc::new(LocalStorage::new(&config.storage_root));
    let amqp = Arc::new(dms::queue::connect(&config.amqp_url).await?);
    let prefetch = cli.prefetch.unwrap_or(config.worker_prefetch);

    info!(kind = ?kind, prefetch, "queue-worker starting");

    let ctx = WorkerContext {
        db,
        storage,
        config: config.clone(),
        amqp: amqp.clone(),
    };

    let run_fn: fn(WorkerContext, dms::queue::JobMessage) -> futures_util::future::BoxFuture<'static, dms::queue::JobOutcome> = match kind {
        JobKind::Poster => poster::run,
        JobKind::Thumbnail => thumbnail::run,
        JobKind::Ocr => ocr::run,
        JobKind::Describe => describe::run,
        JobKind::ExtractTags => extract_tags::run,
        JobKind::SimilaritySearch => {
            return Err(anyhow::anyhow!("similarity_search worker not yet implemented"));
        }
    };

    let handler = DispatchHandler { ctx, run: run_fn };
    run_worker(amqp.as_ref(), kind, prefetch, handler).await?;
    Ok(())
}