weavegraph 0.3.0

use std::io;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use tokio::sync::{broadcast, oneshot};
use tokio::task;

use super::diagnostics::{DiagnosticsStream, HealthState, SinkDiagnostic, SinkHealth};
use super::emitter::EventEmitter;
use super::hub::{EventHub, EventHubMetrics, EventStream};
use super::sink::{EventSink, StdOutSink};
use chrono::Utc;

/// Central event broadcasting system for workflow execution events.
///
/// `EventBus` receives events from workflow nodes and broadcasts them to multiple
/// sinks (stdout, channels, files, monitoring systems, etc.). It's the backbone
/// of Weavegraph's observability and streaming capabilities.
///
/// # Architecture
///
/// The EventBus is owned by [`AppRunner`](crate::runtimes::runner::AppRunner), not
/// [`App`](crate::app::App). This design allows:
/// - Multiple runners to share the same graph with different event configurations
/// - Per-request event isolation in web servers
/// - Flexible sink composition
///
/// ```text
/// Workflow Nodes
///     │ ctx.emit()
///     ▼
/// EventBus
///     │ broadcast
///     ├─────┬─────┬─────┐
///     ▼     ▼     ▼     ▼
/// StdOut Channel File Custom
///  Sink   Sink   Sink  Sink
/// ```
///
/// # Usage Patterns
///
/// ## Default EventBus (Stdout Only)
///
/// When using [`App::invoke()`](crate::app::App::invoke), a default EventBus
/// with stdout sink is created automatically:
///
/// ```rust,no_run
/// # use weavegraph::app::App;
/// # use weavegraph::state::VersionedState;
/// # async fn example(app: App) -> Result<(), Box<dyn std::error::Error>> {
/// // Events automatically go to stdout
/// let result = app.invoke(VersionedState::new_with_user_message("Hello")).await?;
/// # Ok(())
/// # }
/// ```
///
/// ## Custom EventBus (Streaming to Web Clients)
///
/// For streaming events to web clients, create a custom EventBus and pass it to
/// [`AppRunner`](crate::runtimes::runner::AppRunner):
///
/// ```rust,no_run
/// use weavegraph::event_bus::{EventBus, ChannelSink, StdOutSink};
/// use weavegraph::runtimes::{AppRunner, CheckpointerType};
/// use weavegraph::state::VersionedState;
/// # use weavegraph::app::App;
/// # async fn example(app: App) -> Result<(), Box<dyn std::error::Error>> {
///
/// // Create channel for streaming
/// let (tx, rx) = flume::unbounded();
///
/// // Create EventBus with multiple sinks
/// let bus = EventBus::with_sinks(vec![
///     Box::new(StdOutSink::default()),  // Server logs
///     Box::new(ChannelSink::new(tx)),   // Client streaming
/// ]);
///
/// // Pass EventBus to AppRunner
/// let mut runner = AppRunner::with_options_and_bus(
///     app,
///     CheckpointerType::InMemory,
///     false,
///     bus,  // Custom EventBus
///     true,
/// ).await;
///
/// let session_id = "client-123".to_string();
/// runner.create_session(
///     session_id.clone(),
///     VersionedState::new_with_user_message("Process this")
/// ).await?;
///
/// // Consume events from channel
/// tokio::spawn(async move {
///     while let Ok(event) = rx.recv_async().await {
///         // Send to web client via SSE, WebSocket, etc.
///         println!("Event: {:?}", event);
///     }
/// });
///
/// runner.run_until_complete(&session_id).await?;
/// # Ok(())
/// # }
/// ```
///
/// ## Per-Request Isolation (Web Server Pattern)
///
/// Create a new EventBus for each HTTP request to isolate events:
///
/// ```rust,no_run
/// use std::sync::Arc;
/// use weavegraph::event_bus::{EventBus, ChannelSink};
/// use weavegraph::runtimes::{AppRunner, CheckpointerType};
/// use weavegraph::state::VersionedState;
/// # use weavegraph::app::App;
/// # async fn handle_request(app: Arc<App>) -> Result<(), Box<dyn std::error::Error>> {
///
/// // Each request gets its own EventBus and channel
/// let (tx, rx) = flume::unbounded();
/// let bus = EventBus::with_sinks(vec![Box::new(ChannelSink::new(tx))]);
///
/// // Reuse the App, create new runner with isolated EventBus
/// let mut runner = AppRunner::with_options_and_bus(
///     Arc::try_unwrap(app).unwrap_or_else(|arc| (*arc).clone()),
///     CheckpointerType::InMemory,
///     false,
///     bus,  // Isolated EventBus for this request
///     true,
/// ).await;
///
/// // Run workflow - events are isolated to this request
/// let session_id = uuid::Uuid::new_v4().to_string();
/// runner.create_session(
///     session_id.clone(),
///     VersionedState::new_with_user_message("User query")
/// ).await?;
/// runner.run_until_complete(&session_id).await?;
/// # Ok(())
/// # }
/// ```
///
/// # Available Sinks
///
/// - [`StdOutSink`](crate::event_bus::StdOutSink) - Write to stdout (default)
/// - [`ChannelSink`](crate::event_bus::ChannelSink) - Stream to async channels
/// - [`MemorySink`](crate::event_bus::MemorySink) - Capture for testing
/// - Custom sinks implementing [`EventSink`](crate::event_bus::EventSink)
///
/// # See Also
///
/// - [`AppRunner::with_options_and_bus()`](crate::runtimes::runner::AppRunner::with_options_and_bus) - How to use custom EventBus
/// - [`ChannelSink`](crate::event_bus::ChannelSink) - For streaming events
/// - Example: `examples/streaming_events.rs` - Complete streaming demonstration
const DEFAULT_BUFFER_CAPACITY: usize = 1024;

