oxirs-vec 0.2.4

//! Distributed Vector Search - Version 1.1 Roadmap Feature
//!
//! This module implements distributed vector search capabilities for scaling
//! vector operations across multiple nodes and data centers.

use crate::{
    advanced_analytics::VectorAnalyticsEngine,
    similarity::{SimilarityMetric, SimilarityResult},
    Vector,
};

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::Mutex;
use tracing::{debug, error, info};

/// Distributed node configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedNodeConfig {
    /// Unique node identifier
    pub node_id: String,
    /// Node endpoint URL
    pub endpoint: String,
    /// Node region/datacenter
    pub region: String,
    /// Node capacity (max vectors)
    pub capacity: usize,
    /// Current load factor (0.0 to 1.0)
    pub load_factor: f32,
    /// Network latency to this node (ms)
    pub latency_ms: u64,
    /// Node health status
    pub health_status: NodeHealthStatus,
    /// Replication factor
    pub replication_factor: usize,
}

/// Node health status
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum NodeHealthStatus {
    Healthy,
    Degraded,
    Unhealthy,
    Offline,
}

/// Distributed search query
#[derive(Debug, Clone)]
pub struct DistributedQuery {
    pub id: String,
    pub query_vector: Vector,
    pub k: usize,
    pub similarity_metric: SimilarityMetric,
    pub filters: HashMap<String, String>,
    pub timeout: Duration,
    pub consistency_level: ConsistencyLevel,
}

/// Consistency levels for distributed queries
#[derive(Debug, Clone, Copy)]
pub enum ConsistencyLevel {
    /// Read from any available node
    Eventual,
    /// Read from majority of nodes
    Quorum,
    /// Read from all nodes
    Strong,
}

/// Search result from a distributed node
#[derive(Debug, Clone)]
pub struct NodeSearchResult {
    pub node_id: String,
    pub results: Vec<SimilarityResult>,
    pub latency_ms: u64,
    pub error: Option<String>,
}

/// Distributed search response
#[derive(Debug, Clone)]
pub struct DistributedSearchResponse {
    pub query_id: String,
    pub merged_results: Vec<SimilarityResult>,
    pub node_results: Vec<NodeSearchResult>,
    pub total_latency_ms: u64,
    pub nodes_queried: usize,
    pub nodes_responded: usize,
}

/// Partitioning strategy for vector distribution
#[derive(Debug, Clone)]
pub enum PartitioningStrategy {
    /// Hash-based partitioning
    Hash,
    /// Range-based partitioning
    Range,
    /// Consistent hashing
    ConsistentHash,
    /// Geography-based partitioning
    Geographic,
    /// Custom partitioning function
    Custom(fn(&Vector) -> String),
}

/// Distributed vector search coordinator
pub struct DistributedVectorSearch {
    /// Node registry
    nodes: Arc<RwLock<HashMap<String, DistributedNodeConfig>>>,
    /// Partitioning strategy
    partitioning_strategy: PartitioningStrategy,
    /// Load balancer
    load_balancer: Arc<Mutex<LoadBalancer>>,
    /// Replication manager
    replication_manager: Arc<Mutex<ReplicationManager>>,
    /// Query router
    query_router: Arc<QueryRouter>,
    /// Health monitor
    health_monitor: Arc<Mutex<HealthMonitor>>,
    /// Performance analytics
    analytics: Arc<Mutex<VectorAnalyticsEngine>>,
}

/// Load balancer for distributed queries
#[derive(Debug)]
pub struct LoadBalancer {
    /// Load balancing algorithm
    algorithm: LoadBalancingAlgorithm,
    /// Node usage statistics
    node_stats: HashMap<String, NodeStats>,
}

/// Load balancing algorithms
#[derive(Debug, Clone)]
pub enum LoadBalancingAlgorithm {
    RoundRobin,
    LeastConnections,
    WeightedRoundRobin,
    LatencyBased,
    ResourceBased,
}

/// Node statistics for load balancing
#[derive(Debug, Clone)]
pub struct NodeStats {
    pub active_queries: u64,
    pub average_latency_ms: f64,
    pub success_rate: f64,
    pub last_updated: SystemTime,
}

/// Replication manager for data consistency
#[derive(Debug)]
pub struct ReplicationManager {
    /// Replication configurations per partition
    partition_replicas: HashMap<String, Vec<String>>,
    /// Consistency policies
    consistency_policies: HashMap<String, ConsistencyLevel>,
}

