mq-bridge 0.2.15

//  mq-bridge
//  © Copyright 2025, by Marco Mengelkoch
//  Licensed under MIT License, see License file for more details
//  git clone https://github.com/marcomq/mq-bridge

use crate::endpoints::{create_consumer_from_route, create_publisher_from_route};
use crate::errors::ProcessingError;
pub use crate::models::Route;
use crate::models::{Endpoint, EndpointType, RouteOptions};
use crate::traits::{
    BatchCommitFunc, ConsumerError, Handler, HandlerError, MessageConsumer, MessageDisposition,
    MessagePublisher, PublisherError, SentBatch,
};
use async_channel::{bounded, Sender};
use serde::de::DeserializeOwned;
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, OnceLock, RwLock};
use tokio::{
    select,
    task::{JoinHandle, JoinSet},
};
use tracing::{debug, error, info, trace, warn};

// Re-export extensions for backward compatibility and internal usage
pub use crate::extensions::{
    get_endpoint_factory, get_middleware_factory, register_endpoint_factory,
    register_middleware_factory,
};

#[derive(Debug)]
pub struct RouteHandle((JoinHandle<()>, Sender<()>));

impl RouteHandle {
    pub async fn stop(&self) {
        let _ = self.0 .1.send(()).await;
        self.0 .1.close();
    }

    pub async fn join(self) -> Result<(), tokio::task::JoinError> {
        self.0 .0.await
    }
}

async fn run_publisher_connect_hook(
    route_name: &str,
    publisher: &Arc<dyn MessagePublisher>,
) -> anyhow::Result<()> {
    if let Some(hook) = publisher.on_connect_hook() {
        hook.await.map_err(|err| {
            anyhow::anyhow!(
                "Publisher on_connect hook failed for route '{}': {}",
                route_name,
                err
            )
        })?;
    }
    Ok(())
}

async fn run_consumer_connect_hook(
    route_name: &str,
    consumer: &dyn MessageConsumer,
) -> anyhow::Result<()> {
    if let Some(hook) = consumer.on_connect_hook() {
        hook.await.map_err(|err| {
            anyhow::anyhow!(
                "Consumer on_connect hook failed for route '{}': {}",
                route_name,
                err
            )
        })?;
    }
    Ok(())
}

async fn run_publisher_disconnect_hook(route_name: &str, publisher: &Arc<dyn MessagePublisher>) {
    if let Some(hook) = publisher.on_disconnect_hook() {
        if let Err(err) = hook.await {
            warn!(
                "Publisher on_disconnect hook failed for route '{}': {}",
                route_name, err
            );
        }
    }
}

async fn run_consumer_disconnect_hook(route_name: &str, consumer: &dyn MessageConsumer) {
    if let Some(hook) = consumer.on_disconnect_hook() {
        if let Err(err) = hook.await {
            warn!(
                "Consumer on_disconnect hook failed for route '{}': {}",
                route_name, err
            );
        }
    }
}

impl From<(JoinHandle<()>, Sender<()>)> for RouteHandle {
    fn from(tuple: (JoinHandle<()>, Sender<()>)) -> Self {
        RouteHandle(tuple)
    }
}

struct ActiveRoute {
    route: Route,
    handle: RouteHandle,
}

static ROUTE_REGISTRY: OnceLock<RwLock<HashMap<String, ActiveRoute>>> = OnceLock::new();
static ENDPOINT_REF_REGISTRY: OnceLock<RwLock<HashMap<String, Endpoint>>> = OnceLock::new();

/// Registers a named endpoint that can be referenced by other endpoints using `ref: "name"`.
/// This will overwrite any existing endpoint with the same name.
pub fn register_endpoint(name: &str, endpoint: Endpoint) {
    let registry = ENDPOINT_REF_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
    let mut writer = registry
        .write()
        .expect("Named endpoint registry lock poisoned");
    if writer.insert(name.to_string(), endpoint).is_some() {
        debug!("Overwriting a registered endpoint named '{}'", name);
    }
}

/// Retrieves a registered endpoint by name.
pub fn get_endpoint(name: &str) -> Option<Endpoint> {
    let registry = ENDPOINT_REF_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
    let reader = registry
        .read()
        .expect("Named endpoint registry lock poisoned");
    reader.get(name).cloned()
}

impl Route {
    /// Creates a new route with default concurrency (1) and batch size (128).
    ///
    /// # Arguments
    /// * `input` - The input/source endpoint for the route
    /// * `output` - The output/sink endpoint for the route
    pub fn new(input: Endpoint, output: Endpoint) -> Self {
        Self {
            input,
            output,
            ..Default::default()
        }
    }

    /// Retrieves a registered (and running) route by name.
    pub fn get(name: &str) -> Option<Self> {
        let registry = ROUTE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
        let map = registry.read().expect("Route registry lock poisoned");
        map.get(name).map(|active| active.route.clone())
    }

    /// Returns a list of all registered route names.
    pub fn list() -> Vec<String> {
        let registry = ROUTE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
        let map = registry.read().expect("Route registry lock poisoned");
        map.keys().cloned().collect()
    }

    /// Returns true if the input is of type ref (and the output isn't)
    pub fn is_ref(&self) -> bool {
        matches!(self.input.endpoint_type, EndpointType::Ref(_))
            && !matches!(self.output.endpoint_type, EndpointType::Ref(_))
    }

    /// Registers the route's output endpoint under the given name.
    /// This allows other routes to reference this output using `ref: "name"`.
    pub fn register_output_endpoint(&self, name: Option<&str>) -> Result<(), anyhow::Error> {
        match name {
            Some(name) => {
                register_endpoint(name, self.output.clone());
            }
            None => {
                if let EndpointType::Ref(name) = &self.input.endpoint_type {
                    register_endpoint(name, self.output.clone());
                } else {
                    return Err(anyhow::anyhow!(
                        "No name and input is not a reference endpoint"
                    ));
                }
            }
        };
        Ok(())
    }

    /// Registers the route and starts it.
    /// If a route with the same name is already running, it will be stopped first.
    ///    
    /// # Examples
    /// ```
    /// use mq_bridge::{Route, models::Endpoint};
    ///
    /// let route = Route::new(Endpoint::new_memory("in", 10), Endpoint::new_memory("out", 10));
    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
    /// route.deploy("global_route").await.unwrap();
    /// assert!(Route::get("global_route").is_some());
    /// # });
    /// ```
    pub async fn deploy(&self, name: &str) -> anyhow::Result<()> {
        Self::stop(name).await;

        let handle = self.run(name).await?;
        let active = ActiveRoute {
            route: self.clone(),
            handle,
        };

        let registry = ROUTE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
        let mut map = registry.write().expect("Route registry lock poisoned");
        map.insert(name.to_string(), active);
        Ok(())
    }