pub struct EventBus {
    sinks: Arc<Mutex<Vec<SinkEntry>>>,
    hub: Arc<EventHub>,
    started: AtomicBool,
    /// Generation counter tracking listener restarts so stale workers exit.
    generation: Arc<AtomicU64>,
    /// Diagnostics broadcast channel for sink errors.
    diagnostics_tx: broadcast::Sender<SinkDiagnostic>,
    /// In-memory health tracking per sink name.
    health: Arc<Mutex<std::collections::HashMap<String, HealthState>>>,
    diagnostics_enabled: bool,
    diagnostics_emit_to_events: bool,
}

impl Default for EventBus {
    fn default() -> Self {
        Self::with_sink(StdOutSink::default())
    }
}

impl EventBus {
    pub fn with_sink<T>(sink: T) -> Self
    where
        T: EventSink + 'static,
    {
        Self::with_sinks(vec![Box::new(sink)])
    }

    pub fn with_sinks(sinks: Vec<Box<dyn EventSink>>) -> Self {
        Self::with_capacity(sinks, DEFAULT_BUFFER_CAPACITY)
    }

    pub(crate) fn with_capacity(sinks: Vec<Box<dyn EventSink>>, buffer_capacity: usize) -> Self {
        Self::with_capacity_and_diag(sinks, buffer_capacity, buffer_capacity, true, false)
    }

    pub(crate) fn with_capacity_and_diag(
        sinks: Vec<Box<dyn EventSink>>,
        buffer_capacity: usize,
        diagnostics_capacity: usize,
        diagnostics_enabled: bool,
        diagnostics_emit_to_events: bool,
    ) -> Self {
        let hub = EventHub::new(buffer_capacity);
        let entries = sinks.into_iter().map(SinkEntry::new).collect();
        let (diagnostics_tx, _) = if diagnostics_enabled {
            broadcast::channel(diagnostics_capacity.max(1))
        } else {
            // Create a tiny channel and immediately drop the sender at drop time when EventBus drops.
            // Publishing will be skipped when disabled, so the capacity is largely irrelevant.
            broadcast::channel(1)
        };
        Self {
            sinks: Arc::new(Mutex::new(entries)),
            hub,
            started: AtomicBool::new(false),
            generation: Arc::new(AtomicU64::new(0)),
            diagnostics_tx,
            health: Arc::new(Mutex::new(std::collections::HashMap::new())),
            diagnostics_enabled,
            diagnostics_emit_to_events,
        }
    }

