intercom-rs 1.1.1

//! Work queue support for JetStream.
//!
//! This module provides a convenient interface for creating and using work queues
//! with JetStream, including support for interest-based consumers.

use std::{
    marker::PhantomData,
    pin::Pin,
    task::{Context, Poll},
};

use futures::{Sink, Stream};
use pin_project_lite::pin_project;
use serde::{de::DeserializeOwned, Serialize};

use crate::{
    codec::CodecType,
    error::{Error, Result},
    jetstream::{
        consumer::{JetStreamMessage, PullBatch, PullConsumer, PullMessages},
        stream::{RetentionPolicy, Stream as JsStream, StreamBuilder},
    },
};

/// A typed work queue backed by JetStream with configurable codec.
///
/// Work queues provide at-least-once delivery with automatic acknowledgment tracking.
/// Messages are removed from the queue once acknowledged.
///
/// # Type Parameters
///
/// * `T` - The message type for this queue
/// * `C` - The codec type used for serialization
///
/// # Example
///
/// ```no_run
/// use intercom::{Client, MsgPackCodec, jetstream::queue::WorkQueue};
/// use serde::{Deserialize, Serialize};
/// use futures::StreamExt;
///
/// #[derive(Serialize, Deserialize, Debug)]
/// struct Job {
///     id: u64,
///     payload: String,
/// }
///
/// # async fn example() -> intercom::Result<()> {
/// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
/// let jetstream = client.jetstream();
///
/// // Create a work queue
/// let queue = WorkQueue::<Job, MsgPackCodec>::builder(&jetstream, "jobs")
///     .max_messages(10_000)
///     .create()
///     .await?;
///
/// // Push a job
/// queue.push(&Job { id: 1, payload: "do work".into() }).await?;
///
/// // Option 1: Use as a Stream (pulls one job at a time)
/// let mut queue = queue.into_stream().await?;
/// while let Some(result) = queue.next().await {
///     let job = result?;
///     println!("Processing: {:?}", job.payload);
///     job.ack().await?;
/// }
///
/// // Option 2: Pull a batch of jobs
/// // let mut jobs = queue.pull(10).await?;
/// // while let Some(result) = jobs.next().await {
/// //     let job = result?;
/// //     job.ack().await?;
/// // }
/// # Ok(())
/// # }
/// ```
pub struct WorkQueue<T, C: CodecType> {
    stream: JsStream<C>,
    consumer: PullConsumer<T, C>,
    subject: String,
    context: async_nats::jetstream::Context,
    _marker: PhantomData<(T, C)>,
}

impl<T, C: CodecType> WorkQueue<T, C> {
    /// Create a work queue builder.
    pub fn builder(
        context: &crate::jetstream::context::JetStreamContext<C>,
        name: &str,
    ) -> WorkQueueBuilder<T, C> {
        WorkQueueBuilder::new(context.inner().clone(), name.to_string())
    }

    /// Get the stream backing this queue.
    pub fn stream(&self) -> &JsStream<C> {
        &self.stream
    }

    /// Get the consumer for this queue.
    pub fn consumer(&self) -> &PullConsumer<T, C> {
        &self.consumer
    }

    /// Get the subject for this queue.
    pub fn subject(&self) -> &str {
        &self.subject
    }
}

impl<T: Serialize, C: CodecType> WorkQueue<T, C> {
    /// Push a message to the queue.
    pub async fn push(&self, message: &T) -> Result<u64> {
        let data = C::encode(message)?;
        let ack = self
            .context
            .publish(self.subject.clone(), data.into())
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?;
        Ok(ack.sequence)
    }

    /// Create a sink for pushing messages.
    pub fn sink(&self) -> WorkQueueSink<T, C> {
        WorkQueueSink::new(self.context.clone(), self.subject.clone())
    }
}

impl<T: DeserializeOwned, C: CodecType> WorkQueue<T, C> {
    /// Pull a batch of messages from the queue.
    pub async fn pull(&self, batch_size: usize) -> Result<PullBatch<T, C>> {
        self.consumer.fetch(batch_size).await
    }

    /// Get a continuous stream of messages.
    pub async fn messages(&self) -> Result<crate::jetstream::consumer::PullMessages<T, C>> {
        self.consumer.messages().await
    }

