//! Node registry and management for CDN control plane

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

/// Node health status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum NodeStatus {
    Online,
    Offline,
    Maintenance,
    Overloaded,
}

impl Default for NodeStatus {
    fn default() -> Self {
        Self::Offline
    }
}

/// OSTP server node in the CDN network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node {
    pub node_id: Uuid,
    pub name: String,
    pub address: String,           // "ip:port"
    pub country_code: String,
    pub status: NodeStatus,
    pub cpu_load: f32,             // 0.0 - 1.0
    pub active_connections: u32,
    pub max_connections: u32,
    pub bytes_tx: u64,
    pub bytes_rx: u64,
    pub last_checkin: DateTime<Utc>,
    pub registered_at: DateTime<Utc>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub psk_hash: Option<String>,  // First 8 chars of PSK hash for identification
}

impl Node {
    pub fn new(name: impl Into<String>, address: impl Into<String>, country_code: impl Into<String>) -> Self {
        let now = Utc::now();
        Self {
            node_id: Uuid::new_v4(),
            name: name.into(),
            address: address.into(),
            country_code: country_code.into(),
            status: NodeStatus::Offline,
            cpu_load: 0.0,
            active_connections: 0,
            max_connections: 1000,
            bytes_tx: 0,
            bytes_rx: 0,
            last_checkin: now,
            registered_at: now,
            psk_hash: None,
        }
    }

    /// Check if node is healthy and can accept connections
    pub fn is_available(&self) -> bool {
        self.status == NodeStatus::Online && 
        self.active_connections < self.max_connections &&
        self.cpu_load < 0.9
    }

    /// Calculate node load score (lower is better)
    pub fn load_score(&self) -> f32 {
        let conn_ratio = self.active_connections as f32 / self.max_connections.max(1) as f32;
        (self.cpu_load + conn_ratio) / 2.0
    }

    /// Check if node is stale (no check-in for timeout period)
    pub fn is_stale(&self, timeout_secs: i64) -> bool {
        let elapsed = Utc::now().signed_duration_since(self.last_checkin);
        elapsed.num_seconds() > timeout_secs
    }
}

/// Node check-in request from OSTP server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeCheckin {
    pub node_id: Uuid,
    pub cpu_load: f32,
    pub active_connections: u32,
    pub bytes_tx: u64,
    pub bytes_rx: u64,
    #[serde(default)]
    pub sni_list: Vec<String>,     // Currently loaded SNIs
}

/// Node registry managing all OSTP servers
pub struct NodeRegistry {
    nodes: Arc<RwLock<HashMap<Uuid, Node>>>,
    checkin_timeout_secs: i64,
}

impl NodeRegistry {
    pub fn new(checkin_timeout_secs: i64) -> Self {
        Self {
            nodes: Arc::new(RwLock::new(HashMap::new())),
            checkin_timeout_secs,
        }
    }

    /// Register a new node
    pub async fn register(&self, node: Node) -> Uuid {
        let id = node.node_id;
        self.nodes.write().await.insert(id, node);
        tracing::info!("Registered node {}", id);
        id
    }

    /// Process node check-in (heartbeat + status update)
    pub async fn checkin(&self, checkin: NodeCheckin) -> Option<Node> {
        let mut nodes = self.nodes.write().await;
        if let Some(node) = nodes.get_mut(&checkin.node_id) {
            node.cpu_load = checkin.cpu_load;
            node.active_connections = checkin.active_connections;
            node.bytes_tx = checkin.bytes_tx;
            node.bytes_rx = checkin.bytes_rx;
            node.last_checkin = Utc::now();
            
            // Update status based on load
            node.status = if checkin.cpu_load > 0.95 {
                NodeStatus::Overloaded
            } else {
                NodeStatus::Online
            };
            
            Some(node.clone())
        } else {
            None
        }
    }

    /// Get node by ID
    pub async fn get(&self, node_id: &Uuid) -> Option<Node> {
        self.nodes.read().await.get(node_id).cloned()
    }

    /// Get all nodes
    pub async fn list(&self) -> Vec<Node> {
        self.nodes.read().await.values().cloned().collect()
    }

    /// Get online nodes only
    pub async fn list_online(&self) -> Vec<Node> {
        self.nodes.read().await
            .values()
            .filter(|n| n.status == NodeStatus::Online)
            .cloned()
            .collect()
    }

