use crate::common::error::Result;
use crate::common::protocol::Frame;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, warn};
#[async_trait]
pub trait MessageHandler: Send + Sync {
async fn handle(&self, frame: &Frame) -> Result<Option<Frame>>;
}
pub struct MessageRouter {
handlers: HashMap<String, Vec<Arc<dyn MessageHandler>>>,
default_handler: Option<Arc<dyn MessageHandler>>,
}
impl MessageRouter {
pub fn new() -> Self {
Self {
handlers: HashMap::new(),
default_handler: None,
}
}
pub fn register(&mut self, route: impl Into<String>, handler: Arc<dyn MessageHandler>) {
let route = route.into();
self.handlers.entry(route).or_default().push(handler);
}
pub fn set_default_handler(&mut self, handler: Arc<dyn MessageHandler>) {
self.default_handler = Some(handler);
}
pub async fn route(&self, frame: &Frame) -> Result<Vec<Frame>> {
let route_key = Self::extract_route_key(frame);
debug!(
"Routing message: route={}, frame_id={}",
route_key, frame.message_id
);
let mut replies = Vec::new();
if let Some(handlers) = self.handlers.get(&route_key) {
for handler in handlers {
match handler.handle(frame).await {
Ok(Some(reply)) => {
replies.push(reply);
}
Ok(None) => {
}
Err(e) => {
warn!("Handler error for route {}: {}", route_key, e);
}
}
}
} else if let Some(ref default_handler) = self.default_handler {
match default_handler.handle(frame).await {
Ok(Some(reply)) => {
replies.push(reply);
}
Ok(None) => {
}
Err(e) => {
warn!("Default handler error: {}", e);
}
}
} else {
debug!("No handler found for route: {}", route_key);
}
Ok(replies)
}
fn extract_route_key(frame: &Frame) -> String {
if let Some(ref command) = frame.command {
match &command.r#type {
Some(crate::common::protocol::command::Type::System(sys_cmd)) => {
use crate::common::protocol::system_command::Type as SysType;
use std::convert::TryFrom;
match SysType::try_from(sys_cmd.r#type) {
Ok(SysType::Connect) => "system.connect".to_string(),
Ok(SysType::ConnectAck) => "system.connect_ack".to_string(),
Ok(SysType::Close) => "system.close".to_string(),
Ok(SysType::Ping) => "system.ping".to_string(),
Ok(SysType::Pong) => "system.pong".to_string(),
Ok(SysType::Error) => "system.error".to_string(),
Ok(SysType::Event) => "system.event".to_string(),
Ok(SysType::Auth) => "system.auth".to_string(),
Ok(SysType::AuthAck) => "system.auth_ack".to_string(),
_ => "system.unknown".to_string(),
}
}
Some(crate::common::protocol::command::Type::Payload(msg_cmd)) => {
use crate::common::protocol::payload_command::Type as MsgType;
use std::convert::TryFrom;
match MsgType::try_from(msg_cmd.r#type) {
Ok(MsgType::Message) => "payload.message".to_string(),
Ok(MsgType::Event) => "payload.event".to_string(),
Ok(MsgType::Ack) => "payload.ack".to_string(),
Ok(MsgType::Data) => "payload.data".to_string(),
_ => format!("payload.{}", msg_cmd.r#type),
}
}
Some(crate::common::protocol::command::Type::Notification(notif_cmd)) => {
use crate::common::protocol::notification_command::Type as NotifType;
use std::convert::TryFrom;
match NotifType::try_from(notif_cmd.r#type) {
Ok(NotifType::System) => "notification.system".to_string(),
Ok(NotifType::Broadcast) => "notification.broadcast".to_string(),
Ok(NotifType::Alert) => "notification.alert".to_string(),
Ok(NotifType::User) => "notification.user".to_string(),
Ok(NotifType::Connection) => "notification.connection".to_string(),
_ => format!("notification.{}", notif_cmd.r#type),
}
}
Some(crate::common::protocol::command::Type::Custom(custom_cmd)) => {
format!("custom.{}", custom_cmd.name)
}
None => "unknown".to_string(),
}
} else {
"unknown".to_string()
}
}
pub fn remove_route(&mut self, route: &str) {
self.handlers.remove(route);
}
pub fn clear(&mut self) {
self.handlers.clear();
self.default_handler = None;
}
}
impl Default for MessageRouter {
fn default() -> Self {
Self::new()
}
}
pub struct SimpleHandler {
#[allow(clippy::type_complexity)]
handler: Box<dyn Fn(&Frame) -> Result<Option<Frame>> + Send + Sync>,
}
impl SimpleHandler {
pub fn new<F>(handler: F) -> Self
where
F: Fn(&Frame) -> Result<Option<Frame>> + Send + Sync + 'static,
{
Self {
handler: Box::new(handler),
}
}
}
#[async_trait]
impl MessageHandler for SimpleHandler {
async fn handle(&self, frame: &Frame) -> Result<Option<Frame>> {
(self.handler)(frame)
}
}
pub struct AsyncHandler {
#[allow(clippy::type_complexity)]
handler: Arc<
dyn Fn(
&Frame,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Option<Frame>>> + Send + '_>,
> + Send
+ Sync,
>,
}
impl AsyncHandler {
pub fn new<F, Fut>(handler: F) -> Self
where
F: Fn(&Frame) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = Result<Option<Frame>>> + Send + 'static,
{
Self {
handler: Arc::new(move |frame| Box::pin(handler(frame))),
}
}
}
#[async_trait]
impl MessageHandler for AsyncHandler {
async fn handle(&self, frame: &Frame) -> Result<Option<Frame>> {
(self.handler)(frame).await
}
}