velo 0.4.0

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Message dispatcher for active message routing.

use crate::observability::DispatchFailure;
use crate::transports::VeloBackend;
use bytes::Bytes;
use dashmap::DashMap;
use std::sync::{Arc, OnceLock};
use tokio::sync::Semaphore;
use tokio_util::task::TaskTracker;
use tracing::{error, trace, warn};
use velo_ext::WorkerId;

use crate::Messenger;
use crate::messenger::common::events::{EventType, Outcome, encode_event_header};
use crate::messenger::common::messages::ResponseType;
use crate::messenger::common::responses::ResponseId;

/// Context passed to handlers during dispatch.
#[derive(Clone)]
pub(crate) struct HandlerContext {
    /// The response ID for correlation
    pub message_id: ResponseId,

    /// Message payload
    pub payload: Bytes,

    /// Response type (FireAndForget, AckNack, Unary)
    pub response_type: ResponseType,

    /// Optional user headers (for tracing, metadata, etc.)
    pub headers: Option<std::collections::HashMap<String, String>>,

    /// The messenger system for handler access
    pub system: Arc<Messenger>,
}

/// Base trait for active message handlers.
pub(crate) trait ActiveMessageHandler: Send + Sync {
    /// Handle a message asynchronously
    fn handle(
        &self,
        ctx: HandlerContext,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'static>>;

    /// Get the handler name
    fn name(&self) -> &str;
}

/// Trait for dispatching messages to handlers.
pub(crate) trait ActiveMessageDispatcher: Send + Sync {
    /// Get the handler name
    fn name(&self) -> &str;

    /// Dispatch a message to the handler (non-async, kicks off handler execution)
    fn dispatch(&self, ctx: HandlerContext);
}

/// Dispatcher implementation that spawns handlers on a task tracker.
pub(crate) struct SpawnedDispatcher<H: ActiveMessageHandler> {
    handler: Arc<H>,
    task_tracker: TaskTracker,
    semaphore: Option<Arc<Semaphore>>,
}

impl<H: ActiveMessageHandler> SpawnedDispatcher<H> {
    pub fn new(handler: H, task_tracker: TaskTracker) -> Self {
        Self {
            handler: Arc::new(handler),
            task_tracker,
            semaphore: None,
        }
    }

    #[expect(dead_code)]
    #[doc(hidden)]
    pub fn with_concurrency_limit(handler: H, task_tracker: TaskTracker, limit: usize) -> Self {
        Self {
            handler: Arc::new(handler),
            task_tracker,
            semaphore: Some(Arc::new(Semaphore::new(limit))),
        }
    }
}

impl<H: ActiveMessageHandler + 'static> ActiveMessageDispatcher for SpawnedDispatcher<H> {
    fn name(&self) -> &str {
        self.handler.name()
    }

    fn dispatch(&self, ctx: HandlerContext) {
        let handler = self.handler.clone();
        let semaphore = self.semaphore.clone();
        let handler_name = handler.name().to_string();

        self.task_tracker.spawn(async move {
            let _permit = if let Some(sem) = &semaphore {
                Some(sem.acquire().await.expect("semaphore closed"))
            } else {
                None
            };

            trace!(target: "crate::messenger::dispatcher", handler = %handler_name, "Handler task started");
            handler.handle(ctx).await;
            trace!(target: "crate::messenger::dispatcher", handler = %handler_name, "Handler task completed");
        });
    }
}

/// Dispatcher implementation that executes handlers inline on the dispatcher task.
pub(crate) struct InlineDispatcher<H: ActiveMessageHandler> {
    handler: Arc<H>,
    semaphore: Option<Arc<Semaphore>>,
}

impl<H: ActiveMessageHandler> InlineDispatcher<H> {
    pub fn new(handler: H) -> Self {
        Self {
            handler: Arc::new(handler),
            semaphore: None,
        }
    }

    #[expect(dead_code)]
    #[doc(hidden)]
    pub fn with_concurrency_limit(handler: H, limit: usize) -> Self {
        Self {
            handler: Arc::new(handler),
            semaphore: Some(Arc::new(Semaphore::new(limit))),
        }
    }
}

impl<H: ActiveMessageHandler + 'static> ActiveMessageDispatcher for InlineDispatcher<H> {
    fn name(&self) -> &str {
        self.handler.name()
    }

    fn dispatch(&self, ctx: HandlerContext) {
        let handler = self.handler.clone();
        let semaphore = self.semaphore.clone();

        tokio::spawn(async move {
            let _permit = if let Some(sem) = &semaphore {
                Some(sem.acquire().await.expect("semaphore closed"))
            } else {
                None
            };

            handler.handle(ctx).await;
        });
    }
}

/// Main message dispatcher hub that routes messages to handlers.
pub(crate) struct DispatcherHub {
    /// Handler registry (lock-free for fast dispatch).
    /// Shared with HandlerManager so registration is immediately visible.
    handlers: Arc<DashMap<String, Arc<dyn ActiveMessageDispatcher>>>,

    /// Backend for sending messages
    backend: Arc<VeloBackend>,

