use super::{BroadcastMsg, GlobalRegistry, PubSubMessage, ThreadLocalPubSub};
use crate::network::{Connection, PubSubOp};
use crate::protocol::resp::{write_resp_value, RespValue};
use std::sync::Arc;
pub fn handle_pubsub_operation(
pubsub_manager: &mut ThreadLocalPubSub,
pubsub_registry: &Arc<GlobalRegistry>,
conn_id: usize,
op: PubSubOp,
connection: &mut Connection,
thread_id: usize,
) -> Vec<(usize, PubSubMessage)> {
let mut local_deliveries = Vec::new();
match op {
PubSubOp::Subscribe(channels) => {
let messages = pubsub_manager.subscribe(conn_id, channels);
for message in messages {
if let PubSubMessage::Subscribe { count, .. } = &message {
connection.set_subscription_count(*count);
}
connection.queue_pubsub_message(message);
}
}
PubSubOp::Unsubscribe(channels) => {
let messages = pubsub_manager.unsubscribe(conn_id, channels);
for message in messages {
if let PubSubMessage::Unsubscribe { count, .. } = &message {
connection.set_subscription_count(*count);
}
connection.queue_pubsub_message(message);
}
}
PubSubOp::PSubscribe(patterns) => {
let messages = pubsub_manager.psubscribe(conn_id, patterns);
for message in messages {
if let PubSubMessage::PSubscribe { count, .. } = &message {
connection.set_subscription_count(*count);
}
connection.queue_pubsub_message(message);
}
}
PubSubOp::PUnsubscribe(patterns) => {
let messages = pubsub_manager.punsubscribe(conn_id, patterns);
for message in messages {
if let PubSubMessage::PUnsubscribe { count, .. } = &message {
connection.set_subscription_count(*count);
}
connection.queue_pubsub_message(message);
}
}
PubSubOp::Publish { channel, message } => {
local_deliveries =
pubsub_manager.publish_local(&channel, &bytes::Bytes::from(message.clone()));
let channel_threads = pubsub_registry.get_channel_threads(&channel);
let pattern_threads = pubsub_registry.get_all_pattern_threads();
let msg = BroadcastMsg::Publish {
channel: channel.clone(),
message: bytes::Bytes::from(message.clone()),
exclude_thread: Some(thread_id),
};
pubsub_registry.broadcast_to_threads(msg, &channel_threads);
let pattern_msg = BroadcastMsg::PatternPublish {
channel: channel.clone(),
message: bytes::Bytes::from(message),
exclude_thread: Some(thread_id),
};
let pattern_thread_vec: Vec<_> = pattern_threads.into_iter().collect();
pubsub_registry.broadcast_to_threads(pattern_msg, &pattern_thread_vec);
let total_count = pubsub_registry.get_channel_subscriber_count(&channel)
+ pubsub_registry.get_total_pattern_matches(&channel);
let resp = RespValue::Integer(total_count as i64);
write_resp_value(&mut connection.write_buffer, &resp);
}
PubSubOp::PubSubChannels { pattern } => {
let all_channels = pubsub_registry.get_all_channels();
let filtered = if let Some(pat) = pattern {
all_channels
.into_iter()
.filter(|ch| GlobalRegistry::glob_match(&pat, ch))
.map(|ch| RespValue::BulkString(Some(ch.into())))
.collect()
} else {
all_channels
.into_iter()
.map(|ch| RespValue::BulkString(Some(ch.into())))
.collect()
};
let resp = RespValue::Array(Some(filtered));
write_resp_value(&mut connection.write_buffer, &resp);
}
PubSubOp::PubSubNumSub { channels } => {
let mut results = Vec::new();
for channel in channels {
results.push(RespValue::BulkString(Some(channel.clone().into())));
let count = pubsub_registry.get_channel_subscriber_count(&channel);
results.push(RespValue::Integer(count as i64));
}
let resp = RespValue::Array(Some(results));
write_resp_value(&mut connection.write_buffer, &resp);
}
PubSubOp::PubSubNumPat => {
let count = pubsub_registry.get_pattern_count();
let resp = RespValue::Integer(count as i64);
write_resp_value(&mut connection.write_buffer, &resp);
}
}
local_deliveries
}