ollie 0.2.10

//! This module provides a RabbitMQ router for consuming and handling messages
//! with asynchronous handlers.

use async_trait::async_trait;
use futures::StreamExt;
use lapin::BasicProperties;
use lapin::{options::*, types::FieldTable, Connection, ConnectionProperties, Consumer};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::exchange::Exchange;
use crate::rabbit_result::RabbitResult;

#[async_trait]
/// Trait for handling asynchronous message processing.
/// - `S` represents the shared state that can be used by the handler.
pub trait AsyncHandler<S>: Send + Sync {
    /// Asynchronously handle incoming data.
    /// - `data`: The message data as a byte vector.
    /// - `state`: Shared state for processing messages.
    /// - Returns an optional `RabbitResult`.
    async fn handle(&self, data: Vec<u8>, state: Arc<Mutex<S>>) -> Option<RabbitResult>;
}

#[async_trait]
/// Implementation of `AsyncHandler` for closures.
/// This allows closures to be used as message handlers.
impl<S, F, Fut> AsyncHandler<S> for F
where
    S: Send + Sync + 'static,
    F: Fn(Vec<u8>, Arc<Mutex<S>>) -> Fut + Send + Sync,
    Fut: std::future::Future<Output = Option<RabbitResult>> + Send + 'static,
{
    async fn handle(&self, data: Vec<u8>, state: Arc<Mutex<S>>) -> Option<RabbitResult> {
        (self)(data, state).await
    }
}

/// Type alias for the shared state used by the router.
/// This is a mutex-protected shared state that can be accessed by the handlers.
/// Defined here so that it can be easily referenced in end user code.
/// - `S` represents the shared state.
///
pub type RabbitRouterState<S> = Arc<Mutex<S>>;

/// A router for managing RabbitMQ connections, channels, and routes.
pub struct RabbitRouter<S> {
    /// Shared RabbitMQ connection.
    pub connection: Arc<Connection>,
    /// Shared RabbitMQ channel.
    pub channel: Arc<lapin::Channel>,
    /// Map of queue names to their respective handlers.
    pub routes: Arc<Mutex<HashMap<String, Arc<dyn AsyncHandler<S>>>>>,
    /// Shared application state.
    pub state: Arc<Mutex<S>>,
}