    /// Messenger system reference (late-bound via OnceLock)
    system: OnceLock<Arc<Messenger>>,

    /// Notifies waiters when `system` has been set
    system_ready: tokio::sync::Notify,
}

impl DispatcherHub {
    /// Create a new dispatcher hub
    pub fn new(backend: Arc<VeloBackend>) -> Self {
        Self {
            handlers: Arc::new(DashMap::new()),
            backend,
            system: OnceLock::new(),
            system_ready: tokio::sync::Notify::new(),
        }
    }

    /// Initialize the system reference (must be called exactly once before dispatching)
    pub fn set_system(&self, system: Arc<Messenger>) -> anyhow::Result<()> {
        self.system
            .set(system)
            .map_err(|_| anyhow::anyhow!("System already initialized"))?;
        self.system_ready.notify_waiters();
        Ok(())
    }

    /// Get the system reference (panics if not initialized)
    pub(crate) fn system(&self) -> &Arc<Messenger> {
        self.system
            .get()
            .expect("System must be initialized before dispatching messages")
    }

    /// Wait until the system reference is available, then return it.
    pub(crate) async fn wait_for_system(&self) -> &Arc<Messenger> {
        // Fast path: already initialized
        if let Some(system) = self.system.get() {
            return system;
        }
        // Register interest before re-checking to avoid missed notification
        let notified = self.system_ready.notified();
        if let Some(system) = self.system.get() {
            return system;
        }
        notified.await;
        self.system
            .get()
            .expect("system must be set after notification")
    }

    /// Get a clone of the handlers Arc for sharing with HandlerManager.
    pub(crate) fn handlers_arc(&self) -> Arc<DashMap<String, Arc<dyn ActiveMessageDispatcher>>> {
        self.handlers.clone()
    }

    /// Get a list of all registered handler names
    pub(crate) fn list_handlers(&self) -> Vec<String> {
        self.handlers
            .iter()
            .map(|entry| entry.key().clone())
            .collect()
    }

    /// Dispatch a message to the appropriate handler
    pub fn dispatch_message(&self, handler_name: &str, ctx: HandlerContext) {
        match self.handlers.get(handler_name) {
            Some(dispatcher) => {
                dispatcher.dispatch(ctx);
            }
            None => {
                self.handle_unknown_handler(handler_name, ctx);
            }
        }
    }

    /// Handle messages for unknown handlers
    fn handle_unknown_handler(&self, handler_name: &str, ctx: HandlerContext) {
        if let Some(metrics) = ctx.system.observability().as_ref() {
            metrics.record_dispatch_failure(DispatchFailure::DispatchUnknownHandler);
        }
        error!(
            target: "crate::messenger::dispatcher",
            handler = %handler_name,
            message_id = %ctx.message_id,
            "No handler registered for message"
        );

        let backend = self.backend.clone();
        let message_id = ctx.message_id;
        let handler_name = handler_name.to_string();

        match ctx.response_type {
            ResponseType::AckNack | ResponseType::Unary => {
                let error_message = format!("Handler '{}' not found", handler_name);
                tokio::spawn(async move {
                    if let Err(e) =
                        Self::send_error_response_static(&backend, message_id, error_message).await
                    {
                        error!(
                            target: "crate::messenger::dispatcher",
                            "Failed to send error response for unknown handler: {}", e
                        );
                    }
                });
            }
            ResponseType::FireAndForget => {
                warn!(
                    target: "crate::messenger::dispatcher",
                    handler = %handler_name,
                    "Fire-and-forget message to unknown handler, no response sent"
                );
            }
        }
    }

    /// Send an error response back to the sender.
    pub(crate) async fn send_error_response(
        &self,
        response_id: ResponseId,
        error_message: String,
    ) -> anyhow::Result<()> {
        Self::send_error_response_static(&self.backend, response_id, error_message).await
    }

    /// Send an error response back to the sender (static method)
    async fn send_error_response_static(
        backend: &VeloBackend,
        response_id: ResponseId,
        error_message: String,
    ) -> anyhow::Result<()> {
        use crate::transports::MessageType;

        let header = encode_event_header(EventType::Ack(response_id, Outcome::Error));
        let payload = Bytes::from(error_message.into_bytes());

        struct DispatcherErrorHandler;
        impl crate::transports::TransportErrorHandler for DispatcherErrorHandler {
            fn on_error(&self, _header: Bytes, _payload: Bytes, error: String) {
                error!(target: "crate::messenger::dispatcher", "Failed to send error response: {}", error);
            }
        }

        static ERROR_HANDLER: std::sync::OnceLock<
            Arc<dyn crate::transports::TransportErrorHandler>,
        > = std::sync::OnceLock::new();
        let error_handler = ERROR_HANDLER
            .get_or_init(|| Arc::new(DispatcherErrorHandler))
            .clone();

        let outcome = backend.send_message_to_worker(
            WorkerId::from_u64(response_id.worker_id()),
            header,
            payload,
            MessageType::Ack,
            error_handler,
        )?;
        if let crate::transports::SendOutcome::Backpressured(bp) = outcome {
            bp.await;
        }

        Ok(())
    }
}