poimen-memory/crates/mem-cli/src/federation.rs

/// Federation Layer: Coordinate queries across multiple memory instances/projects
///
/// Provides:
/// - Multi-instance coordination (round-robin, load-balancing)
/// - Project federation (query across related projects)
/// - Result merging and deduplication
/// - Distributed ranking
/// - Failure resilience (fallback to other instances)

use anyhow::Result;
use std::sync::Arc;
use std::collections::HashMap;

/// Instance metadata
#[derive(Debug, Clone)]
pub struct InstanceMetadata {
    pub id: String,
    pub name: String,
    pub region: String,
    pub is_healthy: bool,
    pub latency_ms: u64,
    pub load_percent: f32,
}

impl InstanceMetadata {
    pub fn new(id: &str, name: &str, region: &str) -> Self {
        Self {
            id: id.to_string(),
            name: name.to_string(),
            region: region.to_string(),
            is_healthy: true,
            latency_ms: 0,
            load_percent: 0.0,
        }
    }

    /// Calculate health score (0-1)
    pub fn health_score(&self) -> f32 {
        if !self.is_healthy {
            return 0.0;
        }

        let latency_penalty = (self.latency_ms as f32 / 1000.0).min(1.0);
        let load_penalty = self.load_percent / 100.0;

        ((1.0 - latency_penalty) * 0.6 + (1.0 - load_penalty) * 0.4).max(0.0)
    }
}

/// Distributed query result
#[derive(Debug, Clone)]
pub struct FederatedResult {
    pub instance_id: String,
    pub result_id: String,
    pub text: String,
    pub score: f32,
    pub latency_ms: u64,
}

impl FederatedResult {
    pub fn new(instance_id: &str, result_id: &str, text: &str, score: f32) -> Self {
        Self {
            instance_id: instance_id.to_string(),
            result_id: result_id.to_string(),
            text: text.to_string(),
            score,
            latency_ms: 0,
        }
    }
}

/// Result deduplicator
pub struct ResultDeduplicator {
    similarity_threshold: f32,
}

impl ResultDeduplicator {
    pub fn new(similarity_threshold: f32) -> Self {
        Self {
            similarity_threshold,
        }
    }

    /// Simple Jaccard similarity
    fn similarity(&self, text_a: &str, text_b: &str) -> f32 {
        let text_a_lower = text_a.to_lowercase();
        let text_b_lower = text_b.to_lowercase();
        let words_a: std::collections::HashSet<_> = text_a_lower
            .split_whitespace()
            .collect();
        let words_b: std::collections::HashSet<_> = text_b_lower
            .split_whitespace()
            .collect();

        let intersection = words_a.intersection(&words_b).count();
        let union = words_a.union(&words_b).count();

        if union == 0 {
            0.0
        } else {
            intersection as f32 / union as f32
        }
    }

    /// Deduplicate results
    pub fn deduplicate(&self, results: Vec<FederatedResult>) -> Vec<FederatedResult> {
        let mut unique = Vec::new();

        for result in results {
            let is_duplicate = unique.iter().any(|kept: &FederatedResult| {
                self.similarity(&result.text, &kept.text) > self.similarity_threshold
            });

            if !is_duplicate {
                unique.push(result);
            }
        }

        unique
    }
}

/// Instance selector (routing strategy)
pub trait InstanceSelector: Send + Sync {
    fn select<'a>(&self, instances: &'a [InstanceMetadata]) -> Option<&'a InstanceMetadata>;
}

/// Round-robin selector
pub struct RoundRobinSelector {
    counter: std::sync::atomic::AtomicUsize,
}

impl RoundRobinSelector {
    pub fn new() -> Self {
        Self {
            counter: std::sync::atomic::AtomicUsize::new(0),
        }
    }
}

impl InstanceSelector for RoundRobinSelector {
    fn select<'a>(&self, instances: &'a [InstanceMetadata]) -> Option<&'a InstanceMetadata> {
        if instances.is_empty() {
            return None;
        }

        let healthy: Vec<_> = instances.iter().filter(|i| i.is_healthy).collect();

        if healthy.is_empty() {
            return None;
        }

        let idx = self
            .counter
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
            % healthy.len();

        Some(healthy[idx])
    }
}

