//! Cluster creation methods.

use crate::error::{ChaosError, Result};
use crate::fault_injection::{ClockController, NetworkController};
use crate::harness::ChaosNode;
use std::sync::Arc;
use stemedb_cluster::membership::NodeId;
use tracing::{info, instrument};

use super::types::TestCluster;

/// Extension trait for cluster creation.
pub trait TestClusterCreationExt: Sized {
    /// Creates a new test cluster with the specified number of nodes.
    ///
    /// All nodes are initialized with sharding configured.
    /// Default: 4 shards, replication factor of 2.
    ///
    /// # Arguments
    ///
    /// * `node_count` - Number of nodes in the cluster (must be >= 1)
    fn spawn(node_count: usize) -> impl std::future::Future<Output = Result<Self>> + Send;

    /// Creates a new test cluster with custom shard configuration.
    ///
    /// # Arguments
    ///
    /// * `node_count` - Number of nodes in the cluster
    /// * `num_shards` - Number of shards
    /// * `replication_factor` - Number of replicas per shard
    fn spawn_with_config(
        node_count: usize,
        num_shards: u32,
        replication_factor: u32,
    ) -> impl std::future::Future<Output = Result<Self>> + Send;
}

impl TestClusterCreationExt for TestCluster {
    async fn spawn(node_count: usize) -> Result<Self> {
        Self::spawn_with_config(node_count, 4, 2).await
    }

    #[instrument(skip_all, fields(nodes = node_count, shards = num_shards, rf = replication_factor))]
    async fn spawn_with_config(
        node_count: usize,
        num_shards: u32,
        replication_factor: u32,
    ) -> Result<Self> {
        if node_count == 0 {
            return Err(ChaosError::InvalidConfig("Node count must be >= 1".to_string()));
        }

        info!("Spawning test cluster");

        let network = Arc::new(NetworkController::new());
        let clock = Arc::new(ClockController::new());

        let mut nodes = Vec::with_capacity(node_count);

        for i in 0..node_count {
            let node = ChaosNode::new(i, network.clone(), clock.clone())?;
            nodes.push(node);
        }

        // Collect node IDs for shard initialization
        let node_ids: Vec<NodeId> = nodes.iter().map(|n| n.node_id()).collect();

        // Initialize sharding on all nodes
        for node in &nodes {
            node.init_shards(&node_ids, num_shards, replication_factor);
        }

        Ok(Self { nodes, network, clock, num_shards, replication_factor })
    }
}
