//! Custom broadcast handler for actual state propagation

use chainfire_types::NodeId;
use dashmap::DashMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::SystemTime;
use tracing::debug;

/// Actual state data broadcast via gossip
///
/// This is the "Actual State" mentioned in the design - things like
/// CPU usage, memory, running tasks, etc. that are eventually consistent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActualState {
    /// Node ID this state is from
    pub node_id: NodeId,
    /// Timestamp when this state was generated
    pub timestamp: u64,
    /// CPU usage percentage (0-100)
    pub cpu_usage: f32,
    /// Memory usage percentage (0-100)
    pub memory_usage: f32,
    /// Disk usage percentage (0-100)
    pub disk_usage: f32,
    /// Custom status fields (e.g., "vm-a" -> "running")
    pub status: HashMap<String, String>,
}

impl ActualState {
    /// Create a new actual state
    pub fn new(node_id: NodeId) -> Self {
        let timestamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();

        Self {
            node_id,
            timestamp,
            cpu_usage: 0.0,
            memory_usage: 0.0,
            disk_usage: 0.0,
            status: HashMap::new(),
        }
    }

    /// Set CPU usage
    pub fn with_cpu(mut self, usage: f32) -> Self {
        self.cpu_usage = usage;
        self
    }

    /// Set memory usage
    pub fn with_memory(mut self, usage: f32) -> Self {
        self.memory_usage = usage;
        self
    }

    /// Set disk usage
    pub fn with_disk(mut self, usage: f32) -> Self {
        self.disk_usage = usage;
        self
    }

    /// Add a status entry
    pub fn with_status(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.status.insert(key.into(), value.into());
        self
    }

    /// Update timestamp to now
    pub fn touch(&mut self) {
        self.timestamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();
    }
}

/// Broadcast handler for actual state propagation
pub struct ActualStateBroadcast {
    /// Current node's actual state
    local_state: RwLock<Option<ActualState>>,
    /// Collected states from other nodes
    cluster_state: DashMap<NodeId, ActualState>,
}

impl ActualStateBroadcast {
    /// Create a new broadcast handler
    pub fn new() -> Self {
        Self {
            local_state: RwLock::new(None),
            cluster_state: DashMap::new(),
        }
    }

    /// Set the local node's state
    pub fn set_local_state(&self, state: ActualState) {
        *self.local_state.write() = Some(state);
    }

    /// Get the local node's state
    pub fn local_state(&self) -> Option<ActualState> {
        self.local_state.read().clone()
    }

    /// Get a node's state
    pub fn get_state(&self, node_id: NodeId) -> Option<ActualState> {
        self.cluster_state.get(&node_id).map(|r| r.clone())
    }

    /// Get all cluster states
    pub fn all_states(&self) -> Vec<ActualState> {
        self.cluster_state.iter().map(|r| r.clone()).collect()
    }

    /// Remove a node's state (on member down)
    pub fn remove_state(&self, node_id: NodeId) {
        self.cluster_state.remove(&node_id);
    }
}

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

impl ActualStateBroadcast {
    /// Receive and process state from another node
    /// Returns true if the state was newer and accepted
    pub fn receive_state(&self, state: ActualState) -> bool {
        let node_id = state.node_id;

        // Check if we should update
        if let Some(existing) = self.cluster_state.get(&node_id) {
            if existing.timestamp >= state.timestamp {
                return false; // Stale data
            }
        }

        debug!(node_id, timestamp = state.timestamp, "Received actual state");
        self.cluster_state.insert(node_id, state);
        true
    }
}

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

    #[test]
    fn test_actual_state_creation() {
        let state = ActualState::new(1)
            .with_cpu(50.0)
            .with_memory(75.0)
            .with_status("vm-1", "running");

        assert_eq!(state.node_id, 1);
        assert_eq!(state.cpu_usage, 50.0);
        assert_eq!(state.memory_usage, 75.0);
        assert_eq!(state.status.get("vm-1"), Some(&"running".to_string()));
    }

    #[test]
    fn test_receive_state() {
        let handler = ActualStateBroadcast::new();

        // Receive first state
        let state1 = ActualState::new(1).with_cpu(50.0);
        let result = handler.receive_state(state1.clone());
        assert!(result); // Should accept

        // Receive newer state
        let mut state2 = ActualState::new(1).with_cpu(60.0);
        state2.timestamp = state1.timestamp + 1;
        let result = handler.receive_state(state2.clone());
        assert!(result); // Should accept

        // Receive older state
        let mut state3 = ActualState::new(1).with_cpu(40.0);
        state3.timestamp = state1.timestamp - 1;
        let result = handler.receive_state(state3);
        assert!(!result); // Should reject stale data

        // Verify final state
        let stored = handler.get_state(1).unwrap();
        assert_eq!(stored.cpu_usage, 60.0);
    }

    #[test]
    fn test_cluster_state_collection() {
        let handler = ActualStateBroadcast::new();

        handler
            .cluster_state
            .insert(1, ActualState::new(1).with_cpu(50.0));
        handler
            .cluster_state
            .insert(2, ActualState::new(2).with_cpu(60.0));

        let states = handler.all_states();
        assert_eq!(states.len(), 2);

        handler.remove_state(1);
        assert_eq!(handler.all_states().len(), 1);
    }
}
