//! Mock implementations for testing

use super::client::NodeClientTrait;
use super::registry::{NodeInfo, NodeRegistry};
use super::{NodeError, NodeResult};
use async_trait::async_trait;
use bytes::Bytes;
use dashmap::DashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

/// Mock storage node client for testing
///
/// Stores chunks in memory and allows simulation of failures.
pub struct MockNodeClient {
    node_id: String,
    endpoint: String,
    chunks: DashMap<String, Vec<u8>>,
    healthy: AtomicBool,
    // Counters for verification
    put_count: AtomicU64,
    get_count: AtomicU64,
    delete_count: AtomicU64,
    // Failure injection
    fail_puts: AtomicBool,
    fail_gets: AtomicBool,
    fail_deletes: AtomicBool,
}

impl MockNodeClient {
    /// Create a new mock node client
    pub fn new(node_id: impl Into<String>, endpoint: impl Into<String>) -> Self {
        Self {
            node_id: node_id.into(),
            endpoint: endpoint.into(),
            chunks: DashMap::new(),
            healthy: AtomicBool::new(true),
            put_count: AtomicU64::new(0),
            get_count: AtomicU64::new(0),
            delete_count: AtomicU64::new(0),
            fail_puts: AtomicBool::new(false),
            fail_gets: AtomicBool::new(false),
            fail_deletes: AtomicBool::new(false),
        }
    }

    /// Set the health status of this mock node
    pub fn set_healthy(&self, healthy: bool) {
        self.healthy.store(healthy, Ordering::SeqCst);
    }

    /// Enable/disable put failures
    pub fn set_fail_puts(&self, fail: bool) {
        self.fail_puts.store(fail, Ordering::SeqCst);
    }

    /// Enable/disable get failures
    pub fn set_fail_gets(&self, fail: bool) {
        self.fail_gets.store(fail, Ordering::SeqCst);
    }

    /// Enable/disable delete failures
    pub fn set_fail_deletes(&self, fail: bool) {
        self.fail_deletes.store(fail, Ordering::SeqCst);
    }

    /// Get the count of put operations
    pub fn put_count(&self) -> u64 {
        self.put_count.load(Ordering::SeqCst)
    }

    /// Get the count of get operations
    pub fn get_count(&self) -> u64 {
        self.get_count.load(Ordering::SeqCst)
    }

    /// Get the count of delete operations
    pub fn delete_count(&self) -> u64 {
        self.delete_count.load(Ordering::SeqCst)
    }

    /// Get all stored chunk IDs
    pub fn chunk_ids(&self) -> Vec<String> {
        self.chunks.iter().map(|r| r.key().clone()).collect()
    }

    /// Get the number of stored chunks
    pub fn chunk_count(&self) -> usize {
        self.chunks.len()
    }

    /// Clear all stored chunks
    pub fn clear(&self) {
        self.chunks.clear();
    }

    /// Reset all counters
    pub fn reset_counters(&self) {
        self.put_count.store(0, Ordering::SeqCst);
        self.get_count.store(0, Ordering::SeqCst);
        self.delete_count.store(0, Ordering::SeqCst);
    }
}

#[async_trait]
impl NodeClientTrait for MockNodeClient {
    fn node_id(&self) -> &str {
        &self.node_id
    }

    fn endpoint(&self) -> &str {
        &self.endpoint
    }

    async fn is_healthy(&self) -> bool {
        self.healthy.load(Ordering::SeqCst)
    }

    async fn put_chunk(
        &self,
        chunk_id: &str,
        _shard_index: u32,
        _is_parity: bool,
        data: Bytes,
    ) -> NodeResult<()> {
        self.put_count.fetch_add(1, Ordering::SeqCst);

        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        if self.fail_puts.load(Ordering::SeqCst) {
            return Err(NodeError::RpcFailed("Simulated put failure".into()));
        }

        self.chunks.insert(chunk_id.to_string(), data.to_vec());
        Ok(())
    }

    async fn get_chunk(
        &self,
        chunk_id: &str,
        _shard_index: u32,
        _is_parity: bool,
    ) -> NodeResult<Bytes> {
        self.get_count.fetch_add(1, Ordering::SeqCst);

        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        if self.fail_gets.load(Ordering::SeqCst) {
            return Err(NodeError::RpcFailed("Simulated get failure".into()));
        }

        self.chunks
            .get(chunk_id)
            .map(|r| Bytes::from(r.value().clone()))
            .ok_or_else(|| NodeError::NotFound(chunk_id.to_string()))
    }

    async fn delete_chunk(&self, chunk_id: &str) -> NodeResult<()> {
        self.delete_count.fetch_add(1, Ordering::SeqCst);

        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        if self.fail_deletes.load(Ordering::SeqCst) {
            return Err(NodeError::RpcFailed("Simulated delete failure".into()));
        }

        self.chunks.remove(chunk_id);
        Ok(())
    }

    async fn chunk_exists(&self, chunk_id: &str) -> NodeResult<bool> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        Ok(self.chunks.contains_key(chunk_id))
    }

    async fn chunk_size(&self, chunk_id: &str) -> NodeResult<Option<u64>> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        Ok(self.chunks.get(chunk_id).map(|r| r.value().len() as u64))
    }

    async fn ping(&self) -> NodeResult<Duration> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        Ok(Duration::from_micros(100)) // Simulated latency
    }
}

/// Mock node registry for testing
pub struct MockNodeRegistry {
    nodes: DashMap<String, Arc<MockNodeClient>>,
}

impl MockNodeRegistry {
    /// Create a new empty mock registry
    pub fn new() -> Self {
        Self {
            nodes: DashMap::new(),
        }
    }