    /// Stops a running route by name and removes it from the registry.
    /// Waits up to 5 seconds for the route task to join; if the timeout elapses
    /// the task is aborted and the implementation awaits the aborted handle to
    /// ensure the background task has fully terminated before returning.
    pub async fn stop(name: &str) -> bool {
        let registry = ROUTE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()));
        let active_opt = {
            let mut map = registry.write().expect("Route registry lock poisoned");
            map.remove(name)
        };

        if let Some(active) = active_opt {
            // Move the handle out so we can operate on its internals.
            let handle = active.handle;

            // Signal the route to stop and close the shutdown channel.
            let _ = handle.0 .1.send(()).await;
            handle.0 .1.close();

            // Extract the JoinHandle so we can monitor and, if needed, abort it.
            let mut join_handle = handle.0 .0;
            tokio::select! {
                res = &mut join_handle => {
                    // The task finished naturally within the 5s window
                    let _ = res;
                }
                _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => {
                    // The 5s timer finished first - abort the task to ensure it doesn't linger.
                    join_handle.abort();
                    // Await the handle one last time to ensure the task has fully shut down.
                    let _ = join_handle.await;
                }
            }

            true
        } else {
            false
        }
    }

    /// Creates a new Publisher configured for this route's output.
    /// This is useful if you want to send messages to the same destination as this route.
    ///
    /// # Examples
    ///
    /// ```
    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
    /// use mq_bridge::{Route, models::Endpoint};
    ///
    /// let route = Route::new(Endpoint::new_memory("in", 10), Endpoint::new_memory("out", 10));
    /// let publisher = route.create_publisher().await;
    /// assert!(publisher.is_ok());
    /// # });
    /// ```
    pub async fn create_publisher(&self) -> anyhow::Result<crate::Publisher> {
        crate::Publisher::new(self.output.clone()).await
    }

    /// Creates a consumer connected to the route's output.
    /// This is primarily useful for integration tests to verify messages reaching the destination.
    pub async fn connect_to_output(
        &self,
        name: &str,
    ) -> anyhow::Result<Box<dyn crate::traits::MessageConsumer>> {
        create_consumer_from_route(name, &self.output).await
    }

    /// Validates the route configuration, checking if endpoints are supported and correctly configured.
    /// Core types like file, memory, and response are always supported.
    /// # Arguments
    /// * `name` - The name of the route
    /// * `allowed_endpoints` - An optional list of allowed endpoint types
    pub fn check(
        &self,
        name: &str,
        allowed_endpoints: Option<&[&str]>,
    ) -> anyhow::Result<Vec<String>> {
        let mut warnings = Vec::new();
        warnings.extend(crate::endpoints::check_consumer(
            name,
            &self.input,
            allowed_endpoints,
        )?);
        warnings.extend(crate::endpoints::check_publisher(
            name,
            &self.output,
            allowed_endpoints,
        )?);
        Ok(warnings)
    }

    /// Runs the message processing route with concurrency, error handling, and graceful shutdown.
    ///
    /// This function spawns the necessary background tasks to process messages. It waits asynchronously
    /// until the route is successfully initialized (i.e., connections are established) or until
    /// a timeout occurs.
    /// The name_str parameter is just used for logging and tracing.
    ///
    /// It returns a `JoinHandle` for the main route task and a `Sender` channel
    /// that can be used to signal a graceful shutdown. The result is typically converted into a
    /// [`RouteHandle`] for easier management.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use mq_bridge::{Route, route::RouteHandle, models::Endpoint};
    /// # async fn example() -> anyhow::Result<()> {
    /// let route = Route::new(Endpoint::new_memory("in", 10), Endpoint::new_memory("out", 10));
    ///
    /// // Start the route (blocks until initialized) and convert to RouteHandle
    /// let handle: RouteHandle = route.run("my_route").await?.into();
    ///
    /// // Stop the route later
    /// handle.stop().await;
    /// handle.join().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn run(&self, name_str: &str) -> anyhow::Result<RouteHandle> {
        let warnings = self.check(name_str, None)?;
        for warning in warnings {
            tracing::warn!(route = name_str, "Configuration warning: {}", warning);
        }
        let (shutdown_tx, shutdown_rx) = bounded(1);
        let (ready_tx, ready_rx) = bounded(1);
        // Use `Arc` so route/name clones are cheap (pointer copy) in the reconnect loop.
        let route = Arc::new(self.clone());
        let name = Arc::new(name_str.to_string());

        let handle = tokio::spawn(async move {
            loop {
                let route_arc = Arc::clone(&route);
                let name_arc = Arc::clone(&name);
                // Create a new, per-iteration internal shutdown channel.
                // This avoids a race where both this loop and the inner task
                // try to consume the same external shutdown signal.
                let (internal_shutdown_tx, internal_shutdown_rx) = bounded(1);
                let ready_tx_clone = ready_tx.clone();

                // The actual route logic is in `run_until_err`.
                let mut run_task = tokio::spawn(async move {
                    route_arc
                        .run_until_err(&name_arc, Some(internal_shutdown_rx), Some(ready_tx_clone))
                        .await
                });

                select! {
                    _ = shutdown_rx.recv() => {
                        info!("Shutdown signal received for route '{}'.", name);
                        // Notify the inner task to shut down.
                        let _ = internal_shutdown_tx.send(()).await;
                        // Wait for the inner task to finish gracefully.
                        let _ = run_task.await;
                        break;
                    }
                    res = &mut run_task => {
                        match res {
                            Ok(Ok(should_continue)) if !should_continue => {
                                info!("Route '{}' completed gracefully. Shutting down.", name);
                                break;
                            }
                            Ok(Err(e)) => {
                                let is_permanent =
                                    e.downcast_ref::<ProcessingError>().is_some_and(|pe| matches!(pe, ProcessingError::NonRetryable(_)))
                                    || e.downcast_ref::<ConsumerError>().is_some_and(|ce| matches!(ce, ConsumerError::EndOfStream));

                                if is_permanent {
                                    error!("Route '{}' failed with a permanent error: {}. Shutting down.", name, e);
                                    break;
                                }

                                warn!("Route '{}' failed: {}. Reconnecting in 5 seconds...", name, e);
                                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                            }
                            Err(e) => {
                                error!("Route '{}' task panicked: {}. Reconnecting in 5 seconds...", name, e);
                                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                            }
                            _ => {} // The route should continue running.
                        }
                    }
                }
            }
        });

        match tokio::time::timeout(std::time::Duration::from_secs(5), ready_rx.recv()).await {
            Ok(Ok(_)) => Ok(RouteHandle((handle, shutdown_tx))),
            _ => {
                handle.abort();
                Err(anyhow::anyhow!(
                    "Route '{}' failed to start within 5 seconds or encountered an error",
                    name_str
                ))
            }
        }
    }

    /// The core logic of running the route, designed to be called within a reconnect loop.
    pub async fn run_until_err(
        &self,
        name: &str,
        shutdown_rx: Option<async_channel::Receiver<()>>,
        ready_tx: Option<Sender<()>>,
    ) -> anyhow::Result<bool> {
        let (_internal_shutdown_tx, internal_shutdown_rx) = bounded(1);
        let shutdown_rx = shutdown_rx.unwrap_or(internal_shutdown_rx);
        if self.options.concurrency == 1 {
            self.run_sequentially(name, shutdown_rx, ready_tx).await
        } else {
            self.run_concurrently(name, shutdown_rx, ready_tx).await
        }
    }

    /// A simplified, sequential runner for when concurrency is 1.
    async fn run_sequentially(
        &self,
        name: &str,
        shutdown_rx: async_channel::Receiver<()>,
        ready_tx: Option<Sender<()>>,
    ) -> anyhow::Result<bool> {
        let publisher = create_publisher_from_route(name, &self.output).await?;
        let mut consumer = create_consumer_from_route(name, &self.input).await?;
        if let Err(err) = run_publisher_connect_hook(name, &publisher).await {
            run_publisher_disconnect_hook(name, &publisher).await;
            return Err(err);
        }
        if let Err(err) = run_consumer_connect_hook(name, consumer.as_ref()).await {
            run_consumer_disconnect_hook(name, consumer.as_ref()).await;
            run_publisher_disconnect_hook(name, &publisher).await;
            return Err(err);
        }
        let (err_tx, err_rx) = bounded(1);
        let mut commit_tasks = JoinSet::new();

        // Sequencer setup to ensure ordered commits even with parallel commit tasks
        let (seq_tx, sequencer_handle) = spawn_sequencer(self.options.commit_concurrency_limit);
        let mut seq_counter = 0u64;

        if let Some(tx) = ready_tx {
            let _ = tx.send(()).await;
        }
        let mut message_ids = Vec::with_capacity(self.options.batch_size);
        // Check if retry middleware is present on output
        let has_retry_middleware = self.output.has_retry_middleware();
        let run_result = loop {
            select! {
                Ok(err) = err_rx.recv() => break Err(err),

                _ = shutdown_rx.recv() => {
                    info!("Shutdown signal received in sequential runner for route '{}'.", name);
                    break Ok(true); // Stopped by shutdown signal
                }
                res = consumer.receive_batch(self.options.batch_size) => {
                    let received_batch = match res {
                        Ok(batch) => {
                            if batch.messages.is_empty() {
                                continue; // No messages, loop to select! again
                            }
                            batch
                        }
                        Err(ConsumerError::EndOfStream) => {
                            info!("Consumer for route '{}' reached end of stream. Shutting down.", name);
                            break Ok(false); // Graceful exit
                        }
                        Err(ConsumerError::Connection(e)) => {
                            // Propagate error to trigger reconnect by the outer loop
                            break Err(e);
                        },
                        Err(ConsumerError::Gap { requested, base }) => {
                            // Propagate gap error to trigger reconnect by the outer loop
                            break Err(anyhow::anyhow!("Consumer gap: requested offset {requested} but earliest available is {base}"));
                        }
                    };
                    debug!("Received a batch of {} messages sequentially", received_batch.messages.len());

                    // Process the batch sequentially without spawning a new task
                    let seq = seq_counter;
                    seq_counter += 1;
                    let mut commit_opt = Some(wrap_commit(received_batch.commit, seq, seq_tx.clone()));
                    let batch_len = received_batch.messages.len();
                    message_ids.clear();
                    message_ids.extend(received_batch.messages.iter().map(|m| m.message_id));
                    let request_ids: std::collections::HashSet<u128> = received_batch
                        .messages
                        .iter()
                        .filter(|m| m.metadata.contains_key("reply_to"))
                        .map(|m| m.message_id)
                        .collect();

                    match publisher.send_batch(received_batch.messages).await {
                        Ok(SentBatch::Ack) => {
                            for id in &message_ids {
                                if request_ids.contains(id) {
                                    warn!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Response loop broken.", id);
                                }
                            }
                            let commit = commit_opt.take().expect("Commit already used");
                            let err_tx = err_tx.clone();
                            commit_tasks.spawn(async move {
                                if let Err(e) = commit(vec![MessageDisposition::Ack; batch_len]).await {
                                    error!("Commit failed: {}", e);
                                    match err_tx.try_send(e) {
                                        Ok(_) => trace!("Reported commit error to main task"),
                                        Err(err_send) => warn!(error=?err_send, "Could not send commit error to main task, it might be down or busy."),
                                    }
                                }
                            });
                        }
                        Ok(SentBatch::Partial { responses, failed }) => {
                            // Connection and Retryable are both "transient" errors - treat them the same.
                            // Connection errors from handlers or publishers both indicate a temporary
                            // failure that should either be retried (with middleware) or crash the route.
                            let has_transient = failed.iter().any(|(_, e)| {
                                matches!(e, PublisherError::Retryable(_) | PublisherError::Connection(_))
                            });
                            if has_transient {
                                let (_, first_err) = failed
                                    .iter()
                                    .find(|(_, e)| matches!(e, PublisherError::Retryable(_) | PublisherError::Connection(_)))
                                    .expect("has_transient is true");
                                let err = anyhow::anyhow!(
                                    "Transient error in batch send ({} messages failed). First error: {}",
                                    failed.len(),
                                    first_err
                                );
                                // Ack/nack the batch to fill the sequencer slot.
                                let commit = commit_opt.take().expect("Commit already used");
                                let dispositions =
                                    map_responses_to_dispositions(&message_ids, responses, &failed, &request_ids);
                                if let Err(commit_err) = commit(dispositions).await {
                                    warn!("Commit after transient failure also failed: {}", commit_err);
                                }
                                if !has_retry_middleware {
                                    break Err(err);
                                }
                                warn!("Transient error in batch, message(s) Nack'ed for re-delivery: {}", err);
                                tokio::task::yield_now().await;
                                continue;
                            }
                            // Only non-retryable errors remain - drop them with a log.
                            for (msg, e) in &failed {
                                error!("Dropping message (ID: {:032x}) due to non-retryable error: {}", msg.message_id, e);
                            }
                            let commit = commit_opt.take().expect("Commit already used");
                            let err_tx = err_tx.clone();
                            let ids = std::mem::take(&mut message_ids);
                            let req_ids = request_ids;
                            commit_tasks.spawn(async move {
                                let dispositions = map_responses_to_dispositions(&ids, responses, &failed, &req_ids);
                                if let Err(e) = commit(dispositions).await {
                                    error!("Commit failed: {}", e);
                                    match err_tx.try_send(e) {
                                        Ok(_) => trace!("Reported commit error to main task"),
                                        Err(err_send) => warn!(error=?err_send, "Could not send commit error to main task, it might be down or busy."),
                                    }
                                }
                            });
                        }
                        Err(e) => {
                            // Any direct error from send_batch crashes the route.
                            warn!("Publisher error, sending {} Nacks to commit", batch_len);
                            let commit = commit_opt.take().expect("Commit already used");
                            let nack_result = commit(vec![MessageDisposition::Nack; batch_len]).await;
                            debug!("Nack commit result: {:?}", nack_result);
                            break Err(e.into());
                        }
                    }

                    tokio::task::yield_now().await;
                }
            }
        };

        drop(seq_tx);
        // Drain errors while waiting for tasks to finish to prevent deadlocks and lost errors
        loop {
            select! {
                res = err_rx.recv() => {
                    if let Ok(err) = res {
                        error!("Error reported during shutdown: {}", err);
                    }
                }
                res = commit_tasks.join_next() => {
                    if res.is_none() {
                        break;
                    }
                }
            }
        }
        drop(err_rx);
        let _ = sequencer_handle.await;
        run_consumer_disconnect_hook(name, consumer.as_ref()).await;
        run_publisher_disconnect_hook(name, &publisher).await;
        run_result
    }

    /// The main concurrent runner for when concurrency > 1.
    async fn run_concurrently(
        &self,
        name: &str,
        shutdown_rx: async_channel::Receiver<()>,
        ready_tx: Option<Sender<()>>,
    ) -> anyhow::Result<bool> {
        let publisher = create_publisher_from_route(name, &self.output).await?;
        let mut consumer = create_consumer_from_route(name, &self.input).await?;
        if let Err(err) = run_publisher_connect_hook(name, &publisher).await {
            run_publisher_disconnect_hook(name, &publisher).await;
            return Err(err);
        }
        if let Err(err) = run_consumer_connect_hook(name, consumer.as_ref()).await {
            run_consumer_disconnect_hook(name, consumer.as_ref()).await;
            run_publisher_disconnect_hook(name, &publisher).await;
            return Err(err);
        }
        if let Some(tx) = ready_tx {
            let _ = tx.send(()).await;
        }
        let (err_tx, err_rx) = bounded(1); // For critical, route-stopping errors
                                           // channel capacity: a small buffer proportional to concurrency
        let work_capacity = self
            .options
            .concurrency
            .saturating_mul(self.options.batch_size);
        let (work_tx, work_rx) =
            bounded::<(Vec<crate::CanonicalMessage>, BatchCommitFunc)>(work_capacity);
        // --- Ordered Commit Sequencer ---
        // To prevent data loss with cumulative-ack brokers (Kafka/AMQP), commits must happen in order.
        // We assign a sequence number to each batch and use a sequencer task to enforce order.
        let (seq_tx, sequencer_handle) = spawn_sequencer(self.options.commit_concurrency_limit);

        // --- Worker Pool ---
        let batch_size = self.options.batch_size;
        let mut join_set = JoinSet::new();
        for i in 0..self.options.concurrency {
            let work_rx_clone = work_rx.clone();
            let publisher = Arc::clone(&publisher);
            let err_tx = err_tx.clone();
            let mut commit_tasks = JoinSet::new();
            let has_retry_middleware = self.output.has_retry_middleware();
            join_set.spawn(async move {
                debug!("Starting worker {}", i);
                let mut message_ids = Vec::with_capacity(batch_size);
                while let Ok((messages, commit_func)) = work_rx_clone.recv().await {
                    let mut commit_opt = Some(commit_func);
                    let batch_len = messages.len();
                    message_ids.clear();
                    message_ids.extend(messages.iter().map(|m| m.message_id));
                    let request_ids: std::collections::HashSet<u128> = messages
                        .iter()
                        .filter(|m| m.metadata.contains_key("reply_to"))
                        .map(|m| m.message_id)
                        .collect();

                    match publisher.send_batch(messages).await {
                        Ok(SentBatch::Ack) => {
                            for id in &message_ids {
                                if request_ids.contains(id) {
                                    warn!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Response loop broken.", id);
                                }
                            }
                            let commit = commit_opt.take().expect("Commit already used");
                            let err_tx = err_tx.clone();
                            commit_tasks.spawn(async move {
                                if let Err(e) = commit(vec![MessageDisposition::Ack; batch_len]).await {
                                    error!("Commit failed: {}", e);
                                    match err_tx.try_send(e) {
                                        Ok(_) => trace!("Reported commit error to main task"),
                                        Err(err_send) => warn!(error=?err_send, "Could not send commit error to main task, it might be down or busy."),
                                    }
                                }
                            });
                        }
                        Ok(SentBatch::Partial { responses, failed }) => {
                            // Connection and Retryable are both "transient" errors.
                            let has_transient = failed.iter().any(|(_, e)| {
                                matches!(e, PublisherError::Retryable(_) | PublisherError::Connection(_))
                            });
                            if has_transient {
                                let (_, first_err) = failed
                                    .iter()
                                    .find(|(_, e)| matches!(e, PublisherError::Retryable(_) | PublisherError::Connection(_)))
                                    .expect("has_transient is true");
                                let e = anyhow::anyhow!(
                                    "Transient error in batch send ({} messages failed). First error: {}",
                                    failed.len(),
                                    first_err
                                );
                                let commit = commit_opt.take().expect("Commit already used");
                                // Ack/nack the batch to fill the sequencer slot.
                                let dispositions =
                                    map_responses_to_dispositions(&message_ids, responses, &failed, &request_ids);
                                if let Err(commit_err) = commit(dispositions).await {
                                    warn!("Commit after transient failure also failed: {}", commit_err);
                                }
                                if !has_retry_middleware {
                                    match err_tx.try_send(e) {
                                        Ok(_) => trace!("Reported error to main task"),
                                        Err(err_send) => warn!(error=?err_send, "Could not send error to main task, it might be down or busy."),
                                    }
                                    break;
                                }
                                warn!("Transient error in batch, message(s) Nack'ed for re-delivery: {}", e);
                                tokio::task::yield_now().await;
                                continue;
                            }
                            // Only non-retryable errors remain - drop them with a log.
                            for (msg, e) in &failed {
                                error!("Worker dropping message (ID: {:032x}) due to non-retryable error: {}", msg.message_id, e);
                            }
                            let commit = commit_opt.take().expect("Commit already used");
                            let err_tx = err_tx.clone();
                            let ids = std::mem::take(&mut message_ids);
                            let req_ids = request_ids;
                            commit_tasks.spawn(async move {
                                let dispositions = map_responses_to_dispositions(&ids, responses, &failed, &req_ids);
                                if let Err(e) = commit(dispositions).await {
                                    error!("Commit failed: {}", e);
                                    match err_tx.try_send(e) {
                                        Ok(_) => trace!("Reported commit error to main task"),
                                        Err(err_send) => warn!(error=?err_send, "Could not send commit error to main task, it might be down or busy."),
                                    }
                                }
                            });
                        }
                        Err(e) => {
                            error!("Worker failed to send message batch: {}", e);
                            let commit = commit_opt.take().expect("Commit already used");
                            // Nack the commit to fill the sequencer slot and prevent a deadlock.
                            let nack_result = commit(vec![MessageDisposition::Nack; batch_len]).await;
                            debug!("Nack commit result: {:?}", nack_result);
                            // Send the error back to the main task to tear down the route.
                            match err_tx.try_send(e.into()) {
                                Ok(_) => trace!("Reported error to main task"),
                                Err(err_send) => warn!(error=?err_send, "Could not send error to main task, it might be down or busy."),
                            }
                            break;
                        }
                    }
                }
                // Wait for all in-flight commits to complete
                while commit_tasks.join_next().await.is_some() {}
            });
        }

        let mut seq_counter = 0u64;
        // Holds an error that caused the loop to break, to be returned after graceful shutdown.
        let mut loop_error: Option<anyhow::Error> = None;
        loop {
            select! {
                biased; // Prioritize checking for errors

                Ok(err) = err_rx.recv() => {
                    error!("A worker reported a critical error. Shutting down route.");
                    loop_error = Some(err);
                    break;
                }

                Some(res) = join_set.join_next() => {
                    match res {
                        Ok(_) => {
                            error!("A worker task finished unexpectedly. Shutting down route.");
                            loop_error = Some(anyhow::anyhow!("Worker task finished unexpectedly"));
                        }
                        Err(e) => {
                            error!("A worker task panicked: {}. Shutting down route.", e);
                            loop_error = Some(e.into());
                        }
                    }
                    break;
                }

                _ = shutdown_rx.recv() => {
                    info!("Shutdown signal received in concurrent runner for route '{}'.", name);
                    break;
                }

                res = consumer.receive_batch(self.options.batch_size) => {
                    let (messages, commit) = match res {
                        Ok(batch) => {
                            if batch.messages.is_empty() {
                                continue; // No messages, loop to select! again
                            }
                            (batch.messages, batch.commit)
                        }
                        Err(ConsumerError::EndOfStream) => {
                            info!("Consumer for route '{}' reached end of stream. Shutting down.", name);
                            break; // Graceful exit
                        }
                        Err(ConsumerError::Connection(e)) => {
                            // Propagate error to trigger reconnect by the outer loop
                            loop_error = Some(e);
                            break;
                        }
                        Err(ConsumerError::Gap { requested, base }) => {
                            // Propagate gap error to trigger reconnect by the outer loop
                            loop_error = Some(ConsumerError::Gap { requested, base }.into());
                            break;
                        }
                    };
                    debug!("Received a batch of {} messages concurrently", messages.len());

                    // Wrap the commit function to route it through the sequencer.
                    // Only advance the sequence counter after we've successfully enqueued
                    // the work item to avoid creating sequence gaps if the work channel
                    // is closed while producing batches.
                    let seq = seq_counter;
                    let wrapped_commit = wrap_commit(commit, seq, seq_tx.clone());

                    match work_tx.send((messages, wrapped_commit)).await {
                        Ok(()) => {
                            seq_counter += 1;
                        }
                        Err(e) => {
                            warn!("Work channel closed, cannot process more messages concurrently. Shutting down.");
                            // Recover the moved tuple so we can invoke the wrapped commit
                            // and resolve the batch with a NACK.
                            let (msgs_back, wrapped_commit_back) = e.into_inner();
                            let _ = (wrapped_commit_back)(vec![crate::traits::MessageDisposition::Nack; msgs_back.len()]).await;
                            break;
                        }
                    }

                    tokio::task::yield_now().await;
                }
            }
        }

        // --- Graceful Shutdown ---
        // Close the work channel so workers drain their current messages and exit the loop.
        // This applies on both normal shutdown AND error paths, ensuring in-flight commits
        // are not aborted mid-sequence.
        drop(work_tx);
        // Wait for all worker tasks to complete.
        while join_set.join_next().await.is_some() {}

        // Close sequencer
        drop(seq_tx);
        let _ = sequencer_handle.await;
        run_consumer_disconnect_hook(name, consumer.as_ref()).await;
        run_publisher_disconnect_hook(name, &publisher).await;

        if let Some(err) = loop_error {
            return Err(err);
        }

        if let Ok(err) = err_rx.try_recv() {
            return Err(err);
        }

        // Return true if shutdown was requested (channel is empty means it was closed/consumed),
        // false if we reached end-of-stream naturally.
        Ok(shutdown_rx.is_empty())
    }

    pub fn with_options(mut self, options: RouteOptions) -> Self {
        self.options = options;
        self
    }
    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
        self.options.concurrency = concurrency.max(1);
        self
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.options.batch_size = batch_size.max(1);
        self
    }
    pub fn with_commit_concurrency_limit(mut self, limit: usize) -> Self {
        self.options.commit_concurrency_limit = limit.max(1);
        self
    }

    pub fn with_handler(mut self, handler: impl Handler + 'static) -> Self {
        self.output.handler = Some(Arc::new(handler));
        self
    }

    /// Registers a typed handler for the route.
    ///
    /// The handler can accept either:
    /// - `fn(T) -> Future<Output = Result<Handled, HandlerError>>`
    /// - `fn(T, MessageContext) -> Future<Output = Result<Handled, HandlerError>>`
    ///
    /// # Examples
    ///
    /// ```
    /// # use mq_bridge::{Route, models::Endpoint};
    /// # use serde::Deserialize;
    ///
    /// #[derive(Deserialize)]
    /// struct MyData { id: u32 }
    ///
    /// async fn my_handler(data: MyData) -> anyhow::Result<()> {
    ///     Ok(())
    /// }
    ///
    /// let route = Route::new(Endpoint::new_memory("in", 10), Endpoint::new_memory("out", 10))
    ///     .add_handler("my_type", my_handler);
    /// ```
    pub fn add_handler<T, H, Args>(mut self, type_name: &str, handler: H) -> Self
    where
        T: DeserializeOwned + Send + Sync + 'static,
        H: crate::type_handler::IntoTypedHandler<T, Args>,
        Args: Send + Sync + 'static,
    {
        // Create the wrapper closure that handles deserialization and context extraction
        let handler = Arc::new(handler);
        let wrapper = move |msg: crate::CanonicalMessage| {
            let handler = handler.clone();
            async move {
                let data = msg.parse::<T>().map_err(|e| {
                    HandlerError::NonRetryable(anyhow::anyhow!("Deserialization failed: {}", e))
                })?;
                let ctx = crate::MessageContext::from(msg);
                handler.call(data, ctx).await
            }
        };
        let wrapper = Arc::new(wrapper);

        let prev_handler = self.output.handler.take();

        let new_handler = if let Some(h) = prev_handler {
            if let Some(extended) = h.register_handler(type_name, wrapper.clone()) {
                extended
            } else {
                Arc::new(
                    crate::type_handler::TypeHandler::new()
                        .with_fallback(h)
                        .add_handler(type_name, wrapper),
                )
            }
        } else {
            Arc::new(crate::type_handler::TypeHandler::new().add_handler(type_name, wrapper))
        };

        self.output.handler = Some(new_handler);
        self
    }
    pub fn add_handlers<T, H, Args>(mut self, handlers: HashMap<&str, H>) -> Self
    where
        T: DeserializeOwned + Send + Sync + 'static,
        H: crate::type_handler::IntoTypedHandler<T, Args>,
        Args: Send + Sync + 'static,
    {
        for (type_name, handler) in handlers {
            self = self.add_handler(type_name, handler);
        }
        self
    }
}

