pipeflow 0.0.4

//! Engine for DAG construction and execution

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use futures::stream::{self, StreamExt};
use serde::Serialize;
use tokio::sync::{broadcast, mpsc};
use tokio::task::JoinHandle;
use tokio_stream::wrappers::BroadcastStream;

use crate::common::message::{Message, SharedMessage};
use crate::common::types::{Audit, Event};
use crate::config::Config;
use crate::error::Result;
use crate::sink::Sink;
use crate::source::system::{AUDIT_SOURCE_ID, DLQ_SOURCE_ID, EVENT_SOURCE_ID};
use crate::source::{AuditHook, MessageSender, SendHook, Source};
use crate::transform::Transform;

mod build;
mod system;

pub use system::SystemChannels;

use self::build::{build_sink, build_source, build_transform};
use self::system::{
    build_system_sink, build_system_source, route_error_to_system, route_overflow_to_system,
};

/// Pipeline metrics for observability
///
/// Thread-safe counters tracking message flow through the pipeline.
/// All counters use relaxed ordering for minimal overhead.
#[derive(Debug, Default)]
pub struct PipelineMetrics {
    /// Total messages sent by sources
    pub messages_sent: AtomicU64,
    /// Total messages received by sinks
    pub messages_received: AtomicU64,
    /// Total messages dropped due to buffer overflow
    pub messages_dropped: AtomicU64,
    /// Total processing errors (transform + sink)
    pub errors: AtomicU64,
    /// Total messages sent to DLQ
    pub dlq_sent: AtomicU64,
    /// Total audit records dropped due to channel full
    pub audit_dropped: AtomicU64,
}

