use peat_schema::command::v1::{command_target::Scope, HierarchicalCommand};
use std::collections::HashSet;
pub struct CommandRouter {
node_id: String,
squad_id: Option<String>,
squad_members: Vec<String>,
platoon_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum TargetResolution {
Self_,
Subordinates(Vec<String>),
AllSquadMembers(Vec<String>),
NotApplicable,
}
impl CommandRouter {
pub fn new(
node_id: String,
squad_id: Option<String>,
squad_members: Vec<String>,
platoon_id: Option<String>,
) -> Self {
Self {
node_id,
squad_id,
squad_members,
platoon_id,
}
}
pub fn resolve_target(&self, command: &HierarchicalCommand) -> TargetResolution {
let target = match &command.target {
Some(t) => t,
None => return TargetResolution::NotApplicable,
};
let scope = Scope::try_from(target.scope).unwrap_or(Scope::Unspecified);
match scope {
Scope::Individual => {
let target_ids: HashSet<String> = target.target_ids.iter().cloned().collect();
if target_ids.contains(&self.node_id) {
TargetResolution::Self_
} else {
let subordinate_targets: Vec<String> = self
.squad_members
.iter()
.filter(|m| target_ids.contains(*m))
.cloned()
.collect();
if !subordinate_targets.is_empty() {
TargetResolution::Subordinates(subordinate_targets)
} else {
TargetResolution::NotApplicable
}
}
}
Scope::Squad => {
if let Some(ref my_squad) = self.squad_id {
if target.target_ids.contains(my_squad) {
if !self.squad_members.is_empty() {
TargetResolution::AllSquadMembers(self.squad_members.clone())
} else {
TargetResolution::Self_
}
} else {
TargetResolution::NotApplicable
}
} else {
TargetResolution::NotApplicable
}
}
Scope::Platoon => {
if let Some(ref my_platoon) = self.platoon_id {
if target.target_ids.contains(my_platoon) {
if !self.squad_members.is_empty() {
TargetResolution::AllSquadMembers(self.squad_members.clone())
} else {
TargetResolution::Self_
}
} else {
TargetResolution::NotApplicable
}
} else {
TargetResolution::NotApplicable
}
}
Scope::Broadcast => {
if !self.squad_members.is_empty() {
TargetResolution::AllSquadMembers(self.squad_members.clone())
} else {
TargetResolution::Self_
}
}
Scope::Unspecified => TargetResolution::NotApplicable,
}
}
pub fn should_route(&self, resolution: &TargetResolution) -> bool {
matches!(
resolution,
TargetResolution::Subordinates(_) | TargetResolution::AllSquadMembers(_)
)
}
pub fn get_routing_targets(&self, resolution: &TargetResolution) -> Vec<String> {
match resolution {
TargetResolution::Subordinates(nodes) => nodes.clone(),
TargetResolution::AllSquadMembers(nodes) => nodes.clone(),
_ => Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use peat_schema::command::v1::CommandTarget;
#[test]
fn test_resolve_individual_self() {
let router = CommandRouter::new(
"node-1".to_string(),
Some("squad-alpha".to_string()),
vec!["node-1".to_string(), "node-2".to_string()],
None,
);
let command = HierarchicalCommand {
command_id: "cmd-1".to_string(),
target: Some(CommandTarget {
scope: Scope::Individual as i32,
target_ids: vec!["node-1".to_string()],
}),
..Default::default()
};
let resolution = router.resolve_target(&command);
assert_eq!(resolution, TargetResolution::Self_);
}
#[test]
fn test_resolve_individual_subordinate() {
let router = CommandRouter::new(
"node-1".to_string(),
Some("squad-alpha".to_string()),
vec!["node-1".to_string(), "node-2".to_string()],
None,
);
let command = HierarchicalCommand {
command_id: "cmd-1".to_string(),
target: Some(CommandTarget {
scope: Scope::Individual as i32,
target_ids: vec!["node-2".to_string()],
}),
..Default::default()
};
let resolution = router.resolve_target(&command);
assert_eq!(
resolution,
TargetResolution::Subordinates(vec!["node-2".to_string()])
);
}
#[test]
fn test_resolve_squad() {
let router = CommandRouter::new(
"node-1".to_string(),
Some("squad-alpha".to_string()),
vec![
"node-1".to_string(),
"node-2".to_string(),
"node-3".to_string(),
],
None,
);
let command = HierarchicalCommand {
command_id: "cmd-1".to_string(),
target: Some(CommandTarget {
scope: Scope::Squad as i32,
target_ids: vec!["squad-alpha".to_string()],
}),
..Default::default()
};
let resolution = router.resolve_target(&command);
if let TargetResolution::AllSquadMembers(members) = resolution {
assert_eq!(members.len(), 3);
} else {
panic!("Expected AllSquadMembers resolution");
}
}
#[test]
fn test_should_route() {
let router = CommandRouter::new(
"node-1".to_string(),
Some("squad-alpha".to_string()),
vec!["node-1".to_string(), "node-2".to_string()],
None,
);
assert!(router.should_route(&TargetResolution::Subordinates(vec!["node-2".to_string()])));
assert!(!router.should_route(&TargetResolution::Self_));
assert!(!router.should_route(&TargetResolution::NotApplicable));
}
}