    /// Get best nodes for a specific country (sorted by load)
    pub async fn best_for_country(&self, country_code: &str, limit: usize) -> Vec<Node> {
        let mut nodes: Vec<Node> = self.nodes.read().await
            .values()
            .filter(|n| n.is_available() && n.country_code == country_code)
            .cloned()
            .collect();
        
        nodes.sort_by(|a, b| a.load_score().partial_cmp(&b.load_score()).unwrap());
        nodes.truncate(limit);
        nodes
    }

    /// Get best nodes globally (for any country)
    pub async fn best_global(&self, limit: usize) -> Vec<Node> {
        let mut nodes: Vec<Node> = self.nodes.read().await
            .values()
            .filter(|n| n.is_available())
            .cloned()
            .collect();
        
        nodes.sort_by(|a, b| a.load_score().partial_cmp(&b.load_score()).unwrap());
        nodes.truncate(limit);
        nodes
    }

    /// Remove node
    pub async fn remove(&self, node_id: &Uuid) -> Option<Node> {
        self.nodes.write().await.remove(node_id)
    }

    /// Mark stale nodes as offline
    pub async fn cleanup_stale(&self) -> Vec<Uuid> {
        let mut nodes = self.nodes.write().await;
        let timeout = self.checkin_timeout_secs;
        
        let stale: Vec<Uuid> = nodes
            .iter()
            .filter(|(_, n)| n.is_stale(timeout))
            .map(|(id, _)| *id)
            .collect();
        
        for id in &stale {
            if let Some(node) = nodes.get_mut(id) {
                node.status = NodeStatus::Offline;
            }
        }
        
        stale
    }

    /// Get aggregated network statistics
    pub async fn network_stats(&self) -> NetworkStats {
        let nodes = self.nodes.read().await;
        
        let total_nodes = nodes.len();
        let online_nodes = nodes.values().filter(|n| n.status == NodeStatus::Online).count();
        let total_connections: u32 = nodes.values().map(|n| n.active_connections).sum();
        let total_bytes_tx: u64 = nodes.values().map(|n| n.bytes_tx).sum();
        let total_bytes_rx: u64 = nodes.values().map(|n| n.bytes_rx).sum();
        let avg_load: f32 = if online_nodes > 0 {
            nodes.values()
                .filter(|n| n.status == NodeStatus::Online)
                .map(|n| n.load_score())
                .sum::<f32>() / online_nodes as f32
        } else {
            0.0
        };

        NetworkStats {
            total_nodes,
            online_nodes,
            total_connections,
            total_bytes_tx,
            total_bytes_rx,
            avg_load,
        }
    }
}

/// Aggregated network statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkStats {
    pub total_nodes: usize,
    pub online_nodes: usize,
    pub total_connections: u32,
    pub total_bytes_tx: u64,
    pub total_bytes_rx: u64,
    pub avg_load: f32,
}

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

    #[tokio::test]
    async fn test_node_registry() {
        let registry = NodeRegistry::new(30);
        
        let node = Node::new("test-node", "1.2.3.4:8443", "US");
        let node_id = registry.register(node).await;
        
        // Check-in with updated stats
        let checkin = NodeCheckin {
            node_id,
            cpu_load: 0.3,
            active_connections: 10,
            bytes_tx: 1024,
            bytes_rx: 2048,
            sni_list: vec![],
        };
        
        let updated = registry.checkin(checkin).await;
        assert!(updated.is_some());
        
        let node = updated.unwrap();
        assert_eq!(node.status, NodeStatus::Online);
        assert_eq!(node.active_connections, 10);
    }

    #[tokio::test]
    async fn test_best_nodes() {
        let registry = NodeRegistry::new(30);
        
        // Add nodes with different loads
        let mut node1 = Node::new("low-load", "1.1.1.1:8443", "US");
        node1.status = NodeStatus::Online;
        node1.cpu_load = 0.2;
        
        let mut node2 = Node::new("high-load", "2.2.2.2:8443", "US");
        node2.status = NodeStatus::Online;
        node2.cpu_load = 0.8;
        
        registry.register(node1).await;
        registry.register(node2).await;
        
        let best = registry.best_for_country("US", 1).await;
        assert_eq!(best.len(), 1);
        assert_eq!(best[0].name, "low-load");
    }
}