impl<S> RabbitRouter<S>
where
    S: Send + Sync + 'static,
{
    /// Create a new `RabbitRouter` instance with a RabbitMQ connection.
    /// - `uri`: The RabbitMQ connection URI.
    /// - `initial_state`: The initial shared state.
    pub async fn new(uri: &str, initial_state: S) -> Self {
        let connection = Connection::connect(uri, ConnectionProperties::default())
            .await
            .expect("Failed to connect to RabbitMQ");
        let channel = connection
            .create_channel()
            .await
            .expect("Failed to create channel");

        RabbitRouter {
            connection: Arc::new(connection),
            channel: Arc::new(channel),
            routes: Arc::new(Mutex::new(HashMap::new())),
            state: Arc::new(Mutex::new(initial_state)),
        }
    }

    pub async fn add_rpc_exchange<H>(
        &self,
        exchange: &str,
        routing_key: &str,
        handler: H,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        H: AsyncHandler<S> + 'static,
    {
        let exchange = exchange.to_string();

        let routing_key = routing_key.to_string();
        let handler = Arc::new(handler);

        let channel = self.channel.clone();
        let mut exchange_options = ExchangeDeclareOptions::default();
        exchange_options.durable = false;
        // Declare the exchange
        channel
            .exchange_declare(
                &exchange,
                lapin::ExchangeKind::Topic,
                exchange_options,
                FieldTable::default(),
            )
            .await?;
        let mut options = QueueDeclareOptions::default();
        options.exclusive = true;
        options.durable = true;

        // Declare a temporary queue.
        let result = channel
            .queue_declare("", options, FieldTable::default())
            .await?;
        let queue_name = result.name().as_str().to_string();

        // Register the handler for the queue.
        let mut routes = self.routes.lock().await;
        routes.insert(queue_name.clone(), handler);
        let routes = self.routes.clone();

        let state = self.state.clone();

        // Bind the queue to the exchange.
        channel
            .queue_bind(
                &queue_name,
                &exchange,
                &routing_key,
                QueueBindOptions::default(),
                FieldTable::default(),
            )
            .await?;

        // Start consuming messages.
        tokio::spawn(async move {
            let consumer = channel
                .basic_consume(
                    &queue_name,
                    "",
                    BasicConsumeOptions::default(),
                    FieldTable::default(),
                )
                .await
                .expect("Failed to start consumer");

            Self::consume_rpc(
                queue_name,
                consumer,
                channel,
                routes,
                state,
            )
            .await;
        });
        Ok(())
    }

    /// Add a route to a temporary queue bound to a specific exchange.
    /// - `exchange`: The name of the exchange.
    /// - `routing_key`: The routing key for the queue binding.
    /// - `result_exchange`: Optional exchange to send processing results.
    /// - `handler`: The message handler.
    pub async fn add_route_exchange<H>(
        &self,
        exchange: &str,
        routing_key: &str,
        result_exchange: Option<Exchange>,
        handler: H,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        H: AsyncHandler<S> + 'static,
    {
        let exchange = exchange.to_string();

        let routing_key = routing_key.to_string();
        let handler = Arc::new(handler);

        let channel = self.channel.clone();
        let mut exchange_options = ExchangeDeclareOptions::default();
        exchange_options.durable = false;
        // Declare the exchange
        channel
            .exchange_declare(
                &exchange,
                lapin::ExchangeKind::Topic,
                exchange_options,
                FieldTable::default(),
            )
            .await?;
        let mut options = QueueDeclareOptions::default();
        options.exclusive = true;
        options.durable = true;

        // Declare a temporary queue.
        let result = channel
            .queue_declare("", options, FieldTable::default())
            .await?;
        let queue_name = result.name().as_str().to_string();

        // Register the handler for the queue.
        let mut routes = self.routes.lock().await;
        routes.insert(queue_name.clone(), handler);
        let routes = self.routes.clone();

        let state = self.state.clone();

        // Bind the queue to the exchange.
        channel
            .queue_bind(
                &queue_name,
                &exchange,
                &routing_key,
                QueueBindOptions::default(),
                FieldTable::default(),
            )
            .await?;

        //Create the results exchange if one is provided
        Self::declare_result_exchange(channel.clone(), result_exchange.clone()).await?;

        // Start consuming messages.
        tokio::spawn(async move {
            let consumer = channel
                .basic_consume(
                    &queue_name,
                    "",
                    BasicConsumeOptions::default(),
                    FieldTable::default(),
                )
                .await
                .expect("Failed to start consumer");

            Self::consume_messages(
                queue_name,
                consumer,
                channel,
                result_exchange,
                routes,
                state,
            )
            .await;
        });
        Ok(())
    }

    /// Add a route to a specific queue.
    /// - `queue_name`: The name of the queue.
    /// - `source_exchange`: Optional exchange to bind the queue to to receive messages.
    /// - `result_exchange`: Optional exchange to send processing results.
    /// - `handler`: The message handler.
    pub async fn add_route_queue<H>(
        &self,
        queue_name: &str,
        source_exchange: Option<Exchange>,
        result_exchange: Option<Exchange>,
        handler: H,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        H: AsyncHandler<S> + 'static,
    {
        let queue_name = queue_name.to_string();
        let handler = Arc::new(handler);

        // Register the handler for the queue.
        let mut routes = self.routes.lock().await;
        routes.insert(queue_name.clone(), handler);
        let routes = self.routes.clone();

        let state = self.state.clone();

        // Declare the queue.
        let channel = self.channel.clone();
        let _queue = channel
            .queue_declare(
                &queue_name,
                QueueDeclareOptions::default(),
                FieldTable::default(),
            )
            .await
            .expect("Failed to declare queue");

        // Bind the queue to the exchange if one is provided.
        if let Some(queue_exchange) = source_exchange {
            let mut exchange_options = ExchangeDeclareOptions::default();
            exchange_options.durable = false;
            channel
                .exchange_declare(
                    &queue_exchange.name,
                    lapin::ExchangeKind::Topic,
                    exchange_options,
                    FieldTable::default(),
                )
                .await
                .expect("Failed to declare exchange");
            channel
                .queue_bind(
                    &queue_name,
                    &queue_exchange.name,
                    &queue_exchange.routing_key,
                    QueueBindOptions::default(),
                    FieldTable::default(),
                )
                .await
                .expect("Failed to bind queue to exchange");
        }

        //Create the results exchange if one is provided
        Self::declare_result_exchange(channel.clone(), result_exchange.clone()).await?;

        // Start consuming messages.
        tokio::spawn(async move {
            let consumer = channel
                .basic_consume(
                    &queue_name,
                    "",
                    BasicConsumeOptions::default(),
                    FieldTable::default(),
                )
                .await
                .expect("Failed to start consumer");

            Self::consume_messages(
                queue_name,
                consumer,
                channel,
                result_exchange,
                routes,
                state,
            )
            .await;
        });
        Ok(())
    }

    /// Internal method to consume messages from a queue.
    /// - `queue_name`: The name of the queue.
    /// - `consumer`: The consumer for the queue.
    /// - `channel`: The RabbitMQ channel.
    /// - `result_exchange`: Optional exchange to send processing results.
    /// - `routes`: Map of routes and their handlers.
    /// - `state`: Shared application state.
    async fn consume_messages(
        queue_name: String,
        mut consumer: Consumer,
        channel: Arc<lapin::Channel>,
        result_exchange: Option<Exchange>,
        routes: Arc<Mutex<HashMap<String, Arc<dyn AsyncHandler<S>>>>>,
        state: Arc<Mutex<S>>,
    ) {
        while let Some(delivery) = consumer.next().await {
            match delivery {
                Ok(delivery) => {
                    let data = delivery.data.clone();

                    // Retrieve and invoke the appropriate handler.
                    let routes = routes.lock().await;
                    if let Some(handler) = routes.get(&queue_name) {
                        let result = handler.handle(data, state.clone()).await;
                        if let Some(result) = result {
                            // Publish the result if a result exchange is provided.
                            if let Some(result_exchange) = &result_exchange {
                                let full_key = result.logging_level.to_string()
                                    + "."
                                    + &result.billing_type.to_string()
                                    + "."
                                    + result_exchange.routing_key.as_str();
                                let payload = serde_json::to_vec(&result).unwrap();

                                channel
                                    .basic_publish(
                                        &result_exchange.name,
                                        &full_key,
                                        BasicPublishOptions::default(),
                                        &payload,
                                        BasicProperties::default(),
                                    )
                                    .await
                                    .expect("Failed to publish message");
                            }
                        }
                    }

                    // Acknowledge the message.
                    delivery
                        .ack(BasicAckOptions::default())
                        .await
                        .expect("Failed to ack message");
                }
                Err(error) => eprintln!("Error receiving message: {:?}", error),
            }
        }
    }

    async fn consume_rpc(
        queue_name: String,
        mut consumer: Consumer,
        channel: Arc<lapin::Channel>,
        routes: Arc<Mutex<HashMap<String, Arc<dyn AsyncHandler<S>>>>>,
        state: Arc<Mutex<S>>,
    ) {
        while let Some(delivery) = consumer.next().await {
            match delivery {
                Ok(delivery) => {
                    let data = delivery.data.clone();

                    // Retrieve and invoke the appropriate handler.
                    let routes = routes.lock().await;
                    if let Some(handler) = routes.get(&queue_name) {
                        let result = handler.handle(data, state.clone()).await;
                        if let Some(result) = result {
                            // Publish the result if a result exchange is provided.
                            let routing_key =
                                delivery.properties.reply_to().as_ref().unwrap().as_str();

                            let correlation_id =
                                delivery.properties.correlation_id().clone().unwrap();

                            let payload = serde_json::to_vec(&result).unwrap();

                            channel
                                .basic_publish(
                                    "",
                                    routing_key,
                                    BasicPublishOptions::default(),
                                    &payload,
                                    BasicProperties::default().with_correlation_id(correlation_id),
                                )
                                .await
                                .expect("Failed to publish message");

                            channel
                                .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                                .await
                                .expect("Failed to publish message");
                        }
                    }

                
                }
                Err(error) => eprintln!("Error receiving message: {:?}", error),
            }
        }
    }
    /// Internal method to consume messages from a queue.
    /// - `queue_name`: The name of the queue.
    /// - `consumer`: The consumer for the queue.
    /// - `channel`: The RabbitMQ channel.
    /// - `result_exchange`: Optional exchange to send processing results.
    /// - `routes`: Map of routes and their handlers.
    /// - `state`: Shared application state.
    async fn declare_result_exchange(
        channel: Arc<lapin::Channel>,
        result_exchange: Option<Exchange>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        if let Some(exchange) = result_exchange {
            let mut exchange_options = ExchangeDeclareOptions::default();
            exchange_options.durable = false;
            channel
                .exchange_declare(
                    &exchange.name,
                    lapin::ExchangeKind::Topic,
                    exchange_options,
                    FieldTable::default(),
                )
                .await?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lapin::BasicProperties;
    use tokio::time::{sleep, Duration};

    async fn test_handler(
        data: Vec<u8>,
        state: Arc<Mutex<Arc<Mutex<HashMap<String, i32>>>>>,
    ) -> Option<RabbitResult> {
        let state = state.lock().await;
        let mut inner_state = state.lock().await;
        let count = inner_state.entry("shared_counter".to_string()).or_insert(0);
        *count += 1;

        println!(
            "Handler for queue_1 received: {:?}, updated shared_counter: {}",
            String::from_utf8_lossy(&data),
            *count
        );
        None
    }

    #[tokio::test]
    async fn test_rabbit_router() {
        // Create a router
        let state = HashMap::new();
        let initial_state = Arc::new(Mutex::new(state));
        let router = RabbitRouter::new("amqp://127.0.0.1:5672", initial_state.clone()).await;

        // Add a route and send a test message
        router
            .add_route_queue("test_queue", None, None, test_handler)
            .await
            .expect("Failed to add route for queue_1");

        // Simulate sending a message to the queue
        let channel = router.channel.clone();
        let payload = b"Test message".to_vec();
        channel
            .basic_publish(
                "",
                "test_queue",
                BasicPublishOptions::default(),
                &payload,
                BasicProperties::default(),
            )
            .await
            .expect("Failed to publish message");

        // Give some time for the message to be processed
        sleep(Duration::from_secs(2)).await;
        println!("State: {:?}", initial_state.lock().await);
        assert!(initial_state.lock().await.get("shared_counter").is_some());
        assert!(initial_state.lock().await.get("shared_counter").unwrap() == &1);
    }

    #[tokio::test]
    async fn test_multiple_routes() {
        // Create a router
        let state = HashMap::new();
        let initial_state = Arc::new(Mutex::new(state));
        let router = RabbitRouter::new("amqp://127.0.0.1:5672/%2f", initial_state.clone()).await;

        // Add multiple routes
        let _ = router
            .add_route_queue("queue_1", None, None, test_handler)
            .await;

        let _ = router
            .add_route_queue("queue_2", None, None, test_handler)
            .await;

        // Simulate sending messages to the queues
        let channel = router.channel.clone();
        let payload_1 = b"Message for queue_1".to_vec();
        let payload_2 = b"Message for queue_2".to_vec();

        channel
            .basic_publish(
                "",
                "queue_1",
                BasicPublishOptions::default(),
                &payload_1,
                BasicProperties::default(),
            )
            .await
            .expect("Failed to publish message to queue_1");

        channel
            .basic_publish(
                "",
                "queue_2",
                BasicPublishOptions::default(),
                &payload_2,
                BasicProperties::default(),
            )
            .await
            .expect("Failed to publish message to queue_2");

        // Give some time for the messages to be processed
        sleep(Duration::from_secs(2)).await;
        assert!(initial_state.lock().await.get("shared_counter").is_some());
        assert!(initial_state.lock().await.get("shared_counter").unwrap() == &2);
    }

    #[tokio::test]
    async fn test_exchange_route() {
        // Create a router
        let state = HashMap::new();
        let initial_state = Arc::new(Mutex::new(state));
        let router = RabbitRouter::new("amqp://127.0.0.1:5672", initial_state.clone()).await;
        // Declare the exchange
        let channel = router.channel.clone();
        channel
            .exchange_declare(
                "test_exchange",
                lapin::ExchangeKind::Topic,
                ExchangeDeclareOptions::default(),
                FieldTable::default(),
            )
            .await
            .expect("Failed to declare exchange");
        // Add a route to an exchange
        let _ = router
            .add_route_exchange("test_exchange", "test.routing.key", None, test_handler)
            .await;

        // Simulate sending a message to the exchange
        let channel = router.channel.clone();
        let payload = b"Test exchange message".to_vec();
        channel
            .basic_publish(
                "test_exchange",
                "test.routing.key",
                BasicPublishOptions::default(),
                &payload,
                BasicProperties::default(),
            )
            .await
            .expect("Failed to publish message to exchange");
        println!("Message published to exchange");
        // Give some time for the message to be processed
        sleep(Duration::from_secs(2)).await;
        assert!(initial_state.lock().await.get("shared_counter").is_some());
        assert!(initial_state.lock().await.get("shared_counter").unwrap() == &1);
    }
}