//! Node registry for discovering and tracking storage nodes

use super::client::{NodeClient, NodeClientTrait};
use super::NodeResult;
use async_trait::async_trait;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;

/// Information about a storage node
#[derive(Debug, Clone)]
pub struct NodeInfo {
    /// Unique node identifier
    pub node_id: String,
    /// gRPC endpoint (host:port)
    pub endpoint: String,
    /// Zone/rack identifier for placement
    pub zone: String,
    /// Region identifier
    pub region: String,
    /// Total storage capacity in bytes
    pub capacity_bytes: u64,
    /// Currently used storage in bytes
    pub used_bytes: u64,
    /// Whether the node is healthy
    pub healthy: bool,
}

impl NodeInfo {
    /// Create a new node info with minimal information
    pub fn new(node_id: impl Into<String>, endpoint: impl Into<String>) -> Self {
        Self {
            node_id: node_id.into(),
            endpoint: endpoint.into(),
            zone: String::new(),
            region: String::new(),
            capacity_bytes: 0,
            used_bytes: 0,
            healthy: true,
        }
    }

    /// Builder method to set the zone
    pub fn with_zone(mut self, zone: impl Into<String>) -> Self {
        self.zone = zone.into();
        self
    }

    /// Builder method to set the region
    pub fn with_region(mut self, region: impl Into<String>) -> Self {
        self.region = region.into();
        self
    }

    /// Builder method to set capacity
    pub fn with_capacity(mut self, capacity_bytes: u64) -> Self {
        self.capacity_bytes = capacity_bytes;
        self
    }

    /// Calculate available storage in bytes
    pub fn available_bytes(&self) -> u64 {
        self.capacity_bytes.saturating_sub(self.used_bytes)
    }

    /// Calculate usage percentage
    pub fn usage_percent(&self) -> f64 {
        if self.capacity_bytes == 0 {
            0.0
        } else {
            (self.used_bytes as f64 / self.capacity_bytes as f64) * 100.0
        }
    }
}

/// Trait for node registry implementations
#[async_trait]
pub trait NodeRegistry: Send + Sync {
    /// Get all registered nodes
    async fn get_all_nodes(&self) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>>;

    /// Get only healthy nodes
    async fn get_healthy_nodes(&self) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>>;

    /// Register a new node
    async fn register_node(&self, info: NodeInfo) -> NodeResult<()>;

    /// Remove a node from the registry
    async fn deregister_node(&self, node_id: &str) -> NodeResult<()>;

    /// Update node health status
    async fn update_health(&self, node_id: &str, healthy: bool) -> NodeResult<()>;

    /// Get a specific node by ID
    async fn get_node(&self, node_id: &str) -> NodeResult<Option<Arc<dyn NodeClientTrait>>>;

    /// Get the number of registered nodes
    async fn node_count(&self) -> usize;

    /// Get the number of healthy nodes
    async fn healthy_node_count(&self) -> usize;
}

/// Static node registry that uses a fixed list of endpoints
///
/// Nodes are configured at startup and don't change dynamically.
pub struct StaticNodeRegistry {
    nodes: RwLock<Vec<Arc<dyn NodeClientTrait>>>,
    node_info: RwLock<Vec<NodeInfo>>,
    connection_timeout: Duration,
    request_timeout: Duration,
}

impl StaticNodeRegistry {
    /// Create a new static node registry with the given endpoints
    pub async fn new(endpoints: &[String]) -> NodeResult<Self> {
        Self::new_with_timeouts(
            endpoints,
            Duration::from_secs(5),
            Duration::from_secs(300),
        )
        .await
    }

    /// Create a new static node registry with explicit timeout settings.
    pub async fn new_with_timeouts(
        endpoints: &[String],
        connection_timeout: Duration,
        request_timeout: Duration,
    ) -> NodeResult<Self> {
        let mut nodes: Vec<Arc<dyn NodeClientTrait>> = Vec::new();
        let mut node_info = Vec::new();

        for (i, endpoint) in endpoints.iter().enumerate() {
            let node_id = format!("node-{}", i);
            let client = NodeClient::connect_with_id_and_timeouts(
                &node_id,
                endpoint,
                connection_timeout,
                request_timeout,
            )
            .await?;
            let info = NodeInfo::new(&node_id, endpoint);

            nodes.push(Arc::new(client));
            node_info.push(info);
        }

        Ok(Self {
            nodes: RwLock::new(nodes),
            node_info: RwLock::new(node_info),
            connection_timeout,
            request_timeout,
        })
    }