    /// Convert this queue into a [`Stream`] that continuously pulls messages one at a time.
    ///
    /// This is the most ergonomic way to process queue messages when you want to handle
    /// them one at a time in a loop.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec, jetstream::queue::WorkQueue};
    /// use serde::{Deserialize, Serialize};
    /// use futures::StreamExt;
    ///
    /// #[derive(Serialize, Deserialize, Debug)]
    /// struct Job { id: u64 }
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    /// let queue = WorkQueue::<Job, MsgPackCodec>::builder(&jetstream, "jobs").create().await?;
    ///
    /// // Convert to a Stream and iterate
    /// let mut queue = queue.into_stream().await?;
    /// while let Some(job) = queue.next().await {
    ///     let job = job?;
    ///     println!("Processing job: {:?}", job.payload);
    ///     job.ack().await?;
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn into_stream(self) -> Result<StreamingWorkQueue<T, C>> {
        let messages = self.consumer.messages().await?;
        Ok(StreamingWorkQueue {
            messages,
            context: self.context,
            subject: self.subject,
            _marker: PhantomData,
        })
    }
}

pin_project! {
    /// A work queue that implements [`Stream`] for continuous message processing.
    ///
    /// Created from [`WorkQueue::into_stream`]. This type allows using the queue
    /// directly with `StreamExt` methods like `next()`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec, jetstream::queue::WorkQueue};
    /// use serde::{Deserialize, Serialize};
    /// use futures::StreamExt;
    ///
    /// #[derive(Serialize, Deserialize, Debug)]
    /// struct Job { id: u64 }
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    /// let queue = WorkQueue::<Job, MsgPackCodec>::builder(&jetstream, "jobs").create().await?;
    ///
    /// let mut streaming = queue.into_stream().await?;
    /// while let Some(job) = streaming.next().await {
    ///     let job = job?;
    ///     job.ack().await?;
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub struct StreamingWorkQueue<T, C: CodecType> {
        #[pin]
        messages: PullMessages<T, C>,
        context: async_nats::jetstream::Context,
        subject: String,
        _marker: PhantomData<(T, C)>,
    }
}

impl<T: Serialize, C: CodecType> StreamingWorkQueue<T, C> {
    /// Push a message to the queue.
    ///
    /// Note: This is available even while streaming, allowing you to add work
    /// to the queue while processing existing messages.
    pub async fn push(&self, message: &T) -> Result<u64> {
        let data = C::encode(message)?;
        let ack = self
            .context
            .publish(self.subject.clone(), data.into())
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?;
        Ok(ack.sequence)
    }
}

impl<T: DeserializeOwned, C: CodecType> Stream for StreamingWorkQueue<T, C> {
    type Item = Result<JetStreamMessage<T>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.project().messages.poll_next(cx)
    }
}

/// Builder for work queues.
pub struct WorkQueueBuilder<T, C: CodecType> {
    context: async_nats::jetstream::Context,
    name: String,
    subject: Option<String>,
    max_messages: i64,
    max_bytes: i64,
    max_age: std::time::Duration,
    replicas: usize,
    _marker: PhantomData<(T, C)>,
}

impl<T, C: CodecType> WorkQueueBuilder<T, C> {
    /// Create a new work queue builder.
    fn new(context: async_nats::jetstream::Context, name: String) -> Self {
        Self {
            context,
            name,
            subject: None,
            max_messages: -1,
            max_bytes: -1,
            max_age: std::time::Duration::ZERO,
            replicas: 1,
            _marker: PhantomData,
        }
    }

    /// Set the subject for the queue.
    ///
    /// If not set, defaults to the queue name.
    pub fn subject(mut self, subject: impl Into<String>) -> Self {
        self.subject = Some(subject.into());
        self
    }

    /// Set the maximum number of messages.
    pub fn max_messages(mut self, max: i64) -> Self {
        self.max_messages = max;
        self
    }

    /// Set the maximum bytes.
    pub fn max_bytes(mut self, max: i64) -> Self {
        self.max_bytes = max;
        self
    }

    /// Set the maximum age for messages.
    pub fn max_age(mut self, age: std::time::Duration) -> Self {
        self.max_age = age;
        self
    }

    /// Set the number of replicas.
    pub fn replicas(mut self, replicas: usize) -> Self {
        self.replicas = replicas;
        self
    }

