//! Node lifecycle operations (kill, revive, restart, shutdown).

use crate::error::{ChaosError, Result};
use crate::harness::ChaosNode;
use stemedb_cluster::membership::NodeId;
use tracing::info;

use super::types::TestCluster;

/// Extension trait for node lifecycle operations.
pub trait TestClusterLifecycleExt {
    /// Kills a node (simulates failure).
    fn kill_node(&mut self, idx: usize);

    /// Revives a node (simulates recovery).
    fn revive_node(&mut self, idx: usize);

    /// Creates a fresh node at the given index (simulates restart with empty state).
    ///
    /// The new node joins with no data but will sync from surviving nodes.
    fn restart_node(&mut self, idx: usize) -> Result<()>;

    /// Shuts down the cluster, cleaning up resources.
    fn shutdown(self) -> impl std::future::Future<Output = ()> + Send;
}

impl TestClusterLifecycleExt for TestCluster {
    fn kill_node(&mut self, idx: usize) {
        if idx < self.nodes.len() {
            self.nodes[idx].kill();
        }
    }

    fn revive_node(&mut self, idx: usize) {
        if idx < self.nodes.len() {
            self.nodes[idx].revive();
        }
    }

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

        let new_node = ChaosNode::new(idx, self.network.clone(), self.clock.clone())?;

        // Initialize sharding
        let node_ids: Vec<NodeId> = self.nodes.iter().map(|n| n.node_id()).collect();
        new_node.init_shards(&node_ids, self.num_shards, self.replication_factor);

        self.nodes[idx] = new_node;
        Ok(())
    }

    async fn shutdown(self) {
        info!("Shutting down test cluster");
        // Nodes are dropped, temp directories cleaned up
        drop(self);
    }
}
