rabbit_warren 0.1.1

use std::{future::Future, sync::Arc};

use futures_util::stream::StreamExt;
use lapin::{
    message::Delivery,
    options::{
        BasicAckOptions, BasicConsumeOptions, BasicNackOptions, BasicPublishOptions,
        BasicQosOptions, ConfirmSelectOptions, ExchangeDeclareOptions, QueueBindOptions,
        QueueDeclareOptions,
    },
    publisher_confirm::Confirmation,
    types::FieldTable,
    BasicProperties, Channel, Connection, ConnectionProperties, ExchangeKind,
};
use tokio::task::JoinHandle;
use tracing::{debug, error, warn, Instrument};

use crate::errors::{AckAction, Result};
use crate::errors::{IntoAckAction, ProcessingError};

/// Thin wrapper around a RabbitMQ channel that groups together the most common
/// operations used across services: connection establishment, topology
/// declaration (exchange/queue/bind), publishing with publisher confirms and
/// consuming with automatic ack/nack handling.
#[derive(Clone)]
pub struct RmqClient {
    channel: Arc<Channel>,
}

impl RmqClient {
    /// Establishes a connection to the broker using default connection
    /// properties, creates a channel and configures:
    /// - publisher confirms (needed to handle the `mandatory` publish flag);
    /// - basic QoS with the given prefetch count.
    pub async fn connect(address: &str, prefetch_count: u16) -> Result<Self> {
        Self::connect_with_properties(address, prefetch_count, ConnectionProperties::default())
            .await
    }

    /// Same as [`Self::connect`], but allows custom [`ConnectionProperties`]
    /// (e.g. to override the executor or the heartbeat interval).
    pub async fn connect_with_properties(
        address: &str,
        prefetch_count: u16,
        connection_properties: ConnectionProperties,
    ) -> Result<Self> {
        let connection = Connection::connect(address, connection_properties).await?;
        let channel = connection.create_channel().await?;

        // Enables `mandatory` flag publish handling.
        // Without confirm_select it will always be `NotRequested`.
        channel
            .confirm_select(ConfirmSelectOptions::default())
            .await?;

        channel
            .basic_qos(prefetch_count, BasicQosOptions::default())
            .await?;

        Ok(Self {
            channel: Arc::new(channel),
        })
    }

    /// Wraps an already created channel. Useful when the connection lifecycle
    /// is managed outside of this client.
    pub fn from_channel(channel: Arc<Channel>) -> Self {
        Self { channel }
    }

    /// Returns the underlying channel wrapped into `Arc`.
    pub fn channel(&self) -> Arc<Channel> {
        Arc::clone(&self.channel)
    }

    /// Declares an exchange (creates it if it doesn't exist).
    pub async fn declare_exchange(
        &self,
        name: &str,
        kind: ExchangeKind,
        durable: bool,
    ) -> Result<()> {
        let options = ExchangeDeclareOptions {
            durable,
            ..Default::default()
        };
        self.channel
            .exchange_declare(name, kind, options, FieldTable::default())
            .await?;
        Ok(())
    }

    /// Declares a queue (creates it if it doesn't exist).
    pub async fn declare_queue(&self, name: &str, durable: bool) -> Result<()> {
        let options = QueueDeclareOptions {
            durable,
            ..Default::default()
        };
        self.channel
            .queue_declare(name, options, FieldTable::default())
            .await?;
        Ok(())
    }

    /// Binds a queue to an exchange with the given routing key.
    pub async fn bind_queue(&self, queue: &str, exchange: &str, routing_key: &str) -> Result<()> {
        self.channel
            .queue_bind(
                queue,
                exchange,
                routing_key,
                QueueBindOptions::default(),
                FieldTable::default(),
            )
            .await?;
        Ok(())
    }

    fn handle_confirmation(confirmation: Confirmation) -> Result<()> {
        match confirmation {
            Confirmation::NotRequested => {
                debug!("Publish confirmation not requested");
            }
            Confirmation::Nack(resp) => {
                debug!("Message was not acknowledged by server: {resp:?}");
            }
            Confirmation::Ack(Some(resp)) => {
                debug!(
                    code = resp.reply_code,
                    msg = %resp.reply_text,
                    "Message was acknowledged by server, but response has a message"
                );
                return Err(crate::Error::PublishAck {
                    code: resp.reply_code,
                    text: resp.reply_text,
                });
            }
            Confirmation::Ack(None) => {}
        }
        Ok(())
    }

    /// Publishes a raw payload to the exchange with the given routing key.
    ///
    /// - The message is published with the `mandatory` flag set.
    /// - The message is marked as persistent (`delivery_mode = 2`).
    /// - The publisher confirm result is awaited and handled:
    ///   `Ack(Some(basic_return))` (unroutable mandatory message) is returned
    ///   as an error.
    pub async fn publish(
        &self,
        exchange: &str,
        routing_key: &str,
        payload: &[u8],
        headers: Option<FieldTable>,
    ) -> Result<()> {
        let publish_options = BasicPublishOptions {
            mandatory: true,
            immediate: false,
        };
        let confirmation = self
            .channel
            .basic_publish(
                exchange,
                routing_key,
                publish_options,
                payload,
                BasicProperties::default()
                    .with_delivery_mode(2) // 2 - persistence mode
                    .with_headers(headers.unwrap_or_default()),
            )
            .await?
            .await?;

        Self::handle_confirmation(confirmation)
    }