    /// Create the work queue.
    pub async fn create(self) -> Result<WorkQueue<T, C>> {
        let subject = self.subject.unwrap_or_else(|| self.name.clone());
        let consumer_name = format!("{}-worker", self.name);

        // Create the stream with work queue retention
        let stream_builder = StreamBuilder::<C>::new(self.context.clone(), self.name.clone())
            .subjects(vec![subject.clone()])
            .retention(RetentionPolicy::WorkQueue)
            .max_messages(self.max_messages)
            .max_bytes(self.max_bytes)
            .max_age(self.max_age)
            .replicas(self.replicas);

        let stream = stream_builder.create_or_update().await?;

        // Create a durable pull consumer
        let consumer = stream
            .pull_consumer_builder::<T>(&consumer_name)
            .durable()
            .create_or_update()
            .await?;

        Ok(WorkQueue {
            stream,
            consumer,
            subject,
            context: self.context,
            _marker: PhantomData,
        })
    }
}

/// A sink for pushing messages to a work queue.
pub struct WorkQueueSink<T, C: CodecType> {
    context: async_nats::jetstream::Context,
    subject: String,
    _marker: PhantomData<(T, C)>,
}

impl<T, C: CodecType> WorkQueueSink<T, C> {
    fn new(context: async_nats::jetstream::Context, subject: String) -> Self {
        Self {
            context,
            subject,
            _marker: PhantomData,
        }
    }
}

impl<T: Serialize, C: CodecType> WorkQueueSink<T, C> {
    /// Publish a message directly with error handling.
    ///
    /// Unlike the [`Sink`] implementation, this method returns any publish errors
    /// and waits for JetStream acknowledgment.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use intercom::{Client, MsgPackCodec, WorkQueue};
    /// use serde::{Deserialize, Serialize};
    ///
    /// #[derive(Serialize, Deserialize)]
    /// struct Job { id: u64 }
    ///
    /// # async fn example() -> intercom::Result<()> {
    /// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
    /// let jetstream = client.jetstream();
    /// let queue = WorkQueue::<Job, MsgPackCodec>::builder(&jetstream, "jobs").create().await?;
    /// let sink = queue.sink();
    ///
    /// // Direct publish with acknowledgment
    /// let seq = sink.publish(&Job { id: 1 }).await?;
    /// println!("Published at sequence: {}", seq);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn publish(&self, message: &T) -> Result<u64> {
        let data = C::encode(message)?;
        let ack = self
            .context
            .publish(self.subject.clone(), data.into())
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?;
        Ok(ack.sequence)
    }
}

impl<T, C: CodecType> Clone for WorkQueueSink<T, C> {
    fn clone(&self) -> Self {
        Self {
            context: self.context.clone(),
            subject: self.subject.clone(),
            _marker: PhantomData,
        }
    }
}

impl<T: Serialize, C: CodecType> Sink<T> for WorkQueueSink<T, C> {
    type Error = Error;

    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, item: T) -> Result<()> {
        let data = C::encode(&item)?;
        // Note: Uses fire-and-forget for Sink compatibility.
        // Use WorkQueueSink::publish() for error handling and acknowledgment.
        let context = self.context.clone();
        let subject = self.subject.clone();
        tokio::spawn(async move {
            let _ = context.publish(subject, data.into()).await;
        });
        Ok(())
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        self.poll_flush(cx)
    }
}

// ============================================================================
// Interest-Based Queue
// ============================================================================

/// A typed queue with interest-based retention and configurable codec.
///
/// Messages are retained until acknowledged by all consumers that were
/// active when the message was published.
///
/// # Type Parameters
///
/// * `T` - The message type for this queue
/// * `C` - The codec type used for serialization
///
/// # Example
///
/// ```no_run
/// use intercom::{Client, MsgPackCodec, jetstream::queue::InterestQueue};
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Serialize, Deserialize, Debug)]
/// struct Event {
///     id: u64,
///     data: String,
/// }
///
/// # async fn example() -> intercom::Result<()> {
/// let client = Client::<MsgPackCodec>::connect("nats://localhost:4222").await?;
/// let jetstream = client.jetstream();
///
/// // Create an interest-based queue
/// let queue = InterestQueue::<Event, MsgPackCodec>::builder(&jetstream, "events")
///     .subject("events.>")
///     .create()
///     .await?;
///
/// // Add consumers
/// let consumer1 = queue.add_consumer("service-a").await?;
/// let consumer2 = queue.add_consumer("service-b").await?;
///
/// // Both consumers must acknowledge for message removal
/// # Ok(())
/// # }
/// ```
pub struct InterestQueue<T, C: CodecType> {
    stream: JsStream<C>,
    subject: String,
    context: async_nats::jetstream::Context,
    _marker: PhantomData<(T, C)>,
}