    /// Create an empty registry
    pub fn empty() -> Self {
        Self {
            nodes: RwLock::new(Vec::new()),
            node_info: RwLock::new(Vec::new()),
            connection_timeout: Duration::from_secs(5),
            request_timeout: Duration::from_secs(300),
        }
    }

    /// Get node info for all nodes
    pub async fn get_node_info(&self) -> Vec<NodeInfo> {
        self.node_info.read().await.clone()
    }
}

#[async_trait]
impl NodeRegistry for StaticNodeRegistry {
    async fn get_all_nodes(&self) -> NodeResult<Vec<Arc<dyn NodeClientTrait>>> {
        Ok(self.nodes.read().await.clone())
    }

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

    async fn register_node(&self, info: NodeInfo) -> NodeResult<()> {
        let client = NodeClient::connect_with_id_and_timeouts(
            &info.node_id,
            &info.endpoint,
            self.connection_timeout,
            self.request_timeout,
        )
        .await?;

        self.nodes.write().await.push(Arc::new(client));
        self.node_info.write().await.push(info);

        Ok(())
    }

    async fn deregister_node(&self, node_id: &str) -> NodeResult<()> {
        self.nodes
            .write()
            .await
            .retain(|n| n.node_id() != node_id);
        self.node_info
            .write()
            .await
            .retain(|n| n.node_id != node_id);
        Ok(())
    }

    async fn update_health(&self, node_id: &str, healthy: bool) -> NodeResult<()> {
        let mut info = self.node_info.write().await;
        for node_info in info.iter_mut() {
            if node_info.node_id == node_id {
                node_info.healthy = healthy;
                break;
            }
        }

        // Note: For static registry, we don't actually update the client health
        // as the client manages its own health state. This is mainly for tracking.
        Ok(())
    }

    async fn get_node(&self, node_id: &str) -> NodeResult<Option<Arc<dyn NodeClientTrait>>> {
        let nodes = self.nodes.read().await;
        Ok(nodes.iter().find(|n| n.node_id() == node_id).cloned())
    }

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

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

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

    #[test]
    fn test_node_info_creation() {
        let info = NodeInfo::new("node-1", "http://localhost:9002")
            .with_zone("zone-a")
            .with_region("us-east-1")
            .with_capacity(1024 * 1024 * 1024);

        assert_eq!(info.node_id, "node-1");
        assert_eq!(info.endpoint, "http://localhost:9002");
        assert_eq!(info.zone, "zone-a");
        assert_eq!(info.region, "us-east-1");
        assert_eq!(info.capacity_bytes, 1024 * 1024 * 1024);
    }

    #[test]
    fn test_node_info_usage() {
        let mut info = NodeInfo::new("node-1", "http://localhost:9002").with_capacity(1000);

        info.used_bytes = 250;
        assert_eq!(info.available_bytes(), 750);
        assert!((info.usage_percent() - 25.0).abs() < 0.01);
    }

    #[tokio::test]
    async fn test_static_registry_creation() {
        let endpoints = vec![
            "http://node1:9002".to_string(),
            "http://node2:9002".to_string(),
        ];

        let registry = StaticNodeRegistry::new(&endpoints).await.unwrap();
        assert_eq!(registry.node_count().await, 2);
    }

    #[tokio::test]
    async fn test_static_registry_register_deregister() {
        let registry = StaticNodeRegistry::empty();
        assert_eq!(registry.node_count().await, 0);

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

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

    #[tokio::test]
    async fn test_static_registry_get_node() {
        let endpoints = vec!["http://node1:9002".to_string()];
        let registry = StaticNodeRegistry::new(&endpoints).await.unwrap();

        let node = registry.get_node("node-0").await.unwrap();
        assert!(node.is_some());

        let missing = registry.get_node("nonexistent").await.unwrap();
        assert!(missing.is_none());
    }
}
