use crate::discovery::geo::GeoCoordinate;
use crate::models::capability::Capability;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
pub const DEFAULT_GEOHASH_PRECISION: usize = 7;
pub const BEACON_TTL_SECONDS: u64 = 30;
pub const MIN_SQUAD_SIZE: usize = 2;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeographicBeacon {
pub platform_id: String,
pub position: GeoCoordinate,
pub geohash_cell: String,
pub operational: bool,
pub timestamp: u64,
pub capabilities: Vec<String>,
}
impl GeographicBeacon {
pub fn new(
platform_id: String,
position: GeoCoordinate,
capabilities: Vec<Capability>,
) -> Self {
let geohash_cell = encode_geohash(&position, DEFAULT_GEOHASH_PRECISION);
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
Self {
platform_id,
position,
geohash_cell,
operational: true,
timestamp,
capabilities: capabilities.iter().map(|c| c.id.clone()).collect(),
}
}
pub fn is_expired(&self, current_time: u64) -> bool {
current_time.saturating_sub(self.timestamp) > BEACON_TTL_SECONDS
}
}
#[derive(Debug, Clone)]
pub struct GeographicCluster {
pub geohash_cell: String,
pub platforms: Vec<GeographicBeacon>,
pub center: GeoCoordinate,
}
impl GeographicCluster {
pub fn new(geohash_cell: String) -> Result<Self, &'static str> {
let center = decode_geohash(&geohash_cell)?;
Ok(Self {
geohash_cell,
platforms: Vec::new(),
center,
})
}
pub fn add_beacon(&mut self, beacon: GeographicBeacon) {
self.platforms.push(beacon);
}
pub fn remove_expired(&mut self, current_time: u64) {
self.platforms.retain(|b| !b.is_expired(current_time));
}
pub fn can_form_squad(&self, min_size: usize) -> bool {
self.platforms.len() >= min_size
}
pub fn platform_ids(&self) -> Vec<String> {
self.platforms
.iter()
.map(|b| b.platform_id.clone())
.collect()
}
}
pub fn encode_geohash(coord: &GeoCoordinate, precision: usize) -> String {
crate::geohash::encode(coord.lon, coord.lat, precision).expect("Valid coordinate")
}
pub fn decode_geohash(hash: &str) -> Result<GeoCoordinate, &'static str> {
let ((lon, lat), _, _) = crate::geohash::decode(hash).map_err(|_| "Invalid geohash")?;
GeoCoordinate::new(lat, lon, 0.0)
}
pub struct GeographicDiscovery {
clusters: HashMap<String, GeographicCluster>,
my_platform_id: String,
}
impl GeographicDiscovery {
pub fn new(platform_id: String) -> Self {
Self {
clusters: HashMap::new(),
my_platform_id: platform_id,
}
}
pub fn process_beacon(&mut self, beacon: GeographicBeacon) {
let geohash = beacon.geohash_cell.clone();
self.clusters
.entry(geohash.clone())
.or_insert_with(|| GeographicCluster::new(geohash).unwrap())
.add_beacon(beacon);
}
pub fn cleanup_expired(&mut self) {
let current_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
for cluster in self.clusters.values_mut() {
cluster.remove_expired(current_time);
}
self.clusters
.retain(|_, cluster| !cluster.platforms.is_empty());
}
pub fn find_formable_squads(&self, min_size: usize) -> Vec<&GeographicCluster> {
self.clusters
.values()
.filter(|c| c.can_form_squad(min_size))
.collect()
}
pub fn my_cluster(&self) -> Option<&GeographicCluster> {
self.clusters.values().find(|c| {
c.platforms
.iter()
.any(|b| b.platform_id == self.my_platform_id)
})
}
pub fn should_initiate_squad_formation(&self) -> bool {
if let Some(cluster) = self.my_cluster() {
if cluster.can_form_squad(MIN_SQUAD_SIZE) {
if let Some(min_id) = cluster.platforms.iter().map(|b| &b.platform_id).min() {
return min_id == &self.my_platform_id;
}
}
}
false
}
pub fn get_squad_members(&self, max_size: usize) -> Option<Vec<String>> {
if let Some(cluster) = self.my_cluster() {
if cluster.can_form_squad(MIN_SQUAD_SIZE) {
let mut members = cluster.platform_ids();
members.sort(); members.truncate(max_size);
return Some(members);
}
}
None
}
pub fn total_platforms(&self) -> usize {
self.clusters.values().map(|c| c.platforms.len()).sum()
}
pub fn cluster_count(&self) -> usize {
self.clusters.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::CapabilityExt;
#[test]
fn test_geohash_encoding() {
let coord = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let hash = encode_geohash(&coord, 7);
assert_eq!(hash.len(), 7);
assert!(hash.starts_with("9q8yy")); }
#[test]
fn test_geohash_decoding() {
let hash = "9q8yyk8";
let coord = decode_geohash(hash).unwrap();
assert!((coord.lat - 37.77).abs() < 0.01);
assert!((coord.lon - (-122.41)).abs() < 0.01);
}
#[test]
fn test_beacon_creation() {
use crate::models::capability::CapabilityType;
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let caps = vec![
Capability::new(
"intel1".to_string(),
"Intelligence".to_string(),
CapabilityType::Sensor,
0.9,
),
Capability::new(
"comms1".to_string(),
"Communications".to_string(),
CapabilityType::Communication,
0.95,
),
];
let beacon = GeographicBeacon::new("node_1".to_string(), pos, caps);
assert_eq!(beacon.platform_id, "node_1");
assert_eq!(beacon.position, pos);
assert!(beacon.geohash_cell.starts_with("9q8yy"));
assert!(beacon.operational);
assert_eq!(beacon.capabilities.len(), 2);
}
#[test]
fn test_beacon_expiration() {
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let beacon = GeographicBeacon::new("node_1".to_string(), pos, vec![]);
assert!(!beacon.is_expired(beacon.timestamp));
assert!(beacon.is_expired(beacon.timestamp + 31));
}
#[test]
fn test_cluster_creation() {
let cluster = GeographicCluster::new("9q8yyk8".to_string()).unwrap();
assert_eq!(cluster.geohash_cell, "9q8yyk8");
assert_eq!(cluster.platforms.len(), 0);
assert!(!cluster.can_form_squad(2));
}
#[test]
fn test_cluster_beacon_management() {
let mut cluster = GeographicCluster::new("9q8yyk8".to_string()).unwrap();
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let beacon1 = GeographicBeacon::new("node_1".to_string(), pos, vec![]);
let beacon2 = GeographicBeacon::new("node_2".to_string(), pos, vec![]);
cluster.add_beacon(beacon1);
cluster.add_beacon(beacon2);
assert_eq!(cluster.platforms.len(), 2);
assert!(cluster.can_form_squad(2));
let ids = cluster.platform_ids();
assert!(ids.contains(&"node_1".to_string()));
assert!(ids.contains(&"node_2".to_string()));
}
#[test]
fn test_discovery_basic_operations() {
let mut discovery = GeographicDiscovery::new("node_1".to_string());
assert_eq!(discovery.total_platforms(), 0);
assert_eq!(discovery.cluster_count(), 0);
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let beacon = GeographicBeacon::new("node_2".to_string(), pos, vec![]);
discovery.process_beacon(beacon);
assert_eq!(discovery.total_platforms(), 1);
assert_eq!(discovery.cluster_count(), 1);
}
#[test]
fn test_discovery_squad_formation() {
let mut discovery = GeographicDiscovery::new("node_1".to_string());
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let beacon1 = GeographicBeacon::new("node_1".to_string(), pos, vec![]);
discovery.process_beacon(beacon1);
let beacon2 = GeographicBeacon::new("node_2".to_string(), pos, vec![]);
discovery.process_beacon(beacon2);
assert_eq!(discovery.total_platforms(), 2);
let formable = discovery.find_formable_squads(2);
assert_eq!(formable.len(), 1);
assert!(discovery.should_initiate_squad_formation());
let members = discovery.get_squad_members(5).unwrap();
assert_eq!(members.len(), 2);
assert!(members.contains(&"node_1".to_string()));
assert!(members.contains(&"node_2".to_string()));
}
#[test]
fn test_discovery_multiple_clusters() {
let mut discovery = GeographicDiscovery::new("node_1".to_string());
let pos1 = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let beacon1 = GeographicBeacon::new("node_1".to_string(), pos1, vec![]);
let beacon2 = GeographicBeacon::new("node_2".to_string(), pos1, vec![]);
let pos2 = GeoCoordinate::new(34.0522, -118.2437, 100.0).unwrap();
let beacon3 = GeographicBeacon::new("node_3".to_string(), pos2, vec![]);
let beacon4 = GeographicBeacon::new("platform_4".to_string(), pos2, vec![]);
discovery.process_beacon(beacon1);
discovery.process_beacon(beacon2);
discovery.process_beacon(beacon3);
discovery.process_beacon(beacon4);
assert_eq!(discovery.total_platforms(), 4);
assert_eq!(discovery.cluster_count(), 2);
let formable = discovery.find_formable_squads(2);
assert_eq!(formable.len(), 2);
}
#[test]
fn test_discovery_cleanup_expired() {
let mut discovery = GeographicDiscovery::new("node_1".to_string());
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
let mut beacon = GeographicBeacon::new("node_2".to_string(), pos, vec![]);
beacon.timestamp = 0;
discovery.process_beacon(beacon);
assert_eq!(discovery.total_platforms(), 1);
discovery.cleanup_expired();
assert_eq!(discovery.total_platforms(), 0);
assert_eq!(discovery.cluster_count(), 0);
}
#[test]
fn test_deterministic_leader_selection() {
let mut discovery1 = GeographicDiscovery::new("platform_a".to_string());
let mut discovery2 = GeographicDiscovery::new("platform_b".to_string());
let mut discovery3 = GeographicDiscovery::new("platform_c".to_string());
let pos = GeoCoordinate::new(37.7749, -122.4194, 100.0).unwrap();
for id in ["platform_a", "platform_b", "platform_c"] {
let beacon = GeographicBeacon::new(id.to_string(), pos, vec![]);
discovery1.process_beacon(beacon.clone());
discovery2.process_beacon(beacon.clone());
discovery3.process_beacon(beacon);
}
assert!(discovery1.should_initiate_squad_formation());
assert!(!discovery2.should_initiate_squad_formation());
assert!(!discovery3.should_initiate_squad_formation());
}
}