neiam /dms

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

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

use lapin::{Connection, ConnectionProperties};

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::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, opts, FieldTable::default()).await {
            Ok(q) => (q.message_count(), q.consumer_count()),
            Err(_) => (0, 0),
        }
    }
    let mut out = Vec::new();
    for kind in [
        JobKind::Thumbnail,
        JobKind::Ocr,
        JobKind::Describe,
        JobKind::ExtractTags,
        JobKind::SimilaritySearch,
    ] {
        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)
}