/// Query router for distributed search
pub struct QueryRouter {
    /// Routing table
    routing_table: Arc<RwLock<HashMap<String, Vec<String>>>>,
    /// Query execution strategy
    execution_strategy: QueryExecutionStrategy,
}

/// Query execution strategies
#[derive(Debug, Clone)]
pub enum QueryExecutionStrategy {
    /// Execute on all relevant nodes in parallel
    Parallel,
    /// Execute on nodes sequentially with early termination
    Sequential,
    /// Adaptive execution based on query characteristics
    Adaptive,
}

/// Health monitor for distributed nodes
#[derive(Debug)]
pub struct HealthMonitor {
    /// Health check interval
    check_interval: Duration,
    /// Health check timeout
    check_timeout: Duration,
    /// Node health history
    health_history: HashMap<String, Vec<HealthCheckResult>>,
}

/// Health check result
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
    pub timestamp: SystemTime,
    pub latency_ms: u64,
    pub success: bool,
    pub error_message: Option<String>,
}

impl DistributedVectorSearch {
    /// Create new distributed vector search coordinator
    pub fn new(partitioning_strategy: PartitioningStrategy) -> Result<Self> {
        Ok(Self {
            nodes: Arc::new(RwLock::new(HashMap::new())),
            partitioning_strategy,
            load_balancer: Arc::new(Mutex::new(LoadBalancer::new(
                LoadBalancingAlgorithm::LatencyBased,
            ))),
            replication_manager: Arc::new(Mutex::new(ReplicationManager::new())),
            query_router: Arc::new(QueryRouter::new(QueryExecutionStrategy::Adaptive)),
            health_monitor: Arc::new(Mutex::new(HealthMonitor::new(
                Duration::from_secs(30),
                Duration::from_secs(5),
            ))),
            analytics: Arc::new(Mutex::new(VectorAnalyticsEngine::new())),
        })
    }

    /// Register a new node in the cluster
    pub async fn register_node(&self, config: DistributedNodeConfig) -> Result<()> {
        {
            let mut nodes = self
                .nodes
                .write()
                .expect("nodes lock should not be poisoned");
            info!("Registering node {} at {}", config.node_id, config.endpoint);
            nodes.insert(config.node_id.clone(), config.clone());
        } // Drop nodes lock before await

        // Update load balancer
        let mut load_balancer = self.load_balancer.lock().await;
        load_balancer.add_node(&config.node_id);

        // Update replication manager
        let mut replication_manager = self.replication_manager.lock().await;
        replication_manager.add_node(&config.node_id, config.replication_factor);

        // Start health monitoring for the new node
        let mut health_monitor = self.health_monitor.lock().await;
        health_monitor.start_monitoring(&config.node_id);

        Ok(())
    }

    /// Remove a node from the cluster
    pub async fn deregister_node(&self, node_id: &str) -> Result<()> {
        let config = {
            let mut nodes = self
                .nodes
                .write()
                .expect("nodes lock should not be poisoned");
            nodes.remove(node_id)
        }; // Drop nodes lock before await

        if let Some(config) = config {
            info!("Deregistering node {} at {}", node_id, config.endpoint);

            // Update load balancer
            let mut load_balancer = self.load_balancer.lock().await;
            load_balancer.remove_node(node_id);

            // Update replication manager
            let mut replication_manager = self.replication_manager.lock().await;
            replication_manager.remove_node(node_id);

            // Stop health monitoring
            let mut health_monitor = self.health_monitor.lock().await;
            health_monitor.stop_monitoring(node_id);
        }

        Ok(())
    }

