//! Data placement strategies for distributed storage
//!
//! This module provides strategies for selecting which nodes should store
//! which data, supporting consistent hashing, random placement, and
//! zone-aware placement.

use crate::node::{NodeClientTrait, NodeResult};
use async_trait::async_trait;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::Arc;

/// Trait for node selection strategies
#[async_trait]
pub trait NodeSelector: Send + Sync {
    /// Select N nodes for storing data
    ///
    /// Returns a vector of selected nodes, in order of preference.
    async fn select_nodes(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>>;

    /// Select N nodes for storing data associated with a stable key.
    ///
    /// Implementations can override this to keep write placement aligned with
    /// later keyed reads.
    async fn select_nodes_for_key(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
        _key: &str,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        self.select_nodes(available_nodes, count).await
    }

    /// Select a single node for reading data
    ///
    /// The key is used to deterministically select the same node
    /// for the same data (for cache efficiency).
    async fn select_for_read(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        key: &str,
    ) -> NodeResult<Arc<dyn NodeClientTrait>>;
}

/// Consistent hash-based node selector
///
/// Uses consistent hashing to select nodes, ensuring minimal data movement
/// when nodes are added or removed.
pub struct ConsistentHashSelector {
    /// Number of virtual nodes per physical node
    virtual_nodes: usize,
}

impl ConsistentHashSelector {
    /// Create a new consistent hash selector with default settings
    pub fn new() -> Self {
        Self { virtual_nodes: 100 }
    }

    /// Create a new consistent hash selector with custom virtual node count
    pub fn with_virtual_nodes(virtual_nodes: usize) -> Self {
        Self { virtual_nodes }
    }

    /// Hash a string to a u64
    fn hash_key(key: &str) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Get the hash ring position for a node and virtual node index
    fn node_position(&self, node_id: &str, vnode_index: usize) -> u64 {
        Self::hash_key(&format!("{}:{}", node_id, vnode_index))
    }
}

impl Default for ConsistentHashSelector {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl NodeSelector for ConsistentHashSelector {
    async fn select_nodes(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        if available_nodes.is_empty() {
            return Ok(vec![]);
        }

        let count = count.min(available_nodes.len());

        // Build the hash ring with virtual nodes
        let mut ring: Vec<(u64, usize)> = Vec::new();
        for (node_idx, node) in available_nodes.iter().enumerate() {
            for vnode_idx in 0..self.virtual_nodes {
                let pos = self.node_position(node.node_id(), vnode_idx);
                ring.push((pos, node_idx));
            }
        }
        ring.sort_by_key(|(pos, _)| *pos);

        self.select_nodes_from_ring(available_nodes, &ring, count, None)
    }

    async fn select_nodes_for_key(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
        key: &str,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        if available_nodes.is_empty() {
            return Ok(vec![]);
        }

        let count = count.min(available_nodes.len());
        let mut ring: Vec<(u64, usize)> = Vec::new();
        for (node_idx, node) in available_nodes.iter().enumerate() {
            for vnode_idx in 0..self.virtual_nodes {
                let pos = self.node_position(node.node_id(), vnode_idx);
                ring.push((pos, node_idx));
            }
        }
        ring.sort_by_key(|(pos, _)| *pos);

        self.select_nodes_from_ring(available_nodes, &ring, count, Some(key))
    }

