neiam /dms

pub mod messages;
pub mod publisher;
pub mod consumer;

pub use messages::{JobKind, JobMessage, ALL_KINDS};
pub use publisher::{publish_job, Publisher};
pub use consumer::{run_worker, JobHandler, JobOutcome};

use lapin::{Connection, ConnectionProperties};
use lapin::options::{BasicAckOptions, BasicGetOptions, BasicPublishOptions};

pub const EXCHANGE: &str = "dms.jobs";
pub const DLX_EXCHANGE: &str = "dms.jobs.dlx";

pub async fn connect(amqp_url: &str) -> Result<Connection, lapin::Error> {
    Connection::connect(amqp_url, ConnectionProperties::default()).await
}

/// Routing key + queue name for a given job kind. Each kind has its own queue
/// so workers can be scaled per-kind. All queues bind to the same exchange.
pub fn routing_key(kind: JobKind) -> &'static str {
    match kind {
        JobKind::Poster => "job.poster",
        JobKind::Preview => "job.preview",
        JobKind::Thumbnail => "job.thumbnail",
        JobKind::Ocr => "job.ocr",
        JobKind::Describe => "job.describe",
        JobKind::ExtractTags => "job.extract_tags",
        JobKind::SimilaritySearch => "job.similarity",
    }
}

pub fn queue_name(kind: JobKind) -> String {
    format!("dms.{}", routing_key(kind).replace('.', "_"))
}

pub fn dlq_name(kind: JobKind) -> String {
    format!("dms.{}.dlq", routing_key(kind).replace('.', "_"))
}

/// Snapshot of a single RabbitMQ queue. Pulled via passive queue_declare so
/// the broker reports counts for an existing queue without redeclaring it
/// (a non-passive declare with mismatched args would error).
#[derive(Debug, Clone, serde::Serialize)]
pub struct QueueDepth {
    pub kind: &'static str,
    pub queue: String,
    pub messages: u32,
    pub consumers: u32,
    pub dlq: String,
    pub dlq_messages: u32,
}

/// Read message + consumer counts for every job-kind queue and its DLQ. A
/// missing queue (404 from the broker) is reported as zero rather than
/// failing the whole snapshot — broker may not have seen this kind yet on
/// a fresh deploy.
pub async fn queue_depths(conn: &Connection) -> Result<Vec<QueueDepth>, lapin::Error> {
    use lapin::options::QueueDeclareOptions;
    use lapin::types::FieldTable;

    let opts = QueueDeclareOptions { passive: true, ..Default::default() };
    // A passive declare against a non-existent queue raises NOT_FOUND and
    // closes the channel, so each probe opens a fresh channel. Cheap (no
    // server round-trip beyond the channel open) and keeps the loop going
    // when the broker hasn't seen a particular queue yet.
    async fn probe(conn: &Connection, name: &str, opts: QueueDeclareOptions) -> (u32, u32) {
        let Ok(ch) = conn.create_channel().await else { return (0, 0); };
        match ch.queue_declare(name.into(), opts, FieldTable::default()).await {
            Ok(q) => (q.message_count(), q.consumer_count()),
            Err(_) => (0, 0),
        }
    }
    let mut out = Vec::new();
    for kind in ALL_KINDS {
        let qname = queue_name(kind);
        let dlq = dlq_name(kind);
        let (msgs, consumers) = probe(conn, &qname, opts).await;
        let (dlq_msgs, _) = probe(conn, &dlq, opts).await;
        out.push(QueueDepth {
            kind: kind.as_str(),
            queue: qname,
            messages: msgs,
            consumers,
            dlq,
            dlq_messages: dlq_msgs,
        });
    }
    Ok(out)
}

/// Drain the DLQ for a single job kind back into the main per-kind queue.
/// Each message is fetched with `basic_get` and re-published to `EXCHANGE`
/// using the kind's routing key, preserving the original payload + props
/// (including the accumulated `x-death` header — useful for diagnostics).
///
/// Returns the number of messages moved. Stops when the DLQ is empty or when
/// a re-publish fails (the corresponding DLQ message is left un-acked so the
/// broker requeues it).
pub async fn requeue_dlq(conn: &Connection, kind: JobKind) -> Result<usize, lapin::Error> {
    let ch = conn.create_channel().await?;
    let dlq = dlq_name(kind);
    let rk = routing_key(kind);
    let mut moved = 0usize;
    loop {
        let got = ch
            .basic_get(dlq.as_str().into(), BasicGetOptions { no_ack: false })
            .await?;
        let Some(msg) = got else { break };
        ch.basic_publish(
            EXCHANGE.into(),
            rk.into(),
            BasicPublishOptions::default(),
            &msg.data,
            msg.properties.clone(),
        )
        .await?
        .await?;
        ch.basic_ack(msg.delivery_tag, BasicAckOptions::default()).await?;
        moved += 1;
    }
    Ok(moved)
}

/// Convenience: requeue every kind's DLQ. Returns the total moved across all
/// kinds. A failure on one kind aborts the run — the caller can retry.
pub async fn requeue_all_dlqs(conn: &Connection) -> Result<usize, lapin::Error> {
    let mut total = 0usize;
    for kind in [
        JobKind::Thumbnail,
        JobKind::Ocr,
        JobKind::Describe,
        JobKind::ExtractTags,
        JobKind::SimilaritySearch,
    ] {
        total += requeue_dlq(conn, kind).await?;
    }
    Ok(total)
}