impl PipelineMetrics {
    /// Record a message sent by a source
    pub fn record_sent(&self) {
        self.messages_sent.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a message received by a sink
    pub fn record_received(&self) {
        self.messages_received.fetch_add(1, Ordering::Relaxed);
    }

    /// Record messages dropped due to buffer overflow
    pub fn record_dropped(&self, count: u64) {
        self.messages_dropped.fetch_add(count, Ordering::Relaxed);
    }

    /// Record a processing error
    pub fn record_error(&self) {
        self.errors.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a message sent to DLQ
    pub fn record_dlq(&self) {
        self.dlq_sent.fetch_add(1, Ordering::Relaxed);
    }

    /// Record an audit record dropped due to channel full
    pub fn record_audit_dropped(&self) {
        self.audit_dropped.fetch_add(1, Ordering::Relaxed);
    }

    /// Get a snapshot of current metrics
    pub fn snapshot(&self) -> MetricsSnapshot {
        MetricsSnapshot {
            messages_sent: self.messages_sent.load(Ordering::Relaxed),
            messages_received: self.messages_received.load(Ordering::Relaxed),
            messages_dropped: self.messages_dropped.load(Ordering::Relaxed),
            errors: self.errors.load(Ordering::Relaxed),
            dlq_sent: self.dlq_sent.load(Ordering::Relaxed),
            audit_dropped: self.audit_dropped.load(Ordering::Relaxed),
        }
    }
}

/// Immutable snapshot of pipeline metrics
#[derive(Debug, Clone, Serialize)]
pub struct MetricsSnapshot {
    /// Total messages sent by sources
    pub messages_sent: u64,
    /// Total messages received by sinks
    pub messages_received: u64,
    /// Total messages dropped due to buffer overflow
    pub messages_dropped: u64,
    /// Total processing errors
    pub errors: u64,
    /// Total messages sent to DLQ
    pub dlq_sent: u64,
    /// Total audit records dropped due to channel full
    pub audit_dropped: u64,
}

/// Engine represents an executable data pipeline
pub struct Engine {
    config: Config,
    sources: HashMap<String, Arc<dyn Source>>,
    transforms: HashMap<String, Arc<dyn Transform>>,
    sinks: HashMap<String, Arc<dyn Sink>>,
    shutdown_tx: broadcast::Sender<()>,
    /// System channels for DLQ/Event/Notify sources
    system_channels: Option<SystemChannels>,
    /// Pending system receivers to be consumed during build
    /// Key: source ID (e.g., "source::system::dlq"), Value: boxed receiver
    pending_system_receivers: HashMap<String, Box<dyn std::any::Any + Send>>,
    /// Pipeline metrics for observability
    metrics: Arc<PipelineMetrics>,
}

impl Engine {
    /// Create a new pipeline from configuration
    pub fn from_config(mut config: Config) -> Result<Self> {
        // Ensure config is normalized and validated even when constructed programmatically
        // (e.g. via `Config::from_yaml()`).
        config.normalize()?;
        config.validate()?;

        let (shutdown_tx, _) = broadcast::channel(1);

        Ok(Self {
            config,
            sources: HashMap::new(),
            transforms: HashMap::new(),
            sinks: HashMap::new(),
            shutdown_tx,
            system_channels: None,
            pending_system_receivers: HashMap::new(),
            metrics: Arc::new(PipelineMetrics::default()),
        })
    }

    /// Load pipeline from a YAML file
    pub fn from_file(path: impl AsRef<std::path::Path>) -> Result<Self> {
        let config = Config::from_file(path)?;
        Self::from_config(config)
    }

    /// Build the pipeline nodes from configuration.
    ///
    /// Note: `Engine::from_config()` normalizes and validates the config.
    pub async fn build(&mut self) -> Result<()> {
        // Create system channels for built-in system sources
        self.setup_system_channels();

        // Use index-based iteration to avoid cloning entire config vectors.
        // This works around borrow checker limitations while keeping zero-copy access.

        // Build sources
        // Clone config data upfront to avoid borrow conflicts with &mut self in build_source
        let source_configs: Vec<_> = self
            .config
            .pipeline
            .sources
            .iter()
            .map(|cfg| (cfg.id.clone(), cfg.source_type.clone(), cfg.config.clone()))
            .collect();
        for (id, source_type, config) in source_configs {
            let source = build_source(&id, &source_type, &config)?;
            self.sources.insert(id, source);
        }

        // Build system sources
        for system_id in crate::source::system::VALID_SYSTEM_SOURCE_IDS {
            let source = build_system_source(system_id, &mut self.pending_system_receivers)?;
            self.sources.insert(system_id.to_string(), source);
        }

        // Build transforms
        for i in 0..self.config.pipeline.transforms.len() {
            let transform_config = &self.config.pipeline.transforms[i];
            let transform = build_transform(transform_config)?;
            let id = transform_config.id.clone();
            self.transforms.insert(id, transform);
        }

        // Build sinks
        for i in 0..self.config.pipeline.sinks.len() {
            let sink_config = &self.config.pipeline.sinks[i];
            let sink = build_sink(sink_config, &self.config.system).await?;
            let id = sink_config.id.clone();
            self.sinks.insert(id, sink);
        }

        // Build system sinks
        for system_sink_id in crate::sink::system::VALID_SYSTEM_SINK_IDS {
            let sink = build_system_sink(system_sink_id, &self.config.system).await?;
            self.sinks.insert(system_sink_id.to_string(), sink);
        }

        Ok(())
    }

    /// Setup system channels for built-in sources
    fn setup_system_channels(&mut self) {
        let channel_size = self.config.system.channel_size();
        let (dlq_tx, dlq_rx) = mpsc::channel::<Message>(channel_size);
        let (event_tx, event_rx) = mpsc::channel::<Event>(channel_size);
        let (audit_tx, audit_rx) = mpsc::channel::<Audit>(channel_size);

        self.system_channels = Some(SystemChannels {
            dlq: dlq_tx,
            event: event_tx,
            audit: audit_tx,
        });

        self.pending_system_receivers
            .insert(DLQ_SOURCE_ID.to_string(), Box::new(dlq_rx));
        self.pending_system_receivers
            .insert(EVENT_SOURCE_ID.to_string(), Box::new(event_rx));
        self.pending_system_receivers
            .insert(AUDIT_SOURCE_ID.to_string(), Box::new(audit_rx));
    }

    /// Build a lookup map from source ID to source config for O(1) access.
    fn build_source_config_map(&self) -> HashMap<&str, &crate::config::SourceConfig> {
        self.config
            .pipeline
            .sources
            .iter()
            .map(|cfg| (cfg.id.as_str(), cfg))
            .collect()
    }

    /// Build a lookup map from transform ID to transform config for O(1) access.
    fn build_transform_config_map(&self) -> HashMap<&str, &crate::config::TransformConfig> {
        self.config
            .pipeline
            .transforms
            .iter()
            .map(|cfg| (cfg.id.as_str(), cfg))
            .collect()
    }

    /// Create broadcast channels for all nodes (sources and transforms) that can emit messages.
    /// Uses broadcast channels for fan-out to multiple downstream nodes.
    /// SharedMessage (Arc<Message>) ensures efficient cloning during broadcast.
    ///
    /// Each source/transform can override the system output_buffer_size in its config.
    fn create_node_channels(&self) -> HashMap<String, broadcast::Sender<SharedMessage>> {
        let default_capacity = self.config.system.output_buffer_size();
        let source_config_map = self.build_source_config_map();
        let transform_config_map = self.build_transform_config_map();

        let total_nodes = self.sources.len() + self.transforms.len();
        let mut node_senders = HashMap::with_capacity(total_nodes);

        // Create channels for all sources
        for source_id in self.sources.keys() {
            let capacity = source_config_map
                .get(source_id.as_str())
                .and_then(|cfg| cfg.output_buffer_size)
                .unwrap_or(default_capacity);

            let (broadcast_tx, _) = broadcast::channel(capacity);
            node_senders.insert(source_id.clone(), broadcast_tx);
        }

        // Create channels for all transforms
        for transform_id in self.transforms.keys() {
            let capacity = transform_config_map
                .get(transform_id.as_str())
                .and_then(|cfg| cfg.output_buffer_size)
                .unwrap_or(default_capacity);

            let (broadcast_tx, _) = broadcast::channel(capacity);
            node_senders.insert(transform_id.clone(), broadcast_tx);
        }

        node_senders
    }

    /// Spawn source tasks and return their handles
    fn spawn_source_tasks(
        &self,
        node_channels: &HashMap<String, broadcast::Sender<SharedMessage>>,
        status_tx: mpsc::Sender<(String, std::result::Result<(), String>)>,
        system_channels: Option<SystemChannels>,
    ) -> Vec<JoinHandle<()>> {
        let mut handles = Vec::new();

        for (source_id, source) in &self.sources {
            // INVARIANT: create_node_channels() creates a channel for every source in self.sources
            let sender = node_channels
                .get(source_id)
                .expect("source channel must exist: create_node_channels guarantees this")
                .clone();
            let metrics = Arc::clone(&self.metrics);
            let send_hook: SendHook = Arc::new(move || metrics.record_sent());
            let mut sender = MessageSender::new(sender, Some(send_hook));

            // Add audit hook if system channels are available
            // Skip audit hooks for system sources to avoid feedback loops
            let is_system_source = crate::source::system::is_valid_system_source_id(source_id);
            if !is_system_source && let Some(ref channels) = system_channels {
                let audit_sender = channels.audit.clone();
                let source_id_for_audit = source_id.clone();
                let metrics_for_audit = Arc::clone(&self.metrics);
                let audit_hook: AuditHook = Arc::new(move |message_id: &str| {
                    let audit = Audit::new(&source_id_for_audit, "source", "emit")
                        .with_message_id(message_id)
                        .success(0); // Source emit doesn't have meaningful duration
                    // Use try_send to avoid blocking in sync context
                    if audit_sender.try_send(audit).is_err() {
                        metrics_for_audit.record_audit_dropped();
                    }
                });
                sender = sender.with_audit_hook(audit_hook);
            }

            let shutdown_rx = self.shutdown_tx.subscribe();
            let source = Arc::clone(source);
            let source_id_clone = source_id.clone();
            let status_tx = status_tx.clone();

            let handle = tokio::spawn(async move {
                let result = source.run(sender, shutdown_rx).await;

                // Map Error to String to make it simpler to send over channel
                // We do this to avoid introducing a complex Error type into the system channel
                let status = match result {
                    Ok(_) => Ok(()),
                    Err(e) => Err(e.to_string()),
                };

                // Notify engine of completion (success or failure)
                if let Err(e) = status_tx.send((source_id_clone.clone(), status)).await {
                    tracing::debug!(
                        source_id = %source_id_clone,
                        error = %e,
                        "Failed to report source status (engine likely shutting down)"
                    );
                }
            });
            handles.push(handle);
        }

        handles
    }

    /// Spawn transform tasks that subscribe to input channels and emit to output channels.
    ///
    /// Each transform:
    /// 1. Subscribes to its configured input node(s) via BroadcastStream
    /// 2. Processes each incoming message through its transform logic
    /// 3. Emits output messages to its own broadcast channel for downstream nodes
    fn spawn_transform_tasks(
        &self,
        node_channels: &HashMap<String, broadcast::Sender<SharedMessage>>,
        system_channels: Option<SystemChannels>,
        metrics: Arc<PipelineMetrics>,
    ) -> Vec<JoinHandle<()>> {
        let mut handles = Vec::with_capacity(self.transforms.len());
        let transform_config_map = self.build_transform_config_map();

        for (transform_id, transform) in &self.transforms {
            // INVARIANT: build_transform_config_map() maps all transforms from config,
            // and self.transforms is built from the same config in build()
            let transform_config = transform_config_map
                .get(transform_id.as_str())
                .expect("transform config must exist: built from same config source");

            if transform_config.inputs.is_empty() {
                tracing::warn!(transform_id = %transform_id, "Transform has no inputs configured");
                continue;
            }

            // Collect input streams
            let mut streams = Vec::with_capacity(transform_config.inputs.len());
            for input_id in &transform_config.inputs {
                if let Some(broadcast_tx) = node_channels.get(input_id) {
                    let rx = broadcast_tx.subscribe();
                    streams.push(BroadcastStream::new(rx));
                } else {
                    tracing::warn!(
                        transform_id = %transform_id,
                        input_id = %input_id,
                        "Transform input not found"
                    );
                }
            }

            if streams.is_empty() {
                tracing::warn!(transform_id = %transform_id, "Transform connected to no valid inputs");
                continue;
            }

            // INVARIANT: create_node_channels() creates a channel for every transform in self.transforms
            let output_tx = node_channels
                .get(transform_id)
                .expect("transform output channel must exist: create_node_channels guarantees this")
                .clone();

            let transform = Arc::clone(transform);
            let transform_id_clone = transform_id.clone();
            let channels = system_channels.clone();
            let metrics = Arc::clone(&metrics);

            // Check if this transform consumes system sources to prevent feedback loops
            // Computed outside the async block to avoid lifetime issues
            let consumes_system = transform_config.inputs.iter().any(|id| {
                matches!(
                    id.as_str(),
                    crate::source::system::DLQ_SOURCE_ID
                        | crate::source::system::EVENT_SOURCE_ID
                        | crate::source::system::AUDIT_SOURCE_ID
                )
            });

            let handle = tokio::spawn(async move {
                tracing::info!(transform_id = %transform_id_clone, "Transform task started");

                let mut merged = stream::select_all(streams);

                while let Some(result) = merged.next().await {
                    #[allow(unreachable_patterns)]
                    match result {
                        Ok(msg) => {
                            // Clone from Arc for transform processing
                            let msg_owned = (*msg).clone();
                            let message_id = msg.meta.id.to_string();
                            let start = Instant::now();

                            match transform.process(msg_owned).await {
                                Ok(outputs) => {
                                    let duration_ms = start.elapsed().as_millis() as u64;

                                    // Send audit record for successful processing (only if not a system transform)
                                    if !consumes_system && let Some(ref ch) = channels {
                                        let audit =
                                            Audit::new(&transform_id_clone, "transform", "process")
                                                .with_message_id(&message_id)
                                                .success(duration_ms);
                                        let _ = ch.audit.send(audit).await;
                                    }

                                    for output_msg in outputs {
                                        let shared = Arc::new(output_msg);
                                        if output_tx.send(shared).is_err() {
                                            // No receivers, channel closed
                                            tracing::debug!(
                                                transform_id = %transform_id_clone,
                                                "Transform output channel closed"
                                            );
                                        }
                                    }
                                }
                                Err(e) => {
                                    let duration_ms = start.elapsed().as_millis() as u64;
                                    let error_msg = e.to_string();

                                    tracing::error!(
                                        transform_id = %transform_id_clone,
                                        error = %e,
                                        "Transform failed to process message"
                                    );
                                    metrics.record_error();

                                    if let Some(ref ch) = channels {
                                        // Send audit record for failed processing
                                        let audit =
                                            Audit::new(&transform_id_clone, "transform", "process")
                                                .with_message_id(&message_id)
                                                .failure(duration_ms, &error_msg);
                                        let _ = ch.audit.send(audit).await;

                                        route_error_to_system(
                                            ch,
                                            &metrics,
                                            &transform_id_clone,
                                            "transform",
                                            &error_msg,
                                            &msg,
                                        )
                                        .await;
                                    }
                                }
                            }
                        }
                        Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(
                            n,
                        )) => {
                            tracing::warn!(
                                transform_id = %transform_id_clone,
                                lagged = n,
                                "Transform input lagged, {} messages skipped",
                                n
                            );
                            metrics.record_dropped(n);

                            if let Some(ref ch) = channels {
                                route_overflow_to_system(ch, &transform_id_clone, "transform", n)
                                    .await;
                            }
                        }
                        Err(e) => {
                            tracing::debug!(
                                transform_id = %transform_id_clone,
                                error = %e,
                                "Transform input stream error"
                            );
                        }
                    }
                }

                tracing::info!(transform_id = %transform_id_clone, "Transform task completed");
            });
            handles.push(handle);
        }

        handles
    }

    /// Build a lookup map from sink ID to transform IDs that output into it.
    fn build_sink_input_map(&self) -> HashMap<&str, Vec<&str>> {
        let mut inputs: HashMap<&str, Vec<&str>> = HashMap::new();

        for transform in &self.config.pipeline.transforms {
            for output in &transform.outputs {
                inputs
                    .entry(output.as_str())
                    .or_default()
                    .push(transform.id.as_str());
            }
        }

        // Default system wiring: system sources -> system sinks
        inputs
            .entry(crate::sink::system::DLQ_SINK_ID)
            .or_default()
            .push(crate::source::system::DLQ_SOURCE_ID);
        inputs
            .entry(crate::sink::system::EVENT_SINK_ID)
            .or_default()
            .push(crate::source::system::EVENT_SOURCE_ID);
        inputs
            .entry(crate::sink::system::AUDIT_SINK_ID)
            .or_default()
            .push(crate::source::system::AUDIT_SOURCE_ID);

        for values in inputs.values_mut() {
            values.sort_unstable();
            values.dedup();
        }

        inputs
    }

    /// Spawn sink tasks using stream merging (no forwarder tasks).
    ///
    /// Each sink subscribes to multiple broadcast channels and merges them using
    /// `futures::stream::select_all`. This eliminates intermediate forwarder tasks
    /// and reduces message passing overhead.
    fn spawn_sink_tasks(
        &self,
        broadcast_senders: &HashMap<String, broadcast::Sender<SharedMessage>>,
        system_channels: Option<SystemChannels>,
        metrics: Arc<PipelineMetrics>,
    ) -> Vec<JoinHandle<()>> {
        let mut sink_handles = Vec::with_capacity(self.sinks.len());

        let sink_input_map = self.build_sink_input_map();
        let transform_config_map = self.build_transform_config_map();

        // Helper to check upstream dependencies
        let check_system_dependencies = |sink_id: &str| -> (bool, bool, bool) {
            use std::collections::HashSet;
            let mut visited = HashSet::new();
            let mut stack = Vec::new();

            // Initial inputs
            if let Some(inputs) = sink_input_map.get(sink_id) {
                stack.extend(inputs.iter().cloned());
            }

            let mut consumes_dlq = false;
            let mut consumes_event = false;
            let mut consumes_audit = false;

            while let Some(node_id) = stack.pop() {
                if !visited.insert(node_id) {
                    continue;
                }

                match node_id {
                    crate::source::system::DLQ_SOURCE_ID => consumes_dlq = true,
                    crate::source::system::EVENT_SOURCE_ID => consumes_event = true,
                    crate::source::system::AUDIT_SOURCE_ID => consumes_audit = true,
                    _ => {
                        // If it's a transform, push its inputs
                        if let Some(cfg) = transform_config_map.get(node_id) {
                            for input in &cfg.inputs {
                                stack.push(input.as_str());
                            }
                        }
                    }
                }
            }
            (consumes_dlq, consumes_event, consumes_audit)
        };

        for (sink_id, sink) in &self.sinks {
            let Some(input_ids) = sink_input_map.get(sink_id.as_str()) else {
                tracing::warn!(sink_id = %sink_id, "Sink has no upstream transforms configured");
                continue;
            };

            // Detect system dependencies for this sink
            let (consumes_dlq, consumes_event, consumes_audit) =
                check_system_dependencies(sink_id.as_str());

            if consumes_dlq || consumes_event || consumes_audit {
                tracing::info!(
                    sink_id = %sink_id,
                    dlq = consumes_dlq,
                    event = consumes_event,
                    audit = consumes_audit,
                    "Sink consumes system messages; suppressing specific error loops"
                );
            }

            let mut streams = Vec::with_capacity(input_ids.len());
            let mut connected_inputs = 0;

            for input_id in input_ids {
                let input_id = *input_id;
                if let Some(broadcast_tx) = broadcast_senders.get(input_id) {
                    let rx = broadcast_tx.subscribe();
                    // Wrap broadcast::Receiver as an async Stream
                    let stream = BroadcastStream::new(rx);
                    streams.push(stream);
                    connected_inputs += 1;
                } else {
                    tracing::warn!(
                        sink_id = %sink_id,
                        input_id = %input_id,
                        "Sink input source not found"
                    );
                }
            }

            if connected_inputs == 0 {
                tracing::warn!(sink_id = %sink_id, "Sink connected to no valid sources");
                continue;
            }

            let sink = Arc::clone(sink);
            let sink_id_clone = sink_id.clone();
            let dlq_sender = system_channels.clone();
            let metrics = Arc::clone(&metrics);

            // Spawn sink task with merged streams - no forwarder overhead
            let handle = tokio::spawn(async move {
                tracing::info!(sink_id = %sink_id_clone, inputs = connected_inputs, "Sink task started");

                // Merge all input streams into one using select_all
                let mut merged = stream::select_all(streams);

                while let Some(result) = merged.next().await {
                    // NOTE: At the time of writing, `BroadcastStreamRecvError` only exposes
                    // `Lagged` in our dependency set, so additional error branches can be
                    // flagged as unreachable by clippy. We still allow unreachable patterns
                    // here to keep the loop resilient across `tokio-stream` changes.
                    #[allow(unreachable_patterns)]
                    match result {
                        Ok(msg) => {
                            let message_id = msg.meta.id.to_string();
                            let start = Instant::now();

                            if let Err(e) = sink.process(msg.clone()).await {
                                let duration_ms = start.elapsed().as_millis() as u64;
                                let error_msg = e.to_string();

                                tracing::error!(
                                    sink_id = %sink_id_clone,
                                    error = %e,
                                    "Sink failed to process message"
                                );
                                metrics.record_error();

                                if let Some(ref channels) = dlq_sender {
                                    // Send audit record for failed write (only if not an audit sink)
                                    if !consumes_audit {
                                        let audit = Audit::new(&sink_id_clone, "sink", "write")
                                            .with_message_id(&message_id)
                                            .failure(duration_ms, &error_msg);
                                        let _ = channels.audit.send(audit).await;
                                    }

                                    // Manually inline route_error_to_system logic to support granular suppression
                                    // If consumes_event, suppress Event generation
                                    if !consumes_event {
                                        let event = Event::new(
                                            "sink_error",
                                            serde_json::json!({
                                                "node_id": sink_id_clone,
                                                "node_type": "sink",
                                                "error": error_msg,
                                                "message_id": msg.meta.id.to_string(),
                                            }),
                                        );
                                        let _ = channels.event.send(event).await;
                                    }

                                    // If consumes_dlq, suppress DLQ routing
                                    if !consumes_dlq {
                                        // Use helper from system module but we might need to make it public or duplicate logic
                                        // Since route_error_to_system is async and uses internal helpers, duplication is safer here
                                        // or we just call the helper but we can't easily suppress partial parts.
                                        //
                                        // Looking at route_error_to_system:
                                        // It sends Event AND DLQ.
                                        // Here we split them.

                                        // DLQ Logic:
                                        // We need to import MAX_CHAIN_DEPTH and bump_chain_depth or replicate logic
                                        // Replicating simple logic here to avoid exposing internal helpers if they are private
                                        const MAX_DEPTH: u8 = 8;
                                        if msg.meta.source_node
                                            != crate::source::system::DLQ_SOURCE_ID
                                            && msg.meta.chain_depth < MAX_DEPTH
                                        {
                                            // msg is Arc<Message>, we need to clone the inner Message to modify it
                                            let mut updated = (*msg).clone();
                                            updated.meta.chain_depth =
                                                updated.meta.chain_depth.saturating_add(1);
                                            let _ = channels.dlq.send(updated).await;
                                            metrics.record_dlq();
                                        }
                                    }
                                }
                            } else {
                                let duration_ms = start.elapsed().as_millis() as u64;
                                metrics.record_received();

                                // Send audit record for successful write (only if not an audit sink)
                                if !consumes_audit
                                    && !consumes_event
                                    && !consumes_dlq
                                    && let Some(ref channels) = dlq_sender
                                {
                                    let audit = Audit::new(&sink_id_clone, "sink", "write")
                                        .with_message_id(&message_id)
                                        .success(duration_ms);
                                    let _ = channels.audit.send(audit).await;
                                }
                            }
                        }
                        Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(
                            n,
                        )) => {
                            tracing::warn!(
                                sink_id = %sink_id_clone,
                                lagged = n,
                                "Sink input lagged, {} messages skipped",
                                n
                            );
                            metrics.record_dropped(n);

                            if let Some(ref channels) = dlq_sender
                                && !consumes_event
                            {
                                route_overflow_to_system(channels, &sink_id_clone, "sink", n).await;
                            }
                        }
                        Err(e) => {
                            // Future-proof: handle additional broadcast stream error variants.
                            tracing::debug!(
                                sink_id = %sink_id_clone,
                                error = %e,
                                "Sink input stream returned error"
                            );
                        }
                    }
                }

                tracing::info!(sink_id = %sink_id_clone, "Sink task completed");
            });
            sink_handles.push(handle);
        }

        sink_handles
    }