    async fn select_for_read(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        key: &str,
    ) -> NodeResult<Arc<dyn NodeClientTrait>> {
        if available_nodes.is_empty() {
            return Err(crate::node::NodeError::NotEnoughNodes {
                needed: 1,
                available: 0,
            });
        }

        // Build the hash ring
        let mut ring: Vec<(u64, usize)> = Vec::new();
        for (node_idx, node) in available_nodes.iter().enumerate() {
            for vnode_idx in 0..self.virtual_nodes {
                let pos = self.node_position(node.node_id(), vnode_idx);
                ring.push((pos, node_idx));
            }
        }
        ring.sort_by_key(|(pos, _)| *pos);

        // Find the first node after the key's position
        let key_pos = Self::hash_key(key);
        let idx = ring
            .binary_search_by_key(&key_pos, |(pos, _)| *pos)
            .unwrap_or_else(|i| i % ring.len());

        let node_idx = ring[idx].1;
        Ok(available_nodes[node_idx].clone())
    }
}

impl ConsistentHashSelector {
    fn select_nodes_from_ring(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        ring: &[(u64, usize)],
        count: usize,
        key: Option<&str>,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        let mut selected_indices = Vec::with_capacity(count);
        let mut seen = std::collections::HashSet::new();

        let start_pos = key
            .map(Self::hash_key)
            .unwrap_or_else(|| Self::hash_key(&format!("{:?}", std::time::Instant::now())));
        let start_idx = ring
            .binary_search_by_key(&start_pos, |(pos, _)| *pos)
            .unwrap_or_else(|i| i % ring.len());

        for i in 0..ring.len() {
            let idx = (start_idx + i) % ring.len();
            let node_idx = ring[idx].1;

            if seen.insert(node_idx) {
                selected_indices.push(node_idx);
                if selected_indices.len() >= count {
                    break;
                }
            }
        }

        Ok(selected_indices
            .into_iter()
            .map(|idx| available_nodes[idx].clone())
            .collect())
    }
}

/// Random node selector
///
/// Randomly selects nodes for placement. Simple but doesn't provide
/// consistent placement across operations.
pub struct RandomSelector;

impl RandomSelector {
    pub fn new() -> Self {
        Self
    }
}

impl Default for RandomSelector {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl NodeSelector for RandomSelector {
    async fn select_nodes(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        if available_nodes.is_empty() {
            return Ok(vec![]);
        }

        let count = count.min(available_nodes.len());

        // Shuffle using Fisher-Yates with simple random
        let mut indices: Vec<usize> = (0..available_nodes.len()).collect();
        let seed = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos() as u64;

        let mut rng = seed;
        for i in (1..indices.len()).rev() {
            // Simple LCG random
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
            let j = (rng as usize) % (i + 1);
            indices.swap(i, j);
        }

        Ok(indices
            .into_iter()
            .take(count)
            .map(|idx| available_nodes[idx].clone())
            .collect())
    }

    async fn select_for_read(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        _key: &str,
    ) -> NodeResult<Arc<dyn NodeClientTrait>> {
        if available_nodes.is_empty() {
            return Err(crate::node::NodeError::NotEnoughNodes {
                needed: 1,
                available: 0,
            });
        }

        // Random selection
        let seed = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos() as u64;
        let idx = (seed as usize) % available_nodes.len();

        Ok(available_nodes[idx].clone())
    }
}

/// Round-robin node selector
///
/// Selects nodes in round-robin order. Good for load distribution.
pub struct RoundRobinSelector {
    counter: std::sync::atomic::AtomicUsize,
}

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

impl Default for RoundRobinSelector {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl NodeSelector for RoundRobinSelector {
    async fn select_nodes(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        count: usize,
    ) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        if available_nodes.is_empty() {
            return Ok(vec![]);
        }

        let count = count.min(available_nodes.len());
        let start = self
            .counter
            .fetch_add(count, std::sync::atomic::Ordering::SeqCst);

        Ok((0..count)
            .map(|i| {
                let idx = (start + i) % available_nodes.len();
                available_nodes[idx].clone()
            })
            .collect())
    }

    async fn select_for_read(
        &self,
        available_nodes: &[Arc<dyn NodeClientTrait>],
        _key: &str,
    ) -> NodeResult<Arc<dyn NodeClientTrait>> {
        if available_nodes.is_empty() {
            return Err(crate::node::NodeError::NotEnoughNodes {
                needed: 1,
                available: 0,
            });
        }

        let idx = self
            .counter
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
            % available_nodes.len();
        Ok(available_nodes[idx].clone())
    }
}

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