type SequencerItem = (
    Vec<MessageDisposition>,
    BatchCommitFunc,
    tokio::sync::oneshot::Sender<anyhow::Result<()>>,
);

fn spawn_sequencer(buffer_size: usize) -> (Sender<(u64, SequencerItem)>, JoinHandle<()>) {
    let (seq_tx, seq_rx) = bounded::<(u64, SequencerItem)>(buffer_size);
    let sequencer_handle = tokio::spawn(async move {
        let mut buffer: BTreeMap<u64, SequencerItem> = BTreeMap::new();
        let mut next_seq = 0u64;

        loop {
            // If we have the next item in sequence, execute its commit directly.
            // Using a plain await (no select!) here is essential: if we raced a recv
            // against the commit future, a recv win would drop the commit future and
            // the notify sender, leaving the caller permanently blocked while next_seq
            // stays unadvanced — a deadlock.
            if let Some((dispositions, commit_func, notify)) = buffer.remove(&next_seq) {
                let result = commit_func(dispositions).await;
                let _ = notify.send(result);
                next_seq += 1;
                // Yield to allow other tasks to run, preventing busy-loop when buffer has many messages
                tokio::task::yield_now().await;
                continue;
            }

            // Wait for the next item from any worker.
            match seq_rx.recv().await {
                Ok((seq, item)) => {
                    if seq < next_seq {
                        let (_, _, notify) = item;
                        trace!(
                            seq,
                            next_seq,
                            "Sequencer received late item (seq < next_seq)"
                        );
                        let _ = notify.send(Err(anyhow::anyhow!(
                            "Sequencer received late item (seq {} < next_seq {})",
                            seq,
                            next_seq
                        )));
                    } else {
                        buffer.insert(seq, item);
                    }
                }
                Err(_) => {
                    // seq_tx was dropped — drain and notify any remaining buffered commits.
                    for (_, (_, _, notify)) in buffer {
                        let _ = notify.send(Err(anyhow::anyhow!("Sequencer is shutting down")));
                    }
                    break;
                }
            }
        }
    });
    (seq_tx, sequencer_handle)
}