    /// Add a mock node to the registry
    pub fn add_mock_node(&self, node: Arc<MockNodeClient>) {
        self.nodes.insert(node.node_id().to_string(), node);
    }

    /// Get a mock node by ID
    pub fn get_mock_node(&self, node_id: &str) -> Option<Arc<MockNodeClient>> {
        self.nodes.get(node_id).map(|r| r.value().clone())
    }

    /// Create a registry with N mock nodes
    pub fn with_nodes(count: usize) -> Self {
        let registry = Self::new();
        for i in 0..count {
            let node = Arc::new(MockNodeClient::new(
                format!("node-{}", i),
                format!("http://node-{}:9002", i),
            ));
            registry.add_mock_node(node);
        }
        registry
    }

    /// Get all mock nodes
    pub fn all_mock_nodes(&self) -> Vec<Arc<MockNodeClient>> {
        self.nodes.iter().map(|r| r.value().clone()).collect()
    }
}

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

#[async_trait]
impl NodeRegistry for MockNodeRegistry {
    async fn get_all_nodes(&self) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        Ok(self
            .nodes
            .iter()
            .map(|r| r.value().clone() as Arc<dyn NodeClientTrait>)
            .collect())
    }

    async fn get_healthy_nodes(&self) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        let mut healthy = Vec::new();
        for node_ref in self.nodes.iter() {
            let node = node_ref.value();
            if node.is_healthy().await {
                healthy.push(node.clone() as Arc<dyn NodeClientTrait>);
            }
        }
        Ok(healthy)
    }

    async fn register_node(&self, info: NodeInfo) -> NodeResult<()> {
        let node = Arc::new(MockNodeClient::new(&info.node_id, &info.endpoint));
        self.nodes.insert(info.node_id, node);
        Ok(())
    }

    async fn deregister_node(&self, node_id: &str) -> NodeResult<()> {
        self.nodes.remove(node_id);
        Ok(())
    }

    async fn update_health(&self, node_id: &str, healthy: bool) -> NodeResult<()> {
        if let Some(node) = self.nodes.get(node_id) {
            node.set_healthy(healthy);
        }
        Ok(())
    }

    async fn get_node(&self, node_id: &str) -> NodeResult<Option<Arc<dyn NodeClientTrait>>> {
        Ok(self
            .nodes
            .get(node_id)
            .map(|r| r.value().clone() as Arc<dyn NodeClientTrait>))
    }

    async fn node_count(&self) -> usize {
        self.nodes.len()
    }

    async fn healthy_node_count(&self) -> usize {
        let mut count = 0;
        for node_ref in self.nodes.iter() {
            if node_ref.value().is_healthy().await {
                count += 1;
            }
        }
        count
    }
}

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

    #[tokio::test]
    async fn test_mock_node_put_get() {
        let node = MockNodeClient::new("node-1", "http://localhost:9002");

        let chunk_id = "test-chunk-1";
        let data = Bytes::from(vec![1, 2, 3, 4, 5]);

        node.put_chunk(chunk_id, 0, false, data.clone())
            .await
            .unwrap();
        let retrieved = node.get_chunk(chunk_id, 0, false).await.unwrap();

        assert_eq!(retrieved, data.to_vec());
        assert_eq!(node.put_count(), 1);
        assert_eq!(node.get_count(), 1);
    }

    #[tokio::test]
    async fn test_mock_node_delete() {
        let node = MockNodeClient::new("node-1", "http://localhost:9002");

        let chunk_id = "test-chunk-1";
        let data = Bytes::from(vec![1, 2, 3]);

        node.put_chunk(chunk_id, 0, false, data).await.unwrap();
        assert!(node.chunk_exists(chunk_id).await.unwrap());

        node.delete_chunk(chunk_id).await.unwrap();
        assert!(!node.chunk_exists(chunk_id).await.unwrap());
    }

    #[tokio::test]
    async fn test_mock_node_health() {
        let node = MockNodeClient::new("node-1", "http://localhost:9002");

        assert!(node.is_healthy().await);

        node.set_healthy(false);
        assert!(!node.is_healthy().await);

        let result = node.put_chunk("chunk", 0, false, Bytes::new()).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_mock_node_failure_injection() {
        let node = MockNodeClient::new("node-1", "http://localhost:9002");

        // Normal operation
        node.put_chunk("chunk", 0, false, Bytes::from(vec![1]))
            .await
            .unwrap();

        // Enable failure injection
        node.set_fail_puts(true);
        let result = node.put_chunk("chunk2", 0, false, Bytes::from(vec![2])).await;
        assert!(result.is_err());

        // Disable failure injection
        node.set_fail_puts(false);
        node.put_chunk("chunk3", 0, false, Bytes::from(vec![3]))
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_mock_registry() {
        let registry = MockNodeRegistry::with_nodes(3);

        assert_eq!(registry.node_count().await, 3);
        assert_eq!(registry.healthy_node_count().await, 3);

        // Mark one node unhealthy
        registry.update_health("node-1", false).await.unwrap();
        assert_eq!(registry.healthy_node_count().await, 2);

        // Get healthy nodes
        let healthy = registry.get_healthy_nodes().await.unwrap();
        assert_eq!(healthy.len(), 2);
    }

    #[tokio::test]
    async fn test_mock_registry_register_deregister() {
        let registry = MockNodeRegistry::new();

        let info = NodeInfo::new("new-node", "http://new:9002");
        registry.register_node(info).await.unwrap();
        assert_eq!(registry.node_count().await, 1);

        registry.deregister_node("new-node").await.unwrap();
        assert_eq!(registry.node_count().await, 0);
    }
}