    fn create_mock_nodes(count: usize) -> Vec<Arc<dyn NodeClientTrait>> {
        (0..count)
            .map(|i| {
                Arc::new(MockNodeClient::new(
                    format!("node-{}", i),
                    format!("http://node-{}:9002", i),
                )) as Arc<dyn NodeClientTrait>
            })
            .collect()
    }

    #[tokio::test]
    async fn test_consistent_hash_deterministic_read() {
        let selector = ConsistentHashSelector::new();
        let nodes = create_mock_nodes(5);

        let key = "test-object-123";

        // Same key should always select the same node
        let node1 = selector.select_for_read(&nodes, key).await.unwrap();
        let node2 = selector.select_for_read(&nodes, key).await.unwrap();

        assert_eq!(node1.node_id(), node2.node_id());
    }

    #[tokio::test]
    async fn test_consistent_hash_select_nodes() {
        let selector = ConsistentHashSelector::new();
        let nodes = create_mock_nodes(6);

        let selected = selector.select_nodes(&nodes, 3).await.unwrap();

        assert_eq!(selected.len(), 3);

        // All selected nodes should be unique
        let ids: std::collections::HashSet<_> = selected.iter().map(|n| n.node_id()).collect();
        assert_eq!(ids.len(), 3);
    }

    #[tokio::test]
    async fn test_consistent_hash_select_nodes_for_key_is_deterministic() {
        let selector = ConsistentHashSelector::new();
        let nodes = create_mock_nodes(6);

        let first = selector
            .select_nodes_for_key(&nodes, 3, "stable-object-key")
            .await
            .unwrap();
        let second = selector
            .select_nodes_for_key(&nodes, 3, "stable-object-key")
            .await
            .unwrap();

        let first_ids: Vec<_> = first.iter().map(|n| n.node_id().to_string()).collect();
        let second_ids: Vec<_> = second.iter().map(|n| n.node_id().to_string()).collect();

        assert_eq!(first_ids, second_ids);
    }

    #[tokio::test]
    async fn test_consistent_hash_select_more_than_available() {
        let selector = ConsistentHashSelector::new();
        let nodes = create_mock_nodes(3);

        // Request more nodes than available
        let selected = selector.select_nodes(&nodes, 10).await.unwrap();

        // Should return all available nodes
        assert_eq!(selected.len(), 3);
    }

    #[tokio::test]
    async fn test_random_selector() {
        let selector = RandomSelector::new();
        let nodes = create_mock_nodes(5);

        let selected = selector.select_nodes(&nodes, 3).await.unwrap();
        assert_eq!(selected.len(), 3);
    }

    #[tokio::test]
    async fn test_round_robin_selector() {
        let selector = RoundRobinSelector::new();
        let nodes = create_mock_nodes(3);

        // First selection
        let node1 = selector.select_for_read(&nodes, "key1").await.unwrap();
        // Second selection should be different
        let node2 = selector.select_for_read(&nodes, "key2").await.unwrap();
        // Third selection
        let node3 = selector.select_for_read(&nodes, "key3").await.unwrap();
        // Fourth should wrap around to first
        let node4 = selector.select_for_read(&nodes, "key4").await.unwrap();

        // Verify round-robin behavior
        let ids: Vec<_> = [&node1, &node2, &node3, &node4]
            .iter()
            .map(|n| n.node_id())
            .collect();
        assert_eq!(ids[0], ids[3]); // Wrapped around
    }

    #[tokio::test]
    async fn test_empty_nodes() {
        let selector = ConsistentHashSelector::new();
        let nodes: Vec<Arc<dyn NodeClientTrait>> = vec![];

        let selected = selector.select_nodes(&nodes, 3).await.unwrap();
        assert!(selected.is_empty());

        let result = selector.select_for_read(&nodes, "key").await;
        assert!(result.is_err());
    }
}