    /// Execute distributed vector search
    pub async fn search(&self, query: DistributedQuery) -> Result<DistributedSearchResponse> {
        let start_time = Instant::now();

        // Determine target nodes based on query and partitioning strategy
        let target_nodes = self.select_target_nodes(&query).await?;

        info!(
            "Executing distributed query {} across {} nodes",
            query.id,
            target_nodes.len()
        );

        // Execute query on selected nodes
        let node_results = match self.query_router.execution_strategy {
            QueryExecutionStrategy::Parallel => {
                self.execute_parallel_query(&query, &target_nodes).await?
            }
            QueryExecutionStrategy::Sequential => {
                self.execute_sequential_query(&query, &target_nodes).await?
            }
            QueryExecutionStrategy::Adaptive => {
                self.execute_adaptive_query(&query, &target_nodes).await?
            }
        };

        // Merge results from all nodes
        let merged_results = self.merge_node_results(&node_results, query.k);

        // Update analytics
        let analytics = crate::advanced_analytics::QueryAnalytics {
            query_id: query.id.clone(),
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            query_vector: query.query_vector.as_f32(),
            similarity_metric: "distributed".to_string(),
            top_k: query.k,
            response_time: start_time.elapsed(),
            results_count: merged_results.len(),
            avg_similarity_score: merged_results.iter().map(|r| r.similarity).sum::<f32>()
                / merged_results.len().max(1) as f32,
            min_similarity_score: merged_results
                .iter()
                .map(|r| r.similarity)
                .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .unwrap_or(0.0),
            max_similarity_score: merged_results
                .iter()
                .map(|r| r.similarity)
                .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .unwrap_or(0.0),
            cache_hit: false,
            index_type: "distributed".to_string(),
        };
        let mut analytics_guard = self.analytics.lock().await;
        analytics_guard.record_query(analytics);

        let nodes_responded = node_results.len();
        Ok(DistributedSearchResponse {
            query_id: query.id,
            merged_results,
            node_results,
            total_latency_ms: start_time.elapsed().as_millis() as u64,
            nodes_queried: target_nodes.len(),
            nodes_responded,
        })
    }

    /// Select target nodes for query execution
    async fn select_target_nodes(&self, query: &DistributedQuery) -> Result<Vec<String>> {
        let nodes = self
            .nodes
            .read()
            .expect("nodes lock should not be poisoned")
            .clone();
        let load_balancer = self.load_balancer.lock().await;

        match &self.partitioning_strategy {
            PartitioningStrategy::Hash => {
                // Hash-based selection
                let partition = self.compute_hash_partition(&query.query_vector);
                self.get_nodes_for_partition(&partition, &nodes, &load_balancer)
            }
            PartitioningStrategy::Range => {
                // Range-based selection
                let partition = self.compute_range_partition(&query.query_vector);
                self.get_nodes_for_partition(&partition, &nodes, &load_balancer)
            }
            PartitioningStrategy::ConsistentHash => {
                // Consistent hash selection
                let partition = self.compute_consistent_hash_partition(&query.query_vector);
                self.get_nodes_for_partition(&partition, &nodes, &load_balancer)
            }
            PartitioningStrategy::Geographic => {
                // Geographic-based selection (use all healthy nodes)
                Ok(nodes
                    .iter()
                    .filter(|(_, config)| config.health_status == NodeHealthStatus::Healthy)
                    .map(|(id, _)| id.clone())
                    .collect())
            }
            PartitioningStrategy::Custom(_func) => {
                // Custom partitioning function
                Ok(nodes.keys().cloned().collect())
            }
        }
    }

    /// Execute query in parallel across nodes
    async fn execute_parallel_query(
        &self,
        query: &DistributedQuery,
        target_nodes: &[String],
    ) -> Result<Vec<NodeSearchResult>> {
        let mut handles = Vec::new();

        for node_id in target_nodes {
            let node_id = node_id.clone();
            let query = query.clone();
            let nodes = Arc::clone(&self.nodes);

            let handle =
                tokio::spawn(async move { Self::execute_node_query(node_id, query, nodes).await });

            handles.push(handle);
        }

        let mut results = Vec::new();
        for handle in handles {
            match handle.await {
                Ok(Ok(result)) => results.push(result),
                Ok(Err(e)) => error!("Node query failed: {}", e),
                Err(e) => error!("Task join failed: {}", e),
            }
        }

        Ok(results)
    }

    /// Execute query sequentially across nodes
    async fn execute_sequential_query(
        &self,
        query: &DistributedQuery,
        target_nodes: &[String],
    ) -> Result<Vec<NodeSearchResult>> {
        let mut results = Vec::new();

        for node_id in target_nodes {
            match Self::execute_node_query(node_id.clone(), query.clone(), Arc::clone(&self.nodes))
                .await
            {
                Ok(result) => {
                    results.push(result);
                    // Early termination if we have enough results
                    if results.len() >= query.k {
                        break;
                    }
                }
                Err(e) => {
                    error!("Node query failed for {}: {}", node_id, e);
                    continue;
                }
            }
        }

        Ok(results)
    }