/// Health-based selector (prefer healthier instances)
pub struct HealthBasedSelector;

impl InstanceSelector for HealthBasedSelector {
    fn select<'a>(&self, instances: &'a [InstanceMetadata]) -> Option<&'a InstanceMetadata> {
        instances
            .iter()
            .filter(|i| i.is_healthy)
            .max_by(|a, b| {
                a.health_score()
                    .partial_cmp(&b.health_score())
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|r| r)
    }
}

/// Federation coordinator
pub struct FederationCoordinator {
    instances: HashMap<String, InstanceMetadata>,
    selector: Arc<dyn InstanceSelector>,
    deduplicator: ResultDeduplicator,
}

impl FederationCoordinator {
    pub fn new(selector: Arc<dyn InstanceSelector>) -> Self {
        Self {
            instances: HashMap::new(),
            selector,
            deduplicator: ResultDeduplicator::new(0.7),
        }
    }

    pub fn register_instance(&mut self, instance: InstanceMetadata) {
        self.instances.insert(instance.id.clone(), instance);
    }

    pub fn unregister_instance(&mut self, instance_id: &str) {
        self.instances.remove(instance_id);
    }

    pub fn update_instance_health(&mut self, instance_id: &str, is_healthy: bool) {
        if let Some(instance) = self.instances.get_mut(instance_id) {
            instance.is_healthy = is_healthy;
        }
    }

    pub fn update_instance_metrics(&mut self, instance_id: &str, latency_ms: u64, load_percent: f32) {
        if let Some(instance) = self.instances.get_mut(instance_id) {
            instance.latency_ms = latency_ms;
            instance.load_percent = load_percent;
        }
    }

    /// Select best instance for query
    pub fn select_instance(&self) -> Result<String> {
        let instances: Vec<InstanceMetadata> = self.instances.values().cloned().collect();
        self.selector
            .select(&instances)
            .map(|i| i.id.clone())
            .ok_or_else(|| anyhow::anyhow!("No healthy instances available"))
    }

    /// Merge results from multiple instances
    pub fn merge_results(&self, results: Vec<FederatedResult>, top_k: usize) -> Vec<FederatedResult> {
        // Deduplicate
        let deduplicated = self.deduplicator.deduplicate(results);

        // Sort by score
        let mut sorted = deduplicated;
        sorted.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        sorted.into_iter().take(top_k).collect()
    }

    pub fn get_instance(&self, instance_id: &str) -> Option<&InstanceMetadata> {
        self.instances.get(instance_id)
    }

    pub fn get_healthy_instances(&self) -> Vec<&InstanceMetadata> {
        self.instances
            .values()
            .filter(|i| i.is_healthy)
            .collect()
    }

    pub fn total_instances(&self) -> usize {
        self.instances.len()
    }
}

/// Multi-project query coordinator
pub struct MultiProjectCoordinator {
    projects: HashMap<String, String>, // project_name -> instance_id
    coordinator: Arc<FederationCoordinator>,
}

impl MultiProjectCoordinator {
    pub fn new(coordinator: Arc<FederationCoordinator>) -> Self {
        Self {
            projects: HashMap::new(),
            coordinator,
        }
    }

    pub fn register_project(&mut self, project: &str, instance_id: &str) {
        self.projects.insert(project.to_string(), instance_id.to_string());
    }

    pub fn get_instance_for_project(&self, project: &str) -> Result<Option<&InstanceMetadata>> {
        if let Some(instance_id) = self.projects.get(project) {
            Ok(self.coordinator.get_instance(instance_id))
        } else {
            Ok(None)
        }
    }