    pub fn add_sink<T: EventSink + 'static>(&self, sink: T) {
        self.add_boxed_sink(Box::new(sink));
    }

    /// Attach a new sink to the hub, starting a worker immediately if the bus is live.
    pub fn add_boxed_sink(&self, sink: Box<dyn EventSink>) {
        let mut sinks_guard = self.sinks.lock().expect("EventBus sinks mutex poisoned");
        let mut entry = SinkEntry::new(sink);
        if self.started.load(Ordering::SeqCst) {
            let generation = self.generation.load(Ordering::SeqCst);
            entry.spawn_worker(
                self.hub.clone(),
                Arc::clone(&self.generation),
                generation,
                self.diagnostics_tx.clone(),
                Arc::clone(&self.health),
                self.diagnostics_enabled,
                self.diagnostics_emit_to_events,
            );
        }
        sinks_guard.push(entry);
    }

    pub fn get_emitter(&self) -> Arc<dyn EventEmitter> {
        Arc::new(self.hub.emitter())
    }

    /// Return current hub metrics (buffer capacity and cumulative drop count).
    pub fn metrics(&self) -> EventHubMetrics {
        self.hub.metrics()
    }

    pub fn subscribe(&self) -> EventStream {
        self.listen_for_events();
        self.hub.subscribe()
    }

    /// Access a broadcast stream of sink diagnostics.
    ///
    /// The returned stream mirrors the `EventStream` consumption model but carries
    /// `SinkDiagnostic` entries emitted when sinks report errors. This stream is
    /// isolated from the main event flow to avoid feedback loops.
    pub fn diagnostics(&self) -> DiagnosticsStream {
        DiagnosticsStream::new(self.diagnostics_tx.subscribe())
    }

    /// Return a snapshot of per-sink health counters and last error details.
    pub fn sink_health(&self) -> Vec<SinkHealth> {
        let health = self.health.lock().expect("EventBus health mutex poisoned");
        health
            .iter()
            .map(|(sink, state)| SinkHealth {
                sink: sink.clone(),
                error_count: state.error_count,
                last_error: state.last_error.clone(),
                last_error_at: state.last_error_at,
            })
            .collect()
    }

    /// Spawn workers for every registered sink. Safe to call multiple times.
    pub fn listen_for_events(&self) {
        if self.started.swap(true, Ordering::SeqCst) {
            return;
        }
        let mut sinks = self.sinks.lock().expect("EventBus sinks mutex poisoned");
        let generation = self.generation.load(Ordering::SeqCst);
        for entry in sinks.iter_mut() {
            entry.spawn_worker(
                self.hub.clone(),
                Arc::clone(&self.generation),
                generation,
                self.diagnostics_tx.clone(),
                Arc::clone(&self.health),
                self.diagnostics_enabled,
                self.diagnostics_emit_to_events,
            );
        }
    }

    /// Signal all sink workers to stop pulling from the hub.
    pub async fn stop_listener(&self) {
        if !self.started.swap(false, Ordering::SeqCst) {
            return;
        }
        self.generation.fetch_add(1, Ordering::SeqCst);
        let workers = {
            let mut sinks = self.sinks.lock().expect("EventBus sinks mutex poisoned");
            let mut collected = Vec::with_capacity(sinks.len());
            for entry in sinks.iter_mut() {
                if let Some(worker) = entry.worker.take() {
                    collected.push(worker);
                }
            }
            collected
        };
        for worker in workers {
            let SinkWorker { shutdown, handle } = worker;
            let _ = shutdown.send(());
            let _ = handle.await;
        }
    }

    pub fn close_channel(&self) {
        self.hub.close();
    }
}

impl Drop for EventBus {
    fn drop(&mut self) {
        self.hub.close();
        if self.started.load(Ordering::SeqCst) {
            let mut sinks = self.sinks.lock().expect("EventBus sinks mutex poisoned");
            for entry in sinks.iter_mut() {
                entry.abort_worker();
            }
        }
    }
}

struct SinkEntry {
    sink: Arc<Mutex<Box<dyn EventSink>>>,
    /// Resolved once at registration to avoid recomputing on error paths.
    name: String,
    worker: Option<SinkWorker>,
}