    /// Execute query with adaptive strategy
    async fn execute_adaptive_query(
        &self,
        query: &DistributedQuery,
        target_nodes: &[String],
    ) -> Result<Vec<NodeSearchResult>> {
        // For demonstration, use parallel for small node counts, sequential for large
        if target_nodes.len() <= 5 {
            self.execute_parallel_query(query, target_nodes).await
        } else {
            self.execute_sequential_query(query, target_nodes).await
        }
    }

    /// Execute query on a specific node
    async fn execute_node_query(
        node_id: String,
        query: DistributedQuery,
        nodes: Arc<RwLock<HashMap<String, DistributedNodeConfig>>>,
    ) -> Result<NodeSearchResult> {
        let start_time = Instant::now();

        // In a real implementation, this would make HTTP requests to the node
        // For now, simulate the query execution

        {
            let nodes_guard = nodes.read().expect("nodes lock should not be poisoned");
            let _node_config = nodes_guard
                .get(&node_id)
                .ok_or_else(|| anyhow::anyhow!("Node {} not found", node_id))?;
        } // Drop the guard here

        // Simulate network latency and processing
        tokio::time::sleep(Duration::from_millis(10)).await;

        // Generate sample results
        let mut results = Vec::new();
        for i in 0..query.k.min(10) {
            results.push(SimilarityResult {
                id: format!(
                    "dist_{}_{}_{}",
                    node_id,
                    i,
                    std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_millis()
                ),
                uri: format!("{node_id}:vector_{i}"),
                similarity: 0.9 - (i as f32 * 0.1),
                metadata: Some(HashMap::new()),
                metrics: HashMap::new(),
            });
        }

        Ok(NodeSearchResult {
            node_id,
            results,
            latency_ms: start_time.elapsed().as_millis() as u64,
            error: None,
        })
    }

    /// Merge results from multiple nodes
    fn merge_node_results(
        &self,
        node_results: &[NodeSearchResult],
        k: usize,
    ) -> Vec<SimilarityResult> {
        let mut all_results = Vec::new();

        // Collect all results
        for node_result in node_results {
            all_results.extend(node_result.results.clone());
        }

        // Sort by similarity score (descending)
        all_results.sort_by(|a, b| {
            b.similarity
                .partial_cmp(&a.similarity)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Take top k results
        all_results.truncate(k);
        all_results
    }

    /// Compute hash partition for vector
    fn compute_hash_partition(&self, vector: &Vector) -> String {
        let values = vector.as_f32();
        let mut hash = 0u64;
        for &value in &values {
            hash = hash.wrapping_mul(31).wrapping_add(value.to_bits() as u64);
        }
        format!("partition_{}", hash % 10) // 10 partitions
    }

    /// Compute range partition for vector
    fn compute_range_partition(&self, vector: &Vector) -> String {
        let values = vector.as_f32();
        let sum: f32 = values.iter().sum();
        let partition_id = (sum.abs() % 10.0) as usize;
        format!("partition_{partition_id}")
    }

    /// Compute consistent hash partition for vector
    fn compute_consistent_hash_partition(&self, vector: &Vector) -> String {
        // Simplified consistent hashing
        self.compute_hash_partition(vector)
    }

    /// Get nodes for a specific partition
    fn get_nodes_for_partition(
        &self,
        _partition: &str,
        nodes: &HashMap<String, DistributedNodeConfig>,
        _load_balancer: &LoadBalancer,
    ) -> Result<Vec<String>> {
        // Simplified implementation - return all healthy nodes
        Ok(nodes
            .iter()
            .filter(|(_, config)| config.health_status == NodeHealthStatus::Healthy)
            .map(|(id, _)| id.clone())
            .collect())
    }

    /// Get cluster statistics
    pub fn get_cluster_stats(&self) -> DistributedClusterStats {
        let nodes = self
            .nodes
            .read()
            .expect("nodes lock should not be poisoned");

        let total_nodes = nodes.len();
        let healthy_nodes = nodes
            .values()
            .filter(|config| config.health_status == NodeHealthStatus::Healthy)
            .count();

        let total_capacity: usize = nodes.values().map(|config| config.capacity).sum();
        let average_load_factor = if !nodes.is_empty() {
            nodes.values().map(|config| config.load_factor).sum::<f32>() / nodes.len() as f32
        } else {
            0.0
        };

        DistributedClusterStats {
            total_nodes,
            healthy_nodes,
            total_capacity,
            average_load_factor,
            partitioning_strategy: format!("{:?}", self.partitioning_strategy),
        }
    }
}

/// Cluster statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedClusterStats {
    pub total_nodes: usize,
    pub healthy_nodes: usize,
    pub total_capacity: usize,
    pub average_load_factor: f32,
    pub partitioning_strategy: String,
}

impl LoadBalancer {
    fn new(algorithm: LoadBalancingAlgorithm) -> Self {
        Self {
            algorithm,
            node_stats: HashMap::new(),
        }
    }