impl<T, C: CodecType> InterestQueue<T, C> {
    /// Create an interest queue builder.
    pub fn builder(
        context: &crate::jetstream::context::JetStreamContext<C>,
        name: &str,
    ) -> InterestQueueBuilder<T, C> {
        InterestQueueBuilder::new(context.inner().clone(), name.to_string())
    }

    /// Get the stream backing this queue.
    pub fn stream(&self) -> &JsStream<C> {
        &self.stream
    }

    /// Get the subject for this queue.
    pub fn subject(&self) -> &str {
        &self.subject
    }

    /// Add a durable consumer to this queue.
    pub async fn add_consumer(&self, name: &str) -> Result<PullConsumer<T, C>> {
        self.stream
            .pull_consumer_builder::<T>(name)
            .durable()
            .create_or_update()
            .await
    }

    /// Add a consumer with a filter subject.
    pub async fn add_consumer_filtered(
        &self,
        name: &str,
        filter: &str,
    ) -> Result<PullConsumer<T, C>> {
        self.stream
            .pull_consumer_builder::<T>(name)
            .durable()
            .filter_subject(filter)
            .create_or_update()
            .await
    }
}

impl<T: Serialize, C: CodecType> InterestQueue<T, C> {
    /// Publish a message to the queue.
    pub async fn publish(&self, message: &T) -> Result<u64> {
        let data = C::encode(message)?;
        let ack = self
            .context
            .publish(self.subject.clone(), data.into())
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?;
        Ok(ack.sequence)
    }

    /// Publish a message to a specific subject within the queue's subject space.
    pub async fn publish_to(&self, subject: &str, message: &T) -> Result<u64> {
        let data = C::encode(message)?;
        let ack = self
            .context
            .publish(subject.to_string(), data.into())
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?
            .await
            .map_err(|e| Error::JetStream(e.to_string()))?;
        Ok(ack.sequence)
    }
}

/// Builder for interest-based queues.
pub struct InterestQueueBuilder<T, C: CodecType> {
    context: async_nats::jetstream::Context,
    name: String,
    subject: Option<String>,
    max_messages: i64,
    max_bytes: i64,
    max_age: std::time::Duration,
    replicas: usize,
    _marker: PhantomData<(T, C)>,
}

impl<T, C: CodecType> InterestQueueBuilder<T, C> {
    fn new(context: async_nats::jetstream::Context, name: String) -> Self {
        Self {
            context,
            name,
            subject: None,
            max_messages: -1,
            max_bytes: -1,
            max_age: std::time::Duration::ZERO,
            replicas: 1,
            _marker: PhantomData,
        }
    }

    /// Set the subject for the queue.
    pub fn subject(mut self, subject: impl Into<String>) -> Self {
        self.subject = Some(subject.into());
        self
    }

    /// Set the maximum number of messages.
    pub fn max_messages(mut self, max: i64) -> Self {
        self.max_messages = max;
        self
    }

    /// Set the maximum bytes.
    pub fn max_bytes(mut self, max: i64) -> Self {
        self.max_bytes = max;
        self
    }

    /// Set the maximum age for messages.
    pub fn max_age(mut self, age: std::time::Duration) -> Self {
        self.max_age = age;
        self
    }

    /// Set the number of replicas.
    pub fn replicas(mut self, replicas: usize) -> Self {
        self.replicas = replicas;
        self
    }

    /// Create the interest queue.
    pub async fn create(self) -> Result<InterestQueue<T, C>> {
        let subject = self.subject.unwrap_or_else(|| self.name.clone());

        // Create the stream with interest-based retention
        let stream_builder = StreamBuilder::<C>::new(self.context.clone(), self.name.clone())
            .subjects(vec![subject.clone()])
            .retention(RetentionPolicy::Interest)
            .max_messages(self.max_messages)
            .max_bytes(self.max_bytes)
            .max_age(self.max_age)
            .replicas(self.replicas);

        let stream = stream_builder.create_or_update().await?;

        Ok(InterestQueue {
            stream,
            subject,
            context: self.context,
            _marker: PhantomData,
        })
    }
}

// ============================================================================
// Queue Messages (alias for convenience)
// ============================================================================

/// A message from a work queue or interest queue.
///
/// This is an alias for [`JetStreamMessage`] for convenience.
pub type QueueMessage<T> = JetStreamMessage<T>;