fn wrap_commit(
    commit: BatchCommitFunc,
    seq: u64,
    seq_tx: Sender<(u64, SequencerItem)>,
) -> BatchCommitFunc {
    Box::new(move |dispositions| {
        Box::pin(async move {
            let (notify_tx, notify_rx) = tokio::sync::oneshot::channel();
            if seq_tx
                .send((seq, (dispositions, commit, notify_tx)))
                .await
                .is_ok()
            {
                match notify_rx.await {
                    Ok(res) => res,
                    Err(_) => Err(anyhow::anyhow!(
                        "Sequencer dropped the commit channel unexpectedly"
                    )),
                }
            } else {
                Err(anyhow::anyhow!(
                    "Failed to send commit to sequencer, route is likely shutting down"
                ))
            }
        })
    })
}

fn map_responses_to_dispositions(
    message_ids: &[u128],
    responses: Option<Vec<crate::CanonicalMessage>>,
    failed: &[(crate::CanonicalMessage, PublisherError)],
    request_ids: &std::collections::HashSet<u128>,
) -> Vec<MessageDisposition> {
    let len = message_ids.len();
    if responses.is_none() && failed.is_empty() && (len == 0 || request_ids.is_empty()) {
        return vec![MessageDisposition::Ack; len];
    }

    // Fast path for single message batches (very common for high-concurrency low-batch setups)
    if len == 1 {
        let id = message_ids[0];
        if !failed.is_empty() {
            return vec![MessageDisposition::Nack];
        }
        if let Some(mut resps) = responses {
            if let Some(resp) = resps.pop() {
                return vec![MessageDisposition::Reply(resp)];
            }
        }
        if request_ids.contains(&id) {
            error!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Nacking to avoid committing a lost response.", id);
            return vec![MessageDisposition::Nack];
        }
        return vec![MessageDisposition::Ack];
    }

    let mut dispositions = Vec::with_capacity(len);
    // Build failed_ids manually to avoid collect() overhead
    let mut failed_ids = std::collections::HashSet::with_capacity(failed.len());
    for (m, _) in failed {
        failed_ids.insert(m.message_id);
    }

    // Create a map from message_id to response message for efficient lookup.
    let mut response_map: std::collections::HashMap<u128, crate::CanonicalMessage> = responses
        .unwrap_or_default()
        .into_iter()
        .map(|r| (r.message_id, r))
        .collect();

    for id in message_ids {
        if failed_ids.contains(id) {
            dispositions.push(MessageDisposition::Nack);
        } else if let Some(resp) = response_map.remove(id) {
            // If a response exists for this specific ID, use it.
            dispositions.push(MessageDisposition::Reply(resp));
        } else if request_ids.contains(id) {
            error!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Nacking to avoid committing a lost response.", id);
            dispositions.push(MessageDisposition::Nack);
        } else {
            // Otherwise, it was a successful send that did not produce a response.
            dispositions.push(MessageDisposition::Ack);
        }
    }
    dispositions
}

#[cfg(test)]
fn test_map_responses_to_dispositions_logic() {
    use crate::{traits::PublisherError, CanonicalMessage};
    use anyhow::anyhow;

    let ids = vec![1, 2, 3, 4];

    let mut resp1 = CanonicalMessage::from("resp1");
    resp1.message_id = 1;
    let mut resp4 = CanonicalMessage::from("resp4");
    resp4.message_id = 4;

    let responses = Some(vec![
        resp1, // Corresponds to id 1
        resp4, // Corresponds to id 4
    ]);

    let mut msg2 = CanonicalMessage::from("msg2");
    msg2.message_id = 2;
    let failed = vec![(msg2, PublisherError::NonRetryable(anyhow!("failed")))];

    let mut request_ids = std::collections::HashSet::new();
    request_ids.insert(3); // id 3 expects a reply but won't get one
    let dispositions = map_responses_to_dispositions(&ids, responses, &failed, &request_ids);

    assert_eq!(dispositions.len(), 4);
    assert!(matches!(dispositions[0], MessageDisposition::Reply(_))); // from responses
    assert!(matches!(dispositions[1], MessageDisposition::Nack)); // from failed
    assert!(matches!(dispositions[2], MessageDisposition::Nack)); // missing reply
    assert!(matches!(dispositions[3], MessageDisposition::Reply(_))); // from responses
}

pub fn get_route(name: &str) -> Option<Route> {
    Route::get(name)
}

pub fn list_routes() -> Vec<String> {
    Route::list()
}