    fn add_node(&mut self, node_id: &str) {
        self.node_stats.insert(
            node_id.to_string(),
            NodeStats {
                active_queries: 0,
                average_latency_ms: 0.0,
                success_rate: 1.0,
                last_updated: SystemTime::now(),
            },
        );
    }

    fn remove_node(&mut self, node_id: &str) {
        self.node_stats.remove(node_id);
    }
}

impl ReplicationManager {
    fn new() -> Self {
        Self {
            partition_replicas: HashMap::new(),
            consistency_policies: HashMap::new(),
        }
    }

    fn add_node(&mut self, node_id: &str, _replication_factor: usize) {
        // Add node to replication topology
        debug!("Adding node {} to replication topology", node_id);
    }

    fn remove_node(&mut self, node_id: &str) {
        // Remove node from replication topology
        debug!("Removing node {} from replication topology", node_id);
    }
}

impl QueryRouter {
    fn new(execution_strategy: QueryExecutionStrategy) -> Self {
        Self {
            routing_table: Arc::new(RwLock::new(HashMap::new())),
            execution_strategy,
        }
    }
}

impl HealthMonitor {
    fn new(check_interval: Duration, check_timeout: Duration) -> Self {
        Self {
            check_interval,
            check_timeout,
            health_history: HashMap::new(),
        }
    }

    fn start_monitoring(&mut self, node_id: &str) {
        self.health_history.insert(node_id.to_string(), Vec::new());
        debug!("Started health monitoring for node {}", node_id);
    }

    fn stop_monitoring(&mut self, node_id: &str) {
        self.health_history.remove(node_id);
        debug!("Stopped health monitoring for node {}", node_id);
    }
}

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

    #[tokio::test]
    async fn test_distributed_search_creation() {
        let distributed_search = DistributedVectorSearch::new(PartitioningStrategy::Hash);
        assert!(distributed_search.is_ok());
    }

    #[tokio::test]
    async fn test_node_registration() -> Result<()> {
        let distributed_search = DistributedVectorSearch::new(PartitioningStrategy::Hash)?;

        let config = DistributedNodeConfig {
            node_id: "node1".to_string(),
            endpoint: "http://localhost:8080".to_string(),
            region: "us-west-1".to_string(),
            capacity: 100000,
            load_factor: 0.5,
            latency_ms: 10,
            health_status: NodeHealthStatus::Healthy,
            replication_factor: 3,
        };

        assert!(distributed_search.register_node(config).await.is_ok());

        let stats = distributed_search.get_cluster_stats();
        assert_eq!(stats.total_nodes, 1);
        assert_eq!(stats.healthy_nodes, 1);
        Ok(())
    }

    #[tokio::test]
    async fn test_distributed_query_execution() -> Result<()> {
        let distributed_search = DistributedVectorSearch::new(PartitioningStrategy::Hash)?;

        // Register test nodes
        for i in 0..3 {
            let config = DistributedNodeConfig {
                node_id: format!("node{i}"),
                endpoint: format!("http://localhost:808{i}"),
                region: "us-west-1".to_string(),
                capacity: 100000,
                load_factor: 0.3,
                latency_ms: 5 + i * 2,
                health_status: NodeHealthStatus::Healthy,
                replication_factor: 2,
            };
            distributed_search.register_node(config).await?;
        }

        // Create test query
        let query = DistributedQuery {
            id: "test_query_1".to_string(),
            query_vector: crate::Vector::new(vec![1.0, 0.5, 0.8]),
            k: 10,
            similarity_metric: SimilarityMetric::Cosine,
            filters: HashMap::new(),
            timeout: Duration::from_secs(5),
            consistency_level: ConsistencyLevel::Quorum,
        };

        let response = distributed_search.search(query).await?;

        assert_eq!(response.nodes_queried, 3);
        assert!(response.nodes_responded > 0);
        assert!(!response.merged_results.is_empty());
        Ok(())
    }
}