    /// Publishes with custom properties. Use [`Self::publish`] for defaults.
    pub async fn publish_with_properties(
        &self,
        exchange: &str,
        routing_key: &str,
        payload: &[u8],
        properties: BasicProperties,
    ) -> Result<()> {
        let publish_options = BasicPublishOptions {
            mandatory: true,
            immediate: false,
        };
        let confirmation = self
            .channel
            .basic_publish(exchange, routing_key, publish_options, payload, properties)
            .await?
            .await?;

        Self::handle_confirmation(confirmation)
    }

    /// Serializes `body` to JSON and publishes it. See [`Self::publish`].
    pub async fn publish_json<T: serde::Serialize>(
        &self,
        exchange: &str,
        routing_key: &str,
        body: &T,
        headers: Option<FieldTable>,
    ) -> Result<()> {
        let payload = serde_json::to_vec(body)?;
        self.publish(exchange, routing_key, &payload, headers).await
    }

    /// Starts a consumer task. Handler must return Result<(), ProcessingError<E>>.
    ///
    /// Use `ResultExt` methods to easily convert errors:
    /// - `.requeue_on_err()` - error will cause message to be requeued
    /// - `.discard_on_err()` - error will cause message to be discarded
    pub fn consume<F, Fut, E>(&self, queue: &str, mut handler: F) -> JoinHandle<()>
    where
        F: FnMut(Delivery) -> Fut + Send + 'static,
        Fut: Future<Output = std::result::Result<(), ProcessingError<E>>> + Send + 'static,
        E: std::fmt::Display + Send + 'static,
    {
        let channel = Arc::clone(&self.channel);
        let queue_name = queue.to_string();

        tokio::spawn(async move {
            let mut consumer = match channel
                .basic_consume(
                    &queue_name,
                    "",
                    BasicConsumeOptions::default(),
                    FieldTable::default(),
                )
                .await
            {
                Ok(consumer) => consumer,
                Err(error) => {
                    error!(%error, "[rabbit_warren] Failed to start consumer for queue `{queue_name}`");
                    return;
                }
            };

            while let Some(delivery_result) = consumer.next().await {
                let delivery = match delivery_result {
                    Ok(delivery) => delivery,
                    Err(error) => {
                        error!(%error, "[rabbit_warren] Consumer stream error");
                        continue;
                    }
                };

                let acker = delivery.acker.clone();
                let span = tracing::info_span!("rmq_consume", delivery_tag = delivery.delivery_tag);
                let process_message = handler(delivery).instrument(span);

                match process_message.await {
                    Ok(()) => {
                        if let Err(error) = acker.ack(BasicAckOptions::default()).await {
                            error!(%error, "[rabbit_warren] Failed to ack message");
                        }
                    }
                    Err(processing_error) => {
                        let action = processing_error.ack_action();
                        let error_msg = processing_error.to_string();

                        match action {
                            AckAction::Ack => {
                                warn!(error = %error_msg, "[rabbit_warren] Error returned AckAction::Ack, acknowledging anyway");
                                if let Err(error) = acker.ack(BasicAckOptions::default()).await {
                                    error!(%error, "[rabbit_warren] Failed to ack message");
                                }
                            }
                            AckAction::Requeue => {
                                error!(error = %error_msg, "[rabbit_warren] Message processing failed, requeueing with backoff");
                                tokio::time::sleep(std::time::Duration::from_millis(500)).await;

                                let nack_options = BasicNackOptions {
                                    requeue: true,
                                    ..Default::default()
                                };
                                if let Err(error) = acker.nack(nack_options).await {
                                    error!(%error, "[rabbit_warren] Failed to nack message with requeue");
                                }
                            }
                            AckAction::Discard => {
                                error!(error = %error_msg, "[rabbit_warren] Message processing failed permanently, discarding");
                                let nack_options = BasicNackOptions {
                                    requeue: false,
                                    ..Default::default()
                                };
                                if let Err(error) = acker.nack(nack_options).await {
                                    error!(%error, "[rabbit_warren] Failed to nack message without requeue");
                                }
                            }
                        }
                    }
                }
            }

            error!("[rabbit_warren] Consumer stream ended unexpectedly for queue `{queue_name}`");
        })
    }
}

/// Converts the delivery payload to a UTF-8 string.
pub fn payload_as_utf8(delivery: &Delivery) -> Result<&str> {
    std::str::from_utf8(delivery.data.as_slice()).map_err(Into::into)
}

/// Deserializes the delivery payload from JSON.
pub fn deserialize_delivery<T>(delivery: &Delivery) -> Result<T>
where
    T: serde::de::DeserializeOwned,
{
    let payload = payload_as_utf8(delivery)?;
    serde_json::from_str(payload).map_err(Into::into)
}

/// Returns empty default headers.
///
/// To implement OpenTelemetry trace context, use the standard `opentelemetry` crate,
/// convert its promise to `FieldTable`, and pass it as the `headers` argument
/// to the `publish` / `publish_json` methods.
pub fn current_trace_headers() -> FieldTable {
    FieldTable::default()
}