    /// Gracefully shutdown all tasks with timeout.
    ///
    /// Shutdown order: sources first (stop producing), then transforms, then sinks.
    /// This ensures data flows through the pipeline before shutdown completes.
    async fn graceful_shutdown_all(
        source_handles: Vec<JoinHandle<()>>,
        transform_handles: Vec<JoinHandle<()>>,
        sink_handles: Vec<JoinHandle<()>>,
        shutdown_timeout: tokio::time::Duration,
    ) {
        // Helper to shutdown a batch of handles
        async fn shutdown_batch(
            handles: Vec<JoinHandle<()>>,
            timeout: tokio::time::Duration,
            node_type: &str,
        ) {
            for handle in handles {
                let mut handle = handle;
                match tokio::time::timeout(timeout, &mut handle).await {
                    Ok(join_result) => {
                        if let Err(e) = join_result {
                            tracing::warn!(error = %e, "{} task join failed", node_type);
                        }
                    }
                    Err(_) => {
                        tracing::warn!("{} task did not stop before timeout; aborting", node_type);
                        handle.abort();
                    }
                }
            }
        }

        // 1. Wait for sources to stop producing
        shutdown_batch(source_handles, shutdown_timeout, "Source").await;

        // 2. Wait for transforms to finish processing
        shutdown_batch(transform_handles, shutdown_timeout, "Transform").await;

        // 3. Wait for sinks to drain remaining messages
        shutdown_batch(sink_handles, shutdown_timeout, "Sink").await;
    }

