//! Node access methods.

use crate::error::{ChaosError, Result};
use crate::fault_injection::{ClockController, NetworkController};
use crate::harness::ChaosNode;
use std::sync::Arc;

use super::types::TestCluster;

/// Extension trait for node access.
pub trait TestClusterAccessExt {
    /// Returns the number of nodes in the cluster.
    fn node_count(&self) -> usize;

    /// Gets a reference to a node by index.
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds.
    fn get_node(&self, idx: usize) -> &ChaosNode;

    /// Gets a mutable reference to a node by index.
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds.
    fn get_node_mut(&mut self, idx: usize) -> &mut ChaosNode;

    /// Gets a node by index, returning an error if not found.
    fn try_get_node(&self, idx: usize) -> Result<&ChaosNode>;

    /// Gets a mutable node by index, returning an error if not found.
    fn try_get_node_mut(&mut self, idx: usize) -> Result<&mut ChaosNode>;

    /// Returns the shared network controller.
    fn network(&self) -> &Arc<NetworkController>;

    /// Returns the shared clock controller.
    fn clock(&self) -> &Arc<ClockController>;

    /// Returns all node indices.
    fn node_indices(&self) -> Vec<usize>;

    /// Returns indices of alive nodes only.
    fn alive_node_indices(&self) -> Vec<usize>;
}

impl TestClusterAccessExt for TestCluster {
    fn node_count(&self) -> usize {
        self.nodes.len()
    }

    fn get_node(&self, idx: usize) -> &ChaosNode {
        &self.nodes[idx]
    }

    fn get_node_mut(&mut self, idx: usize) -> &mut ChaosNode {
        &mut self.nodes[idx]
    }

    fn try_get_node(&self, idx: usize) -> Result<&ChaosNode> {
        self.nodes.get(idx).ok_or(ChaosError::NodeNotFound(idx))
    }

    fn try_get_node_mut(&mut self, idx: usize) -> Result<&mut ChaosNode> {
        if idx >= self.nodes.len() {
            return Err(ChaosError::NodeNotFound(idx));
        }
        Ok(&mut self.nodes[idx])
    }

    fn network(&self) -> &Arc<NetworkController> {
        &self.network
    }

    fn clock(&self) -> &Arc<ClockController> {
        &self.clock
    }

    fn node_indices(&self) -> Vec<usize> {
        (0..self.nodes.len()).collect()
    }

    fn alive_node_indices(&self) -> Vec<usize> {
        self.nodes.iter().enumerate().filter(|(_, n)| n.is_alive()).map(|(i, _)| i).collect()
    }
}