    pub fn list_projects(&self) -> Vec<&str> {
        self.projects.keys().map(|s| s.as_str()).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_instance_metadata_creation() {
        let instance = InstanceMetadata::new("inst1", "primary", "us-east");
        assert_eq!(instance.id, "inst1");
        assert!(instance.is_healthy);
    }

    #[test]
    fn test_instance_health_score_healthy() {
        let instance = InstanceMetadata::new("inst1", "primary", "us-east");
        let score = instance.health_score();
        assert!(score > 0.9);
    }

    #[test]
    fn test_instance_health_score_unhealthy() {
        let mut instance = InstanceMetadata::new("inst1", "primary", "us-east");
        instance.is_healthy = false;
        let score = instance.health_score();
        assert_eq!(score, 0.0);
    }

    #[test]
    fn test_federated_result_creation() {
        let result = FederatedResult::new("inst1", "doc1", "text", 0.9);
        assert_eq!(result.instance_id, "inst1");
        assert_eq!(result.score, 0.9);
    }

    #[test]
    fn test_deduplicator_exact_duplicates() {
        let dedup = ResultDeduplicator::new(0.8);
        let results = vec![
            FederatedResult::new("inst1", "doc1", "kubernetes pod debugging", 0.9),
            FederatedResult::new("inst2", "doc2", "kubernetes pod debugging", 0.85),
        ];

        let unique = dedup.deduplicate(results);
        assert_eq!(unique.len(), 1);
    }

    #[test]
    fn test_deduplicator_different() {
        let dedup = ResultDeduplicator::new(0.8);
        let results = vec![
            FederatedResult::new("inst1", "doc1", "kubernetes pod", 0.9),
            FederatedResult::new("inst2", "doc2", "docker container", 0.8),
        ];

        let unique = dedup.deduplicate(results);
        assert_eq!(unique.len(), 2);
    }

    #[test]
    fn test_round_robin_selector() {
        let selector = RoundRobinSelector::new();
        let instances = vec![
            InstanceMetadata::new("inst1", "primary", "us-east"),
            InstanceMetadata::new("inst2", "secondary", "us-west"),
        ];

        let selected1 = selector.select(&instances);
        assert!(selected1.is_some());

        let selected2 = selector.select(&instances);
        assert!(selected2.is_some());

        // Should be different (round-robin)
        assert_ne!(selected1.unwrap().id, selected2.unwrap().id);
    }

    #[test]
    fn test_health_based_selector() {
        let selector = HealthBasedSelector;
        let mut instances = vec![
            InstanceMetadata::new("inst1", "primary", "us-east"),
            InstanceMetadata::new("inst2", "secondary", "us-west"),
        ];

        instances[0].latency_ms = 500; // Slower
        instances[1].latency_ms = 100; // Faster

        let selected = selector.select(&instances);
        assert_eq!(selected.unwrap().id, "inst2"); // Should select faster one
    }

    #[test]
    fn test_federation_coordinator_register() {
        let coordinator = FederationCoordinator::new(Arc::new(RoundRobinSelector::new()));
        let mut coord = coordinator;

        let instance = InstanceMetadata::new("inst1", "primary", "us-east");
        coord.register_instance(instance);

        assert_eq!(coord.total_instances(), 1);
    }

    #[test]
    fn test_federation_coordinator_select() {
        let selector = Arc::new(RoundRobinSelector::new());
        let mut coordinator = FederationCoordinator::new(selector);

        let instance = InstanceMetadata::new("inst1", "primary", "us-east");
        coordinator.register_instance(instance);

        let selected = coordinator.select_instance();
        assert!(selected.is_ok());
    }

    #[test]
    fn test_federation_coordinator_merge_results() {
        let coordinator = FederationCoordinator::new(Arc::new(RoundRobinSelector::new()));

        let results = vec![
            FederatedResult::new("inst1", "doc1", "text1", 0.9),
            FederatedResult::new("inst2", "doc2", "text2", 0.8),
            FederatedResult::new("inst3", "doc3", "text3", 0.7),
        ];

        let merged = coordinator.merge_results(results, 2);
        assert_eq!(merged.len(), 2);
        assert_eq!(merged[0].score, 0.9); // Highest score first
    }

    #[test]
    fn test_multi_project_coordinator() {
        let coordinator = Arc::new(FederationCoordinator::new(Arc::new(RoundRobinSelector::new())));
        let mut multi = MultiProjectCoordinator::new(coordinator);

        multi.register_project("poimen", "inst1");
        multi.register_project("rust-guide", "inst2");

        assert_eq!(multi.list_projects().len(), 2);
    }
}