pub async fn stop_route(name: &str) -> bool {
    Route::stop(name).await
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::{Endpoint, EndpointType, FaultMode, Middleware, RandomPanicMiddleware};
    use crate::traits::{
        CustomMiddlewareFactory, MessageConsumer, MessagePublisher, ReceivedBatch,
    };
    use crate::CanonicalMessage;
    use std::any::Any;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use std::sync::Arc;
    use std::time::Duration;

    #[derive(Debug, Default)]
    struct CommitObservation {
        completed: Mutex<Vec<u64>>,
        active: std::sync::atomic::AtomicUsize,
        max_active: std::sync::atomic::AtomicUsize,
    }

    #[derive(Debug)]
    struct CommitTrackingMiddlewareFactory {
        observation: Arc<CommitObservation>,
    }

    #[derive(Debug)]
    struct ReorderingPublisherMiddlewareFactory;

    struct CommitTrackingConsumer {
        inner: Box<dyn MessageConsumer>,
        observation: Arc<CommitObservation>,
    }

    struct ReorderingPublisher {
        inner: Box<dyn MessagePublisher>,
    }

    #[async_trait::async_trait]
    impl CustomMiddlewareFactory for CommitTrackingMiddlewareFactory {
        async fn apply_consumer(
            &self,
            consumer: Box<dyn MessageConsumer>,
            _route_name: &str,
            _config: &serde_json::Value,
        ) -> anyhow::Result<Box<dyn MessageConsumer>> {
            Ok(Box::new(CommitTrackingConsumer {
                inner: consumer,
                observation: Arc::clone(&self.observation),
            }))
        }
    }

    #[async_trait::async_trait]
    impl CustomMiddlewareFactory for ReorderingPublisherMiddlewareFactory {
        async fn apply_publisher(
            &self,
            publisher: Box<dyn MessagePublisher>,
            _route_name: &str,
            _config: &serde_json::Value,
        ) -> anyhow::Result<Box<dyn MessagePublisher>> {
            Ok(Box::new(ReorderingPublisher { inner: publisher }))
        }
    }

    #[async_trait::async_trait]
    impl MessageConsumer for CommitTrackingConsumer {
        async fn receive_batch(
            &mut self,
            max_messages: usize,
        ) -> Result<ReceivedBatch, ConsumerError> {
            let mut batch = self.inner.receive_batch(max_messages).await?;
            let seq = batch
                .messages
                .first()
                .and_then(|message| message.get_payload_str().parse::<u64>().ok())
                .expect("tracking test expects numeric payloads");
            let original_commit = batch.commit;
            let observation = Arc::clone(&self.observation);
            batch.commit = Box::new(move |dispositions| {
                let observation = Arc::clone(&observation);
                Box::pin(async move {
                    let active_now = observation.active.fetch_add(1, Ordering::SeqCst) + 1;
                    let _ = observation.max_active.fetch_update(
                        Ordering::SeqCst,
                        Ordering::SeqCst,
                        |current| (active_now > current).then_some(active_now),
                    );

                    tokio::time::sleep(Duration::from_millis(20)).await;
                    let result = original_commit(dispositions).await;
                    observation.completed.lock().unwrap().push(seq);
                    observation.active.fetch_sub(1, Ordering::SeqCst);
                    result
                })
            });
            Ok(batch)
        }

        fn as_any(&self) -> &dyn Any {
            self
        }
    }

    #[async_trait::async_trait]
    impl MessagePublisher for ReorderingPublisher {
        async fn send_batch(
            &self,
            messages: Vec<crate::CanonicalMessage>,
        ) -> Result<SentBatch, PublisherError> {
            let seq = messages
                .first()
                .and_then(|message| message.get_payload_str().parse::<u64>().ok())
                .expect("tracking test expects numeric payloads");
            let delay_ms = 10 * (6u64.saturating_sub(seq.min(6)));
            tokio::time::sleep(Duration::from_millis(delay_ms)).await;
            self.inner.send_batch(messages).await
        }

        async fn send(&self, msg: crate::CanonicalMessage) -> Result<Sent, PublisherError> {
            self.inner.send(msg).await
        }

        async fn flush(&self) -> anyhow::Result<()> {
            self.inner.flush().await
        }

        fn as_any(&self) -> &dyn Any {
            self
        }
    }

    async fn assert_route_commits_are_ordered_and_non_overlapping(concurrency: usize) {
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let tracking_name = format!("track_commit_{}", unique_id);
        let reorder_name = format!("reorder_publish_{}", unique_id);
        let in_topic = format!("ordered_commit_in_{}", unique_id);
        let observation = Arc::new(CommitObservation::default());

        register_middleware_factory(
            &tracking_name,
            Arc::new(CommitTrackingMiddlewareFactory {
                observation: Arc::clone(&observation),
            }),
        );
        register_middleware_factory(
            &reorder_name,
            Arc::new(ReorderingPublisherMiddlewareFactory),
        );

        let input = Endpoint::new_memory(&in_topic, 32).add_middleware(Middleware::Custom {
            name: tracking_name,
            config: serde_json::Value::Null,
        });
        let output = Endpoint::new(EndpointType::Null).add_middleware(Middleware::Custom {
            name: reorder_name,
            config: serde_json::Value::Null,
        });

        let route = Route::new(input.clone(), output)
            .with_concurrency(concurrency)
            .with_batch_size(1)
            .with_commit_concurrency_limit(1);

        let input_channel = input.channel().unwrap();
        let messages = (0..6)
            .map(|seq| crate::CanonicalMessage::from(seq.to_string()))
            .collect();
        input_channel.fill_messages(messages).await.unwrap();
        input_channel.close();

        tokio::time::timeout(
            std::time::Duration::from_secs(5),
            route.run_until_err("ordered_commit_regression", None, None),
        )
        .await
        .expect("Route should not hang while draining finite input")
        .expect("Route should complete without commit errors");
        assert_eq!(
            *observation.completed.lock().unwrap(),
            vec![0, 1, 2, 3, 4, 5],
            "Commit execution must follow receive order",
        );
        assert_eq!(
            observation.max_active.load(Ordering::SeqCst),
            1,
            "Broker-facing commit functions must never overlap",
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_sequential_route_commits_are_ordered_and_non_overlapping() {
        assert_route_commits_are_ordered_and_non_overlapping(1).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_concurrent_route_commits_are_ordered_and_non_overlapping() {
        assert_route_commits_are_ordered_and_non_overlapping(4).await;
    }

    // Helper function to run a fault injection test on the consumer side.
    async fn run_consumer_fault_test(
        mode: FaultMode,
        expected_payload: &str,
        route_should_restart: bool,
        concurrency: usize,
    ) {
        let unique_suffix = fast_uuid_v7::gen_id().to_string();
        let in_topic = format!("fault_in_{}_{}_{}", mode, concurrency, unique_suffix);
        let out_topic = format!("fault_out_{}_{}_{}", mode, concurrency, unique_suffix);

        let fault_config = RandomPanicMiddleware {
            mode,
            trigger_on_message: Some(1), // Panic on the first message
            enabled: true,
            ..Default::default()
        };

        let input = Endpoint::new_memory(&in_topic, 10)
            .add_middleware(Middleware::RandomPanic(fault_config));
        let output = Endpoint::new_memory(&out_topic, 10);

        let route_name = format!("fault_test_{}_{}", mode, concurrency);
        let route = Route::new(input.clone(), output.clone()).with_concurrency(concurrency);

        // Start the route
        route
            .deploy(&route_name)
            .await
            .expect("Failed to deploy route");
        // Send a message. The consumer will inject a fault when it tries to receive it.
        let input_ch = input.channel().unwrap();
        input_ch
            .send_message("persistent_msg".into())
            .await
            .unwrap();

        if route_should_restart {
            // The route's worker will fail, and the supervisor will wait 5 seconds before restarting.
            // We wait for a bit longer than that to ensure recovery has happened.
            tokio::time::sleep(std::time::Duration::from_secs(6)).await;
        } else {
            // Route doesn't restart, just wait a bit for the (faulty) message to pass through.
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
        }

        // Verify the outcome.
        let mut verifier = route.connect_to_output("verifier").await.unwrap();
        let received = tokio::time::timeout(std::time::Duration::from_secs(10), verifier.receive())
            .await
            .expect("Timed out waiting for message after fault")
            .expect("Stream closed while waiting for message");

        assert_eq!(received.message.get_payload_str(), expected_payload);
        (received.commit)(MessageDisposition::Ack).await.unwrap();

        // Cleanup
        Route::stop(&route_name).await;
    }

    // Helper function to run a fault injection test on the publisher side.
    async fn run_publisher_fault_test(
        mode: FaultMode,
        expected_payload: &str,
        route_should_restart: bool,
    ) {
        let unique_suffix = fast_uuid_v7::gen_id().to_string();
        let in_topic = format!("pub_fault_in_{}_{}", mode, unique_suffix);
        let out_topic = format!("pub_fault_out_{}_{}", mode, unique_suffix);

        let fault_config = RandomPanicMiddleware {
            mode,
            trigger_on_message: Some(1), // Trigger on the first message
            enabled: true,
            ..Default::default()
        };

        let mut input = Endpoint::new_memory(&in_topic, 10);
        // Enable NACK on input so messages aren't lost when publisher crashes
        if let EndpointType::Memory(ref mut cfg) = input.endpoint_type {
            cfg.enable_nack = true;
        }
        // Apply fault middleware to output
        let output = Endpoint::new_memory(&out_topic, 10)
            .add_middleware(Middleware::RandomPanic(fault_config));

        let route_name = format!("pub_fault_test_{}", mode);
        let route = Route::new(input.clone(), output.clone());

        route
            .deploy(&route_name)
            .await
            .expect("Failed to deploy route");

        let input_ch = input.channel().unwrap();
        input_ch
            .send_message(expected_payload.into())
            .await
            .unwrap();

        if route_should_restart {
            tokio::time::sleep(std::time::Duration::from_secs(6)).await;
        } else {
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
        }

        let mut verifier = route.connect_to_output("verifier").await.unwrap();
        let received = tokio::time::timeout(std::time::Duration::from_secs(10), verifier.receive())
            .await
            .expect("Timed out waiting for message after publisher fault")
            .expect("Stream closed");

        assert_eq!(received.message.get_payload_str(), expected_payload);
        (received.commit)(MessageDisposition::Ack).await.unwrap();

        Route::stop(&route_name).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[ignore = "Takes too much time for regular tests"]
    async fn test_route_recovery_from_faults() {
        let original_payload = "persistent_msg";

        // Test with concurrency > 1
        run_consumer_fault_test(FaultMode::Panic, original_payload, true, 2).await;
        run_consumer_fault_test(FaultMode::Disconnect, original_payload, true, 2).await;
        run_consumer_fault_test(FaultMode::Timeout, original_payload, true, 2).await;
        run_consumer_fault_test(FaultMode::Nack, original_payload, true, 2).await;

        // This fault replaces the message but does not restart the route.
        run_consumer_fault_test(FaultMode::JsonFormatError, "{invalid json}", false, 2).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[ignore = "Takes too much time for regular tests"]
    async fn test_route_recovery_from_faults_sequential() {
        let original_payload = "persistent_msg";

        // Test with concurrency = 1
        run_consumer_fault_test(FaultMode::Panic, original_payload, true, 1).await;
        run_consumer_fault_test(FaultMode::Disconnect, original_payload, true, 1).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[ignore = "Takes too much time for regular tests"]
    async fn test_publisher_recovery_from_faults() {
        let original_payload = "persistent_msg";
        // Test publisher-side faults causing restart/retry.
        // `FaultMode::Panic` is not tested here because the `MemoryConsumer` used for input
        // does not support crash-safe at-least-once delivery. A panic in the publisher
        // worker would cause the in-flight message to be lost.
        run_publisher_fault_test(FaultMode::Disconnect, original_payload, true).await;
        run_publisher_fault_test(FaultMode::Timeout, original_payload, true).await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_route_sequencer_deadlock_fix() {
        // This test ensures that when a worker fails to send a batch (and thus drops the commit handle),
        // the sequencer doesn't deadlock waiting for that sequence number.
        // The fix ensures that even on failure, the commit function is called (with Nack) to fill the sequence gap.

        let unique_id = fast_uuid_v7::gen_id().to_string();
        let factory_name = format!("fail_factory_{}", unique_id);
        let in_topic = format!("deadlock_in_{}", unique_id);
        let out_topic = format!("deadlock_out_{}", unique_id);

        #[derive(Debug)]
        struct FailingMiddlewareFactory {
            fail_flag: Arc<AtomicBool>,
        }

        #[async_trait::async_trait]
        impl CustomMiddlewareFactory for FailingMiddlewareFactory {
            async fn apply_publisher(
                &self,
                publisher: Box<dyn MessagePublisher>,
                _route_name: &str,
                _config: &serde_json::Value,
            ) -> anyhow::Result<Box<dyn MessagePublisher>> {
                Ok(Box::new(FailingPublisher {
                    inner: publisher,
                    fail_flag: self.fail_flag.clone(),
                }))
            }
            async fn apply_consumer(
                &self,
                consumer: Box<dyn MessageConsumer>,
                _route_name: &str,
                _config: &serde_json::Value,
            ) -> anyhow::Result<Box<dyn MessageConsumer>> {
                Ok(consumer)
            }
        }

        struct FailingPublisher {
            inner: Box<dyn MessagePublisher>,
            fail_flag: Arc<AtomicBool>,
        }

        #[async_trait::async_trait]
        impl MessagePublisher for FailingPublisher {
            async fn send_batch(
                &self,
                messages: Vec<crate::CanonicalMessage>,
            ) -> Result<SentBatch, PublisherError> {
                // We want to fail one batch to trigger the error path in the worker.
                // We use compare_exchange to ensure only one failure happens.
                if self
                    .fail_flag
                    .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
                    .is_ok()
                {
                    return Err(PublisherError::Retryable(anyhow::anyhow!(
                        "Simulated failure"
                    )));
                }
                // Add a small delay for successful batches to ensure the failed one (if it created a gap)
                // would block the sequencer if the gap wasn't filled.
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                self.inner.send_batch(messages).await
            }
            async fn send(
                &self,
                msg: crate::CanonicalMessage,
            ) -> Result<crate::traits::Sent, PublisherError> {
                self.inner.send(msg).await
            }
            async fn flush(&self) -> anyhow::Result<()> {
                self.inner.flush().await
            }
            fn as_any(&self) -> &dyn Any {
                self
            }
        }

        let fail_flag = Arc::new(AtomicBool::new(true));
        register_middleware_factory(
            &factory_name,
            Arc::new(FailingMiddlewareFactory {
                fail_flag: fail_flag.clone(),
            }),
        );

        let input = Endpoint::new_memory(&in_topic, 100);
        let output = Endpoint::new_memory(&out_topic, 100).add_middleware(Middleware::Custom {
            name: factory_name,
            config: serde_json::Value::Null,
        });

        // Concurrency > 1 is required to have multiple workers and potential out-of-order completion
        let route = Route::new(input.clone(), output.clone())
            .with_concurrency(2)
            .with_batch_size(1);

        // Send messages
        let input_ch = input.channel().unwrap();
        input_ch.send_message("msg1".into()).await.unwrap();
        input_ch.send_message("msg2".into()).await.unwrap();
        input_ch.send_message("msg3".into()).await.unwrap();

        // Run the route. It should fail eventually due to the simulated error,
        // but it MUST NOT deadlock.
        let run_fut = async {
            let (_shutdown_tx, shutdown_rx) = async_channel::bounded(1);
            route
                .run_until_err("deadlock_test", Some(shutdown_rx), None)
                .await
        };

        // If deadlock exists, this timeout will trigger.
        let result = tokio::time::timeout(std::time::Duration::from_secs(5), run_fut).await;

        match result {
            Ok(res) => {
                // We expect an error because the publisher returns Err.
                assert!(
                    res.is_err(),
                    "Route should have failed with simulated error"
                );
            }
            Err(_) => {
                panic!("Route deadlocked! The sequencer likely didn't receive the Nack for the failed batch.");
            }
        }
    }

    #[tokio::test]
    async fn test_sequencer_ordered_commits() {
        use std::time::Duration;
        use tokio::time::timeout;

        let (seq_tx, sequencer_handle) = spawn_sequencer(16);
        let processed: Arc<Mutex<Vec<u64>>> = Arc::new(Mutex::new(Vec::new()));

        // Send sequences out of order to ensure the sequencer enforces ordering.
        let seqs = [2u64, 0u64, 1u64, 3u64];
        let mut receivers = Vec::new();

        for seq in seqs.iter().cloned() {
            let (notify_tx, notify_rx) = tokio::sync::oneshot::channel();
            let processed_clone = processed.clone();
            let commit: BatchCommitFunc = Box::new(move |_dispositions| {
                let processed = processed_clone.clone();
                Box::pin(async move {
                    // Simulate variable work durations
                    tokio::time::sleep(Duration::from_millis(10 * seq)).await;
                    processed.lock().unwrap().push(seq);
                    Ok(())
                })
            });
            seq_tx
                .send((seq, (Vec::new(), commit, notify_tx)))
                .await
                .unwrap();
            receivers.push(notify_rx);
        }

        // Wait for all commits to complete (with timeout to catch deadlocks)
        for rx in receivers {
            let res = timeout(Duration::from_secs(2), rx)
                .await
                .expect("Sequencer notify timed out");
            assert!(res.is_ok(), "Sequencer reported an error on commit");
            assert!(res.unwrap().is_ok(), "Commit returned an error");
        }

        // Close sender to allow sequencer task to exit and await it.
        drop(seq_tx);
        let _ = sequencer_handle.await;

        let result = processed.lock().unwrap().clone();
        assert_eq!(
            result,
            vec![0u64, 1u64, 2u64, 3u64],
            "Sequencer must process commits in order"
        );
    }

    #[tokio::test]
    async fn test_sequencer_shutdown_notifies_pending() {
        use std::time::Duration;
        use tokio::time::timeout;

        let (seq_tx, sequencer_handle) = spawn_sequencer(8);

        // Prepare two pending items for sequences 1 and 2 while sequence 0 is missing.
        let (notify_tx1, notify_rx1) = tokio::sync::oneshot::channel();
        let (notify_tx2, notify_rx2) = tokio::sync::oneshot::channel();

        let commit1: BatchCommitFunc = Box::new(|_dispositions| {
            Box::pin(async move {
                // Should not be executed because next_seq is missing (0)
                panic!("Commit should not be executed during shutdown drain");
                #[allow(unreachable_code)]
                Ok(())
            })
        });

        let commit2: BatchCommitFunc = Box::new(|_dispositions| {
            Box::pin(async move {
                panic!("Commit should not be executed during shutdown drain");
                #[allow(unreachable_code)]
                Ok(())
            })
        });

        seq_tx
            .send((1u64, (Vec::new(), commit1, notify_tx1)))
            .await
            .unwrap();
        seq_tx
            .send((2u64, (Vec::new(), commit2, notify_tx2)))
            .await
            .unwrap();

        // Trigger shutdown of the sequencer by dropping the sender.
        drop(seq_tx);

        // Sequencer should drain buffered items and reply with an error to the notifiers.
        let r1 = timeout(Duration::from_secs(1), notify_rx1)
            .await
            .expect("Timeout waiting for notify_rx1")
            .expect("Sequencer closed notify channel");
        assert!(
            r1.is_err(),
            "Pending commit should receive Err on sequencer shutdown"
        );

        let r2 = timeout(Duration::from_secs(1), notify_rx2)
            .await
            .expect("Timeout waiting for notify_rx2")
            .expect("Sequencer closed notify channel");
        assert!(
            r2.is_err(),
            "Pending commit should receive Err on sequencer shutdown"
        );

        let _ = sequencer_handle.await;
    }

    use crate::traits::{BoxFuture, CustomEndpointFactory, Sent};
    use std::sync::Mutex;

    type ConsumerBehavior =
        Arc<Mutex<dyn FnMut() -> Result<Box<dyn MessageConsumer>, anyhow::Error> + Send + Sync>>;
    type PublisherBehavior =
        Arc<Mutex<dyn FnMut() -> Result<Box<dyn MessagePublisher>, anyhow::Error> + Send + Sync>>;

    struct MockEndpointFactory {
        create_consumer_fail: bool,
        consumer_behavior: ConsumerBehavior,
        publisher_behavior: PublisherBehavior,
    }

    impl std::fmt::Debug for MockEndpointFactory {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("MockEndpointFactory")
                .field("create_consumer_fail", &self.create_consumer_fail)
                .finish()
        }
    }

    impl MockEndpointFactory {
        fn new() -> Self {
            Self {
                create_consumer_fail: false,
                consumer_behavior: Arc::new(Mutex::new(|| Err(anyhow::anyhow!("Not implemented")))),
                publisher_behavior: Arc::new(Mutex::new(|| {
                    Ok(Box::new(NoOpPublisher) as Box<dyn MessagePublisher>)
                })),
            }
        }
    }

    #[derive(Clone)]
    struct NoOpPublisher;
    #[async_trait::async_trait]
    impl MessagePublisher for NoOpPublisher {
        async fn send_batch(
            &self,
            _: Vec<crate::CanonicalMessage>,
        ) -> Result<SentBatch, PublisherError> {
            Ok(SentBatch::Ack)
        }
        async fn send(&self, _: crate::CanonicalMessage) -> Result<Sent, PublisherError> {
            Ok(Sent::Ack)
        }
        fn as_any(&self) -> &dyn Any {
            self
        }
    }

    #[async_trait::async_trait]
    impl CustomEndpointFactory for MockEndpointFactory {
        async fn create_consumer(
            &self,
            _: &str,
            _: &serde_json::Value,
        ) -> anyhow::Result<Box<dyn MessageConsumer>> {
            if self.create_consumer_fail {
                return Err(anyhow::anyhow!("Endpoint unavailable"));
            }
            (self.consumer_behavior.lock().unwrap())()
        }
        async fn create_publisher(
            &self,
            _: &str,
            _: &serde_json::Value,
        ) -> anyhow::Result<Box<dyn MessagePublisher>> {
            (self.publisher_behavior.lock().unwrap())()
        }
    }

    #[derive(Clone, Default)]
    struct HookState {
        consumer_connects: Arc<AtomicUsize>,
        consumer_disconnects: Arc<AtomicUsize>,
        publisher_connects: Arc<AtomicUsize>,
        publisher_disconnects: Arc<AtomicUsize>,
        shared_mutations: Arc<AtomicUsize>,
        fail_consumer_connect: Arc<AtomicBool>,
        fail_consumer_disconnect: Arc<AtomicBool>,
        fail_publisher_disconnect: Arc<AtomicBool>,
    }

    struct HookConsumer {
        state: HookState,
    }

    struct HookPublisher {
        state: HookState,
    }

    #[async_trait::async_trait]
    impl MessageConsumer for HookConsumer {
        fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
            Some(Box::pin(async move {
                self.state.consumer_connects.fetch_add(1, Ordering::SeqCst);
                self.state.shared_mutations.fetch_add(1, Ordering::SeqCst);
                if self.state.fail_consumer_connect.load(Ordering::SeqCst) {
                    return Err(anyhow::anyhow!("consumer hook failed"));
                }
                Ok(())
            }))
        }

        fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
            Some(Box::pin(async move {
                self.state
                    .consumer_disconnects
                    .fetch_add(1, Ordering::SeqCst);
                if self.state.fail_consumer_disconnect.load(Ordering::SeqCst) {
                    return Err(anyhow::anyhow!("consumer disconnect hook failed"));
                }
                Ok(())
            }))
        }

        async fn receive_batch(&mut self, _max: usize) -> Result<ReceivedBatch, ConsumerError> {
            Err(ConsumerError::EndOfStream)
        }

        fn as_any(&self) -> &dyn Any {
            self
        }
    }

    #[async_trait::async_trait]
    impl MessagePublisher for HookPublisher {
        fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
            Some(Box::pin(async move {
                self.state.publisher_connects.fetch_add(1, Ordering::SeqCst);
                self.state.shared_mutations.fetch_add(1, Ordering::SeqCst);
                Ok(())
            }))
        }

        fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
            Some(Box::pin(async move {
                self.state
                    .publisher_disconnects
                    .fetch_add(1, Ordering::SeqCst);
                if self.state.fail_publisher_disconnect.load(Ordering::SeqCst) {
                    return Err(anyhow::anyhow!("publisher disconnect hook failed"));
                }
                Ok(())
            }))
        }

        async fn send_batch(
            &self,
            _: Vec<crate::CanonicalMessage>,
        ) -> Result<SentBatch, PublisherError> {
            Ok(SentBatch::Ack)
        }

        fn as_any(&self) -> &dyn Any {
            self
        }
    }

    fn hook_route(state: HookState, concurrency: usize) -> Route {
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let factory_name = format!("hooks_{}", unique_id);
        let mut factory = MockEndpointFactory::new();

        let consumer_state = state.clone();
        factory.consumer_behavior = Arc::new(Mutex::new(move || {
            Ok(Box::new(HookConsumer {
                state: consumer_state.clone(),
            }) as Box<dyn MessageConsumer>)
        }));

        let publisher_state = state;
        factory.publisher_behavior = Arc::new(Mutex::new(move || {
            Ok(Box::new(HookPublisher {
                state: publisher_state.clone(),
            }) as Box<dyn MessagePublisher>)
        }));

        register_endpoint_factory(&factory_name, Arc::new(factory));

        let input = Endpoint {
            endpoint_type: EndpointType::Custom {
                name: factory_name.clone(),
                config: serde_json::Value::Null,
            },
            middlewares: vec![],
            handler: None,
        };
        let output = Endpoint {
            endpoint_type: EndpointType::Custom {
                name: factory_name,
                config: serde_json::Value::Null,
            },
            middlewares: vec![],
            handler: None,
        };
        Route::new(input, output).with_concurrency(concurrency)
    }

    #[tokio::test]
    async fn test_lifecycle_hooks_called_once_sequentially() {
        let state = HookState::default();
        let route = hook_route(state.clone(), 1);

        let stopped_by_shutdown = route
            .run_until_err("test_lifecycle_sequential", None, None)
            .await
            .unwrap();

        assert!(!stopped_by_shutdown);
        assert_eq!(state.consumer_connects.load(Ordering::SeqCst), 1);
        assert_eq!(state.consumer_disconnects.load(Ordering::SeqCst), 1);
        assert_eq!(state.publisher_connects.load(Ordering::SeqCst), 1);
        assert_eq!(state.publisher_disconnects.load(Ordering::SeqCst), 1);
        assert_eq!(state.shared_mutations.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn test_lifecycle_hooks_called_once_concurrently() {
        let state = HookState::default();
        let route = hook_route(state.clone(), 4);

        route
            .run_until_err("test_lifecycle_concurrent", None, None)
            .await
            .unwrap();

        assert_eq!(state.consumer_connects.load(Ordering::SeqCst), 1);
        assert_eq!(state.consumer_disconnects.load(Ordering::SeqCst), 1);
        assert_eq!(state.publisher_connects.load(Ordering::SeqCst), 1);
        assert_eq!(state.publisher_disconnects.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_lifecycle_on_connect_failure_stops_route() {
        let state = HookState::default();
        state.fail_consumer_connect.store(true, Ordering::SeqCst);
        let route = hook_route(state.clone(), 1);

        let err = route
            .run_until_err("test_lifecycle_connect_failure", None, None)
            .await
            .unwrap_err();

        assert!(err.to_string().contains("on_connect hook failed"));
        assert_eq!(state.publisher_connects.load(Ordering::SeqCst), 1);
        assert_eq!(state.consumer_connects.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_lifecycle_on_disconnect_failure_does_not_stop_route() {
        let state = HookState::default();
        state.fail_consumer_disconnect.store(true, Ordering::SeqCst);
        state
            .fail_publisher_disconnect
            .store(true, Ordering::SeqCst);
        let route = hook_route(state.clone(), 1);

        let stopped_by_shutdown = route
            .run_until_err("test_lifecycle_disconnect_failure", None, None)
            .await
            .unwrap();

        assert!(!stopped_by_shutdown);
        assert_eq!(state.consumer_disconnects.load(Ordering::SeqCst), 1);
        assert_eq!(state.publisher_disconnects.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_start_fails_on_unavailable_endpoint() {
        // tokio::time::pause();
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let factory_name = format!("unavailable_{}", unique_id);

        let factory = Arc::new(MockEndpointFactory {
            create_consumer_fail: true,
            ..MockEndpointFactory::new()
        });
        register_endpoint_factory(&factory_name, factory);

        let input = Endpoint {
            endpoint_type: EndpointType::Custom {
                name: factory_name,
                config: serde_json::Value::Null,
            },
            middlewares: vec![],
            handler: None,
        };
        let output = Endpoint::new_memory("out", 10);
        let route = Route::new(input, output);

        // The route should fail to start because the input endpoint fails to create.
        // The run() method waits for a ready signal which never comes.
        let result = route.run("test_start_fail").await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("failed to start"));
    }

    #[tokio::test]
    async fn test_reconnect_on_consumer_error() {
        // tokio::time::pause();
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let factory_name = format!("reconnect_{}", unique_id);

        // Shared state to track connection attempts
        let connection_attempts = Arc::new(AtomicUsize::new(0));
        let attempts_clone = connection_attempts.clone();

        let consumer_logic = move || -> Result<Box<dyn MessageConsumer>, anyhow::Error> {
            let attempt = attempts_clone.fetch_add(1, Ordering::SeqCst);

            struct FlakyConsumer {
                attempt: usize,
            }
            #[async_trait::async_trait]
            impl MessageConsumer for FlakyConsumer {
                async fn receive_batch(
                    &mut self,
                    _max: usize,
                ) -> Result<ReceivedBatch, ConsumerError> {
                    if self.attempt == 0 {
                        // First connection works for one batch, then fails
                        self.attempt = 999; // prevent infinite loop in this instance
                        Ok(ReceivedBatch {
                            messages: vec![crate::CanonicalMessage::from("msg1")],
                            commit: Box::new(|_| Box::pin(async { Ok(()) })),
                        })
                    } else if self.attempt == 999 {
                        // Simulate connection drop
                        Err(ConsumerError::Connection(anyhow::anyhow!(
                            "Connection dropped"
                        )))
                    } else {
                        // Subsequent connections work
                        // Sleep a bit to prevent busy loop in test
                        tokio::time::sleep(Duration::from_millis(100)).await;
                        Ok(ReceivedBatch {
                            messages: vec![crate::CanonicalMessage::from("msg2")],
                            commit: Box::new(|_| Box::pin(async { Ok(()) })),
                        })
                    }
                }
                fn as_any(&self) -> &dyn Any {
                    self
                }
            }
            Ok(Box::new(FlakyConsumer { attempt }))
        };

        let mut factory = MockEndpointFactory::new();
        factory.consumer_behavior = Arc::new(Mutex::new(consumer_logic));
        register_endpoint_factory(&factory_name, Arc::new(factory));

        let input = Endpoint {
            endpoint_type: EndpointType::Custom {
                name: factory_name,
                config: serde_json::Value::Null,
            },
            middlewares: vec![],
            handler: None,
        };
        let output = Endpoint::new_memory(&format!("out_{}", unique_id), 10);
        let route = Route::new(input, output.clone());

        route.deploy("test_reconnect").await.unwrap();

        // Wait for reconnection and messages
        let mut verifier = create_consumer_from_route("verifier", &output)
            .await
            .unwrap();

        // Should receive msg1
        let msg1 = tokio::time::timeout(std::time::Duration::from_secs(10), verifier.receive())
            .await
            .expect("Timed out waiting for msg1")
            .unwrap();
        assert_eq!(msg1.message.get_payload_str(), "msg1");

        // Route encounters error, sleeps 5s (skipped by pause), reconnects.
        // Should receive msg2
        let msg2 = tokio::time::timeout(std::time::Duration::from_secs(10), verifier.receive())
            .await
            .expect("Timed out waiting for msg2")
            .unwrap();
        assert_eq!(msg2.message.get_payload_str(), "msg2");

        assert!(connection_attempts.load(Ordering::SeqCst) >= 2);
        Route::stop("test_reconnect").await;
    }

    #[tokio::test]
    async fn test_non_retryable_handler_error_does_not_crash_route() {
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let in_topic = format!("bad_input_in_{}", unique_id);
        let out_topic = format!("bad_input_out_{}", unique_id); // Not used, but good practice

        let input = Endpoint::new_memory(&in_topic, 10);
        let output = Endpoint::new_memory(&out_topic, 10);

        // A handler that fails on specific input
        let handler = |msg: crate::CanonicalMessage| async move {
            if msg.get_payload_str() == "poison" {
                Err(HandlerError::NonRetryable(anyhow::anyhow!("Invalid input")))
            } else {
                Ok(crate::Handled::Publish(msg))
            }
        };

        let route = Route::new(input.clone(), output).with_handler(handler);
        route.deploy("test_invalid_input").await.unwrap();

        let input_ch = input.channel().unwrap();
        let out_channel = route.output.channel().unwrap();

        // 1. Send poison message
        input_ch.send_message("poison".into()).await.unwrap();

        // 2. Send valid message
        input_ch.send_message("valid".into()).await.unwrap();

        // 3. Verify the valid message was processed and published
        let received = tokio::time::timeout(std::time::Duration::from_secs(5), async {
            loop {
                if let Some(msg) = out_channel.drain_messages().pop() {
                    return msg;
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("Timed out waiting for valid message to be processed");
        assert_eq!(received.get_payload_str(), "valid");
        Route::stop("test_invalid_input").await;
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_dlq_and_retry_batch_integration() {
        use crate::models::{DeadLetterQueueMiddleware, Middleware, RetryMiddleware};
        use crate::traits::{MessagePublisher, PublisherError, SentBatch};
        use std::collections::HashMap;
        use std::sync::Mutex;

        // Mock publisher that fails messages with even-numbered IDs
        #[derive(Clone)]
        struct PartialFailPublisher {
            attempts: Arc<Mutex<HashMap<u128, usize>>>,
        }

        #[async_trait::async_trait]
        impl MessagePublisher for PartialFailPublisher {
            async fn send_batch(
                &self,
                messages: Vec<CanonicalMessage>,
            ) -> Result<SentBatch, PublisherError> {
                let mut failed = Vec::new();
                let mut attempts = self.attempts.lock().unwrap();

                for msg in messages {
                    let msg_num: u32 = serde_json::from_slice::<serde_json::Value>(&msg.payload)
                        .unwrap()["id"]
                        .as_u64()
                        .unwrap() as u32;

                    let attempt_count = attempts.entry(msg.message_id).or_insert(0);
                    *attempt_count += 1;

                    if msg_num % 2 == 0 {
                        // Fail even numbers
                        failed.push((
                            msg,
                            PublisherError::Retryable(anyhow::anyhow!("simulated failure")),
                        ));
                    }
                    // Odd numbers succeed implicitly by not being in `failed`
                }

                if failed.is_empty() {
                    Ok(SentBatch::Ack)
                } else {
                    Ok(SentBatch::Partial {
                        responses: None,
                        failed,
                    })
                }
            }
            async fn send(
                &self,
                _msg: CanonicalMessage,
            ) -> Result<crate::traits::Sent, PublisherError> {
                unimplemented!()
            }
            fn as_any(&self) -> &dyn Any {
                self
            }
        }

        // 1. Setup
        let in_topic = "batch_retry_dlq_in";
        let out_topic = "batch_retry_dlq_out";
        let dlq_topic = "batch_retry_dlq_dlq";

        let input = Endpoint::new_memory(in_topic, 10);
        let dlq_endpoint = Endpoint::new_memory(dlq_topic, 10);

        let mock_publisher = PartialFailPublisher {
            attempts: Arc::new(Mutex::new(HashMap::new())),
        };

        let mut output_with_middlewares = Endpoint::new_memory(out_topic, 10);
        output_with_middlewares.middlewares = vec![
            Middleware::Retry(RetryMiddleware {
                max_attempts: 2,
                initial_interval_ms: 1,
                ..Default::default()
            }),
            Middleware::Dlq(Box::new(DeadLetterQueueMiddleware {
                endpoint: dlq_endpoint.clone(),
            })),
        ];

        let route = Route::new(input.clone(), output_with_middlewares).with_batch_size(4);
        // Inject the mock publisher into the route's output
        let final_publisher = crate::middleware::apply_middlewares_to_publisher(
            Box::new(mock_publisher.clone()),
            &route.output,
            "test_route",
        )
        .await
        .unwrap();

        // We need a way to run the route with our mocked publisher.
        // The simplest way is to manually drive the core logic.
        let (work_tx, work_rx) =
            async_channel::bounded::<(Vec<crate::CanonicalMessage>, BatchCommitFunc)>(1);
        let (seq_tx, _sequencer_handle) = spawn_sequencer(1);

        // Spawn a worker to process one batch
        tokio::spawn(async move {
            if let Ok((messages, commit)) = work_rx.recv().await {
                let batch_len = messages.len();
                match final_publisher.send_batch(messages).await {
                    Ok(SentBatch::Ack) => {
                        let _ = commit(vec![MessageDisposition::Ack; batch_len]).await;
                    }
                    Ok(SentBatch::Partial { failed, .. }) => {
                        // In a real route, we'd map responses, but here we just care about failure.
                        let dispositions = if failed.is_empty() {
                            vec![MessageDisposition::Ack; batch_len]
                        } else {
                            // This is a simplification for the test. A real implementation
                            // would map dispositions based on message IDs.
                            vec![MessageDisposition::Nack; batch_len]
                        };
                        let _ = commit(dispositions).await;
                    }
                    Err(_) => {
                        let _ = commit(vec![MessageDisposition::Nack; batch_len]).await;
                    }
                }
            }
        });

        // 2. Send a batch of messages
        let mut messages = Vec::new();
        for i in 1..=4 {
            // 1 (ok), 2 (fail), 3 (ok), 4 (fail)
            messages.push(CanonicalMessage::from_json(serde_json::json!({"id": i})).unwrap());
        }
        let commit = wrap_commit(Box::new(|_| Box::pin(async { Ok(()) })), 0, seq_tx.clone());
        work_tx.send((messages, commit)).await.unwrap();

        // 3. Verify
        let dlq_channel = dlq_endpoint.channel().unwrap();

        let start = std::time::Instant::now();
        while dlq_channel.len() < 2 {
            if start.elapsed() > std::time::Duration::from_secs(5) {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        }

        let dlq_msgs = dlq_channel.drain_messages();

        assert_eq!(dlq_msgs.len(), 2, "Expected 2 messages to go to DLQ");

        let dlq_ids: std::collections::HashSet<u32> = dlq_msgs
            .iter()
            .map(|m| {
                serde_json::from_slice::<serde_json::Value>(&m.payload).unwrap()["id"]
                    .as_u64()
                    .unwrap() as u32
            })
            .collect();

        assert!(dlq_ids.contains(&2));
        assert!(dlq_ids.contains(&4));

        // Verify retry attempts
        let attempts = mock_publisher.attempts.lock().unwrap();
        // Messages 2 and 4 should be tried `max_attempts` times.
        assert_eq!(attempts.values().filter(|&&c| c == 2).count(), 2);
        // Messages 1 and 3 should be tried once.
        assert_eq!(attempts.values().filter(|&&c| c == 1).count(), 2);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_route_dlq_integration() {
        // Setup: Input -> [Panic(Disconnect) -> Retry -> DLQ] -> Output
        // Panic(Disconnect) simulates transient failure.
        // Retry handles it up to N times.
        // If max attempts reached, DLQ catches it.
        // Note: Middleware application order is [Panic, Retry, DLQ] in list to wrap as DLQ(Retry(Panic(Endpoint))).

        let unique_id = fast_uuid_v7::gen_id().to_string();
        let in_topic = format!("dlq_in_{}", unique_id);
        let out_topic = format!("dlq_out_{}", unique_id);
        let dlq_topic = format!("dlq_target_{}", unique_id);
        let input = Endpoint::new_memory(&in_topic, 10);
        let dlq_endpoint = Endpoint::new_memory(&dlq_topic, 10);

        let mut output = Endpoint::new_memory(&out_topic, 10);
        output.middlewares = vec![
            // Inner-most: Fail always
            Middleware::RandomPanic(RandomPanicMiddleware {
                mode: FaultMode::Timeout, // Returns Retryable error, does NOT cause route restart
                trigger_on_message: None, // Fail always
                enabled: true,
                ..Default::default()
            }),
            // Middle: Retry
            Middleware::Retry(crate::models::RetryMiddleware {
                max_attempts: 2,
                initial_interval_ms: 10,
                max_interval_ms: 100,
                multiplier: 1.0,
            }),
            // Outer-most: DLQ
            Middleware::Dlq(Box::new(crate::models::DeadLetterQueueMiddleware {
                endpoint: dlq_endpoint.clone(),
            })),
        ];

        let route = Route::new(input.clone(), output);
        route.deploy("test_dlq_integration").await.unwrap();

        // Send message
        let input_ch = input.channel().unwrap();
        input_ch.send_message("fail_msg".into()).await.unwrap();

        // Verify:
        // 1. Output channel is empty (msg failed to go there)
        // 2. DLQ channel has message

        let dlq_ch = dlq_endpoint.channel().unwrap();

        // Wait for DLQ
        let received = tokio::time::timeout(std::time::Duration::from_secs(5), async {
            loop {
                let batch = dlq_ch.drain_messages();
                if !batch.is_empty() {
                    return batch[0].clone();
                }
                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            }
        })
        .await
        .expect("Timed out waiting for DLQ");

        assert_eq!(received.get_payload_str(), "fail_msg");

        let out_ch_target = mq_bridge::endpoints::memory::get_or_create_channel(
            &mq_bridge::models::MemoryConfig::new(&out_topic, None),
        );
        assert!(out_ch_target.is_empty(), "Message should not reach target");

        Route::stop("test_dlq_integration").await;
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_large_message_handling() {
        let unique_id = fast_uuid_v7::gen_id().to_string();
        let in_topic = format!("large_in_{}", unique_id);
        let out_topic = format!("large_out_{}", unique_id);

        let input = Endpoint::new_memory(&in_topic, 5); // Small capacity
        let output = Endpoint::new_memory(&out_topic, 5);

        let route = Route::new(input.clone(), output.clone());
        route.deploy("test_large_msg").await.unwrap();

        let large_payload = vec![b'x'; 5 * 1024 * 1024]; // 5MB
        let input_ch = input.channel().unwrap();

        input_ch
            .send_message(large_payload.clone().into())
            .await
            .unwrap();

        let mut verifier = route.connect_to_output("verifier").await.unwrap();
        let received = tokio::time::timeout(std::time::Duration::from_secs(10), verifier.receive())
            .await
            .expect("Timed out receiving large message")
            .unwrap();

        assert_eq!(received.message.payload.len(), large_payload.len());
        assert_eq!(received.message.payload, large_payload.as_slice());

        Route::stop("test_large_msg").await;
    }

    #[test]
    fn test_map_responses_to_dispositions_unit() {
        test_map_responses_to_dispositions_logic();
    }
}