    /// Run the pipeline with a custom shutdown signal
    ///
    /// This method is primarily used for testing or when programmatically controlling
    /// the pipeline lifecycle. For standard execution, use `run()`.
    ///
    /// # Arguments
    ///
    /// * `shutdown_signal` - A future that completes when the pipeline should shut down
    pub async fn run_with_signal<F>(self, shutdown_signal: F) -> Result<()>
    where
        F: std::future::Future<Output = ()> + Send + 'static,
    {
        tracing::info!("Starting pipeline execution");

        let node_channels = self.create_node_channels();

        // Create channel for source status updates
        let (status_tx, mut status_rx) = mpsc::channel(100);

        // Subscribe downstream nodes before starting sources to avoid a startup race.
        // Order: sinks first, then transforms, then sources
        let sink_handles = self.spawn_sink_tasks(
            &node_channels,
            self.system_channels.clone(),
            Arc::clone(&self.metrics),
        );
        let transform_handles = self.spawn_transform_tasks(
            &node_channels,
            self.system_channels.clone(),
            Arc::clone(&self.metrics),
        );
        let source_handles =
            self.spawn_source_tasks(&node_channels, status_tx, self.system_channels.clone());

        // Spawn metrics logger task
        let metrics_logger = {
            let metrics = Arc::clone(&self.metrics);
            let interval = self.config.system.metrics_log_interval();

            if interval.is_zero() {
                tokio::spawn(std::future::pending::<()>())
            } else {
                tokio::spawn(async move {
                    let mut interval = tokio::time::interval(interval);
                    // tick immediately returns, so skip the first one to avoid immediate log on startup
                    interval.tick().await;

                    loop {
                        interval.tick().await;
                        let snapshot = metrics.snapshot();
                        tracing::info!(
                            sent = snapshot.messages_sent,
                            received = snapshot.messages_received,
                            dropped = snapshot.messages_dropped,
                            errors = snapshot.errors,
                            dlq = snapshot.dlq_sent,
                            "Pipeline metrics"
                        );
                    }
                })
            }
        };

        tracing::info!(
            sources = self.sources.len(),
            transforms = self.transforms.len(),
            sinks = self.sinks.len(),
            "Pipeline running"
        );

        let is_system_source = |id: &str| crate::source::system::is_valid_system_source_id(id);
        let mut active_sources = self
            .sources
            .keys()
            .filter(|id| !is_system_source(id))
            .count();
        let system_channels = self.system_channels.clone();
        let metrics = Arc::clone(&self.metrics);

        // Main event loop
        tokio::pin!(shutdown_signal);
        loop {
            tokio::select! {
                _ = &mut shutdown_signal => {
                    tracing::info!("Shutdown signal received, stopping pipeline...");
                    break;
                }
                Some((source_id, result)) = status_rx.recv() => {
                    match result {
                        Ok(_) => {
                            tracing::info!(source_id = %source_id, "Source finished successfully");
                            if !is_system_source(&source_id) {
                                active_sources = active_sources.saturating_sub(1);
                            }

                            if active_sources == 0 {
                                tracing::info!("All sources finished, shutting down pipeline...");
                                break;
                            }
                        }
                        Err(e) => {
                            tracing::error!(
                                source_id = %source_id,
                                error = %e,
                                "Source failed, triggering fail-fast shutdown"
                            );
                            metrics.record_error();
                            if let Some(ref channels) = system_channels {
                                let event = Event::new(
                                    "source_error",
                                    serde_json::json!({
                                        "node_id": source_id,
                                        "node_type": "source",
                                        "error": e,
                                    }),
                                );
                                let _ = channels.event.send(event).await;
                            }
                            break;
                        }
                    }
                }
                else => {
                     // status_rx closed (should not happen while sources are running)
                     break;
                }
            }
        }

        // Send shutdown signal to all sources
        let _ = self.shutdown_tx.send(());

        // Stop metrics logging
        metrics_logger.abort();

        // Drop all senders to close channels and unblock receivers
        drop(node_channels);

        // Graceful shutdown: sources -> transforms -> sinks
        let shutdown_timeout = self.config.system.shutdown_timeout();
        Self::graceful_shutdown_all(
            source_handles,
            transform_handles,
            sink_handles,
            shutdown_timeout,
        )
        .await;

        tracing::info!("Pipeline shutdown complete");
        Ok(())
    }

    /// Run the pipeline
    ///
    /// Executes the pipeline by:
    /// 1. Building channel connections between sources and sinks
    /// 2. Spawning async tasks for each source and sink
    /// 3. Waiting for shutdown signal (Ctrl+C)
    /// 4. Gracefully shutting down all tasks
    pub async fn run(self) -> Result<()> {
        let signal = async {
            let _ = tokio::signal::ctrl_c().await;
        };
        self.run_with_signal(signal).await
    }

    /// Get system channels for sending messages to system sources.
    pub fn system_channels(&self) -> Option<&SystemChannels> {
        self.system_channels.as_ref()
    }

    /// Get pipeline metrics for observability.
    pub fn metrics(&self) -> &PipelineMetrics {
        &self.metrics
    }

    /// Shutdown the pipeline gracefully
    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(());
    }
}

#[cfg(test)]
mod tests;