impl SinkEntry {
    fn new(sink: Box<dyn EventSink>) -> Self {
        let candidate = sink.name();
        let default_marker: &str = std::any::type_name::<dyn EventSink>();
        // Prefer implementor override; otherwise fall back to the dynamic concrete type name.
        let name = if candidate == default_marker {
            std::any::type_name_of_val(&*sink).to_string()
        } else {
            candidate
        };
        Self {
            sink: Arc::new(Mutex::new(sink)),
            name,
            worker: None,
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn spawn_worker(
        &mut self,
        hub: Arc<EventHub>,
        generation_state: Arc<AtomicU64>,
        active_generation: u64,
        diagnostics_tx: broadcast::Sender<SinkDiagnostic>,
        health: Arc<Mutex<std::collections::HashMap<String, HealthState>>>,
        diagnostics_enabled: bool,
        diagnostics_emit_to_events: bool,
    ) {
        if self.worker.is_some() {
            return;
        }
        // Each worker holds an `Arc` to the sink so consumers can add/remove sinks without
        // racing the async tasks we spawn here.
        let sink = Arc::clone(&self.sink);
        let sink_name = self.name.clone();
        let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
        let mut stream = hub.subscribe();
        let de_enabled = diagnostics_enabled;
        let de_emit = diagnostics_emit_to_events;
        let hub_clone = Arc::clone(&hub);
        let handle = task::spawn(async move {
            fn record_sink_error(
                health: &Arc<Mutex<std::collections::HashMap<String, HealthState>>>,
                diagnostics_tx: &broadcast::Sender<SinkDiagnostic>,
                sink_name: &str,
                err_msg: &str,
                diagnostics_enabled: bool,
            ) {
                if diagnostics_enabled {
                    let mut map = health.lock().expect("health mutex poisoned");
                    let entry = map.entry(sink_name.to_string()).or_default();
                    entry.error_count = entry.error_count.saturating_add(1);
                    entry.last_error = Some(err_msg.to_string());
                    entry.last_error_at = Some(Utc::now());
                    let occurrence = entry.error_count;
                    drop(map);
                    let _ = diagnostics_tx.send(SinkDiagnostic {
                        sink: sink_name.to_string(),
                        error: err_msg.to_string(),
                        when: Utc::now(),
                        occurrence,
                    });
                }
            }
            loop {
                // Bail out early if the bus has been stopped/restarted since this worker spawned.
                if generation_state.load(Ordering::SeqCst) != active_generation {
                    break;
                }
                tokio::select! {
                    _ = &mut shutdown_rx => break,
                    event = stream.recv() => match event {
                        Ok(event) => {
                            let sink = Arc::clone(&sink);
                            let sink_name = sink_name.clone();
                            let diagnostics_tx = diagnostics_tx.clone();
                            let health = Arc::clone(&health);
                            let diagnostics_enabled = de_enabled;
                            let hub_for_emit = Arc::clone(&hub_clone);
                            let diagnostics_emit_to_events = de_emit;
                            // Dispatch potentially blocking sink logic onto the dedicated
                            // blocking pool so we never park the async runtime thread.
                            let dispatch = task::spawn_blocking(move || -> io::Result<()> {
                                let mut guard = sink.lock().expect("sink mutex poisoned");
                                guard.handle(&event)
                            });
                            match dispatch.await {
                                Ok(Ok(())) => {}
                                Ok(Err(err)) => {
                                    let err_msg = err.to_string();
                                    tracing::error!(
                                        target: "weavegraph::event_bus",
                                        error = %err_msg,
                                        sink = %sink_name,
                                        "event sink reported an error while handling event"
                                    );
                                    record_sink_error(&health, &diagnostics_tx, &sink_name, &err_msg, diagnostics_enabled);
                                    if diagnostics_emit_to_events {
                                        let _ = hub_for_emit.publish(super::event::Event::diagnostic(
                                            "event_bus.sink_error",
                                            format!("{sink}: {err}", sink=sink_name, err=err_msg),
                                        ));
                                    }
                                }
                                Err(err) => {
                                    let err_msg = err.to_string();
                                    tracing::error!(
                                        target: "weavegraph::event_bus",
                                        error = %err_msg,
                                        sink = %sink_name,
                                        "event sink worker task failed to join"
                                    );
                                    // Treat join failures as sink errors for health/diagnostics
                                    record_sink_error(&health, &diagnostics_tx, &sink_name, &err_msg, diagnostics_enabled);
                                    if diagnostics_emit_to_events {
                                        let _ = hub_for_emit.publish(super::event::Event::diagnostic(
                                            "event_bus.sink_join_error",
                                            format!("{sink}: {err}", sink=sink_name, err=err_msg),
                                        ));
                                    }
                                }
                            }
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                    }
                }
            }
        });
        self.worker = Some(SinkWorker {
            shutdown: shutdown_tx,
            handle,
        });
    }

    fn abort_worker(&mut self) {
        if let Some(worker) = self.worker.take() {
            let _ = worker.shutdown.send(());
            worker.handle.abort();
        }
    }
}

struct SinkWorker {
    shutdown: oneshot::Sender<()>,
    handle: task::JoinHandle<()>,
}