//! Shard Router for distributing writes across multiple shards
//!
//! Routes write requests to appropriate shards based on configurable
//! strategies (hash, round-robin, collection-based).

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;

use super::write_shard::WriteShard;
use super::write_queue::WriteResult;
use super::backpressure::{BackpressureController, BackpressureError};
use super::metrics::WriteMetrics;
use crate::types::MemoryNode;

/// Configuration for shard router
#[derive(Debug, Clone)]
pub struct ShardRouterConfig {
    /// Number of shards
    pub shard_count: usize,
    
    /// Routing strategy
    pub strategy: ShardStrategy,
    
    /// Whether to auto-balance load
    pub auto_balance: bool,
}

impl Default for ShardRouterConfig {
    fn default() -> Self {
        Self {
            shard_count: num_cpus::get() * 2,
            strategy: ShardStrategy::Hash,
            auto_balance: true,
        }
    }
}

/// Strategy for routing writes to shards
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShardStrategy {
    /// Route based on document ID hash
    Hash,
    
    /// Round-robin distribution
    RoundRobin,
    
    /// Route based on collection name (if available)
    Collection,
}

/// Shard Router
///
/// Distributes write requests across multiple shards for parallel processing.
pub struct ShardRouter {
    /// Shards
    shards: Vec<Arc<WriteShard>>,
    
    /// Configuration
    config: ShardRouterConfig,
    
    /// Round-robin counter
    round_robin_counter: AtomicUsize,
    
    /// Backpressure controller
    backpressure: BackpressureController,
    
    /// Metrics
    metrics: Arc<WriteMetrics>,
}

impl ShardRouter {
    /// Create a new shard router
    pub fn new(
        shards: Vec<Arc<WriteShard>>,
        config: ShardRouterConfig,
        backpressure: BackpressureController,
        metrics: Arc<WriteMetrics>,
    ) -> Self {
        Self {
            shards,
            config,
            round_robin_counter: AtomicUsize::new(0),
            backpressure,
            metrics,
        }
    }
    
    /// Get number of shards
    pub fn shard_count(&self) -> usize {
        self.shards.len()
    }
    
    /// Route a document to the appropriate shard
    pub fn route(&self, node: &MemoryNode) -> usize {
        match self.config.strategy {
            ShardStrategy::Hash => self.hash_route(&node.id),
            ShardStrategy::RoundRobin => self.round_robin_route(),
            ShardStrategy::Collection => {
                // Use collection from metadata if available, otherwise hash
                if let Some(collection) = node.metadata.get("_collection") {
                    if let Some(collection_str) = collection.as_str() {
                        return self.hash_route(collection_str);
                    }
                }
                self.hash_route(&node.id)
            }
        }
    }
    
    /// Route based on hash
    fn hash_route(&self, key: &str) -> usize {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.shards.len()
    }
    
    /// Route using round-robin
    fn round_robin_route(&self) -> usize {
        let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
        counter % self.shards.len()
    }
    
    /// Write a document (async)
    pub async fn write(&self, node: MemoryNode) -> Result<WriteResult, ShardRouterError> {
        // Check backpressure
        let fill_ratio = self.get_total_fill_ratio();
        self.backpressure.apply(fill_ratio).await
            .map_err(|e| ShardRouterError::Backpressure(e))?;
        
        // Route to shard
        let shard_id = self.route(&node);
        let shard = &self.shards[shard_id];
        
        // Enqueue in shard
        shard.write(node).await
            .map_err(|e| ShardRouterError::ShardError(shard_id, e.to_string()))
    }
    
    /// Write a document synchronously (waits for completion)
    pub async fn write_sync(&self, node: MemoryNode) -> Result<WriteResult, ShardRouterError> {
        // Check backpressure
        let fill_ratio = self.get_total_fill_ratio();
        self.backpressure.apply(fill_ratio).await
            .map_err(|e| ShardRouterError::Backpressure(e))?;
        
        // Route to shard
        let shard_id = self.route(&node);
        let shard = &self.shards[shard_id];
        
        // Enqueue and wait for completion
        shard.write_sync(node).await
            .map_err(|e| ShardRouterError::ShardError(shard_id, e.to_string()))
    }
    
    /// Write multiple documents (async)
    pub async fn write_batch(&self, nodes: Vec<MemoryNode>) -> Result<Vec<WriteResult>, ShardRouterError> {
        // Check backpressure
        let fill_ratio = self.get_total_fill_ratio();
        self.backpressure.apply(fill_ratio).await
            .map_err(|e| ShardRouterError::Backpressure(e))?;
        
        let mut results = Vec::with_capacity(nodes.len());
        
        for node in nodes {
            match self.write(node).await {
                Ok(result) => results.push(result),
                Err(e) => return Err(e),
            }
        }
        
        Ok(results)
    }
    
    /// Get total queue fill ratio across all shards
    pub fn get_total_fill_ratio(&self) -> f64 {
        let total_len: usize = self.shards.iter().map(|s| s.queue_len()).sum();
        let total_capacity: usize = self.shards.iter().map(|s| s.queue_capacity()).sum();
        
        if total_capacity == 0 {
            0.0
        } else {
            total_len as f64 / total_capacity as f64
        }
    }
    
    /// Get shard by ID
    pub fn get_shard(&self, id: usize) -> Option<Arc<WriteShard>> {
        self.shards.get(id).cloned()
    }
    
    /// Get all shards
    pub fn shards(&self) -> &[Arc<WriteShard>] {
        &self.shards
    }
    
    /// Flush all shards
    pub async fn flush(&self) -> Result<(), ShardRouterError> {
        for (i, shard) in self.shards.iter().enumerate() {
            shard.flush().await
                .map_err(|e| ShardRouterError::ShardError(i, e.to_string()))?;
        }
        Ok(())
    }
    
    /// Shutdown all shards
    pub async fn shutdown(&self) -> Result<(), ShardRouterError> {
        for (i, shard) in self.shards.iter().enumerate() {
            shard.shutdown().await
                .map_err(|e| ShardRouterError::ShardError(i, e.to_string()))?;
        }
        Ok(())
    }
}

/// Errors from shard router
#[derive(Debug, thiserror::Error)]
pub enum ShardRouterError {
    #[error("Backpressure: {0}")]
    Backpressure(#[from] BackpressureError),
    
    #[error("Shard {0} error: {1}")]
    ShardError(usize, String),
    
    #[error("No shards available")]
    NoShards,
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_hash_routing() {
        // Test that same ID always routes to same shard
        let config = ShardRouterConfig {
            shard_count: 4,
            strategy: ShardStrategy::Hash,
            auto_balance: false,
        };
        
        // Create minimal test - just verify hash consistency
        let mut hasher1 = DefaultHasher::new();
        "test-id-1".hash(&mut hasher1);
        let shard1 = (hasher1.finish() as usize) % 4;
        
        let mut hasher2 = DefaultHasher::new();
        "test-id-1".hash(&mut hasher2);
        let shard2 = (hasher2.finish() as usize) % 4;
        
        assert_eq!(shard1, shard2);
    }
}
