use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MailboxMessage {
pub id: String,
pub from: String,
pub to: MessageTarget,
pub subject: String,
pub body: String,
pub timestamp: DateTime<Utc>,
pub read: bool,
pub priority: MailboxPriority,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MessageTarget {
Direct(String),
Broadcast,
Team(String),
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
pub enum MailboxPriority {
Low,
Normal,
High,
Urgent,
}
impl MailboxMessage {
pub fn direct(
from: impl Into<String>,
to: impl Into<String>,
subject: impl Into<String>,
body: impl Into<String>,
) -> Self {
Self {
id: Uuid::new_v4().to_string(),
from: from.into(),
to: MessageTarget::Direct(to.into()),
subject: subject.into(),
body: body.into(),
timestamp: Utc::now(),
read: false,
priority: MailboxPriority::Normal,
}
}
pub fn broadcast(
from: impl Into<String>,
subject: impl Into<String>,
body: impl Into<String>,
) -> Self {
Self {
id: Uuid::new_v4().to_string(),
from: from.into(),
to: MessageTarget::Broadcast,
subject: subject.into(),
body: body.into(),
timestamp: Utc::now(),
read: false,
priority: MailboxPriority::Normal,
}
}
pub fn with_priority(mut self, priority: MailboxPriority) -> Self {
self.priority = priority;
self
}
}
#[derive(Debug, Default)]
pub struct AgentMailbox {
inbox: VecDeque<MailboxMessage>,
max_size: usize,
}
impl AgentMailbox {
pub fn new(max_size: usize) -> Self {
Self {
inbox: VecDeque::new(),
max_size,
}
}
pub fn deliver(&mut self, message: MailboxMessage) {
if self.inbox.len() >= self.max_size {
self.inbox.pop_front(); }
self.inbox.push_back(message);
}
pub fn read_unread(&mut self) -> Vec<&MailboxMessage> {
self.inbox.iter().filter(|m| !m.read).collect()
}
pub fn mark_read(&mut self, message_id: &str) {
if let Some(msg) = self.inbox.iter_mut().find(|m| m.id == message_id) {
msg.read = true;
}
}
pub fn unread_count(&self) -> usize {
self.inbox.iter().filter(|m| !m.read).count()
}
pub fn all_messages(&self) -> Vec<&MailboxMessage> {
self.inbox.iter().collect()
}
}
pub struct MailboxSystem {
mailboxes: HashMap<String, AgentMailbox>,
max_mailbox_size: usize,
}
impl MailboxSystem {
pub fn new() -> Self {
Self {
mailboxes: HashMap::new(),
max_mailbox_size: 1000,
}
}
pub fn register_agent(&mut self, agent_id: impl Into<String>) {
let id = agent_id.into();
self.mailboxes
.entry(id)
.or_insert_with(|| AgentMailbox::new(self.max_mailbox_size));
}
pub fn send(&mut self, message: MailboxMessage) {
match &message.to {
MessageTarget::Direct(to) => {
if let Some(mailbox) = self.mailboxes.get_mut(to) {
mailbox.deliver(message);
}
}
MessageTarget::Broadcast => {
let sender = message.from.clone();
for (agent_id, mailbox) in &mut self.mailboxes {
if *agent_id != sender {
mailbox.deliver(message.clone());
}
}
}
MessageTarget::Team(team_id) => {
let sender = message.from.clone();
for (agent_id, mailbox) in &mut self.mailboxes {
if *agent_id != sender {
mailbox.deliver(message.clone());
}
}
let _ = team_id; }
}
}
pub fn get_mailbox(&self, agent_id: &str) -> Option<&AgentMailbox> {
self.mailboxes.get(agent_id)
}
pub fn get_mailbox_mut(&mut self, agent_id: &str) -> Option<&mut AgentMailbox> {
self.mailboxes.get_mut(agent_id)
}
pub fn idle_agents(&self) -> Vec<String> {
self.mailboxes
.iter()
.filter(|(_, mb)| mb.unread_count() == 0)
.map(|(id, _)| id.clone())
.collect()
}
}
impl Default for MailboxSystem {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_direct_message() {
let mut system = MailboxSystem::new();
system.register_agent("agent-1");
system.register_agent("agent-2");
let msg = MailboxMessage::direct("agent-1", "agent-2", "Hello", "Test message");
system.send(msg);
let mailbox = system.get_mailbox("agent-2").unwrap();
assert_eq!(mailbox.unread_count(), 1);
}
#[test]
fn test_broadcast() {
let mut system = MailboxSystem::new();
system.register_agent("lead");
system.register_agent("agent-1");
system.register_agent("agent-2");
let msg = MailboxMessage::broadcast("lead", "Update", "Sprint complete");
system.send(msg);
assert_eq!(system.get_mailbox("lead").unwrap().unread_count(), 0);
assert_eq!(system.get_mailbox("agent-1").unwrap().unread_count(), 1);
assert_eq!(system.get_mailbox("agent-2").unwrap().unread_count(), 1);
}
}