//! DistributedLLM: Trillion-Parameter Inference across Thunderbolt 5 Mac Cluster
//!
//! This demo showcases distributed large language model inference using
//! tensor parallelism and pipeline parallelism across a cluster of Mac
//! systems connected via Thunderbolt 5.
//!
//! ## What is real vs. simulated
//!
//! - **Real**: per-layer projection weight allocation (`ModelShard::load`),
//!   Q/K/V tensor construction, matmul + softmax attention compute
//!   (`DistributedAttention::forward`, via `rtx_tensor::Tensor::
//!   scaled_dot_product_attention`), and all wall-clock timings measured
//!   with `std::time::Instant`. Weight/KV memory sizes are computed from
//!   actual tensor element counts and dtype byte widths, not guessed
//!   constants.
//! - **Simulated**: the network transport layer (`DistributedCommunicator`
//!   models cluster bandwidth/latency but performs no actual inter-process
//!   communication), the tokenizer (word-based, not a real subword
//!   vocabulary), and end-of-sequence detection (a fixed probability, not a
//!   learned distribution). Flash-attention memory savings remain a
//!   documented closed-form estimate rather than a real tiled kernel.
//! - Model sharding across multiple nodes
//! - Distributed attention with sharded KV cache
//! - Pipeline parallelism with micro-batching
//! - High-bandwidth inter-node communication

pub mod attention;
pub mod parallel;
pub mod sample_data;

use distllm_shared::{
    ClusterConfig, FinishReason, GenerationConfig, InferenceRequest, InferenceResult,
    KVCacheConfig, LayerAssignment, ModelConfig, ParallelismConfig, PerformanceMetrics,
};

use attention::{DistributedAttention, KVCache};
use parallel::{LayerPartitioner, PipelineParallel, TensorParallel};
use rtx_tensor::{Device, Tensor};

// ============================================================================
// Model Shard
// ============================================================================

/// Real per-layer projection weights for one transformer layer.
///
/// Shapes are sized (and, under tensor parallelism, sharded) from the
/// model's `ModelConfig` the same way a real transformer layer would be:
/// Q/K/V/O attention projections plus a gated MLP (SwiGLU-style
/// gate/up/down projections).
#[derive(Debug)]
pub struct LayerWeights {
    /// Query projection: `[hidden_dim, sharded_num_heads * head_dim]`.
    pub q_proj: Tensor,
    /// Key projection: `[hidden_dim, sharded_num_kv_heads * head_dim]`.
    pub k_proj: Tensor,
    /// Value projection: `[hidden_dim, sharded_num_kv_heads * head_dim]`.
    pub v_proj: Tensor,
    /// Output projection: `[sharded_num_heads * head_dim, hidden_dim]`.
    pub o_proj: Tensor,
    /// MLP gate projection: `[hidden_dim, sharded_intermediate_dim]`.
    pub gate_proj: Tensor,
    /// MLP up projection: `[hidden_dim, sharded_intermediate_dim]`.
    pub up_proj: Tensor,
    /// MLP down projection: `[sharded_intermediate_dim, hidden_dim]`.
    pub down_proj: Tensor,
}

impl LayerWeights {
    /// Allocate real (randomly initialized) weight tensors for one layer,
    /// sized according to `model_config` and sharded `tensor_world`-ways
    /// for tensor parallelism (column/row split approximation).
    fn allocate(
        model_config: &ModelConfig,
        tensor_world: usize,
        device: &Device,
    ) -> Result<Self, String> {
        let hidden_dim = model_config.hidden_dim;
        let sharded_heads = (model_config.num_heads / tensor_world.max(1)).max(1);
        let sharded_kv_heads = (model_config.num_kv_heads / tensor_world.max(1)).max(1);
        let sharded_intermediate = (model_config.intermediate_dim / tensor_world.max(1)).max(1);
        let head_dim = model_config.head_dim;

        let q_dim = sharded_heads * head_dim;
        let kv_dim = sharded_kv_heads * head_dim;

        Ok(Self {
            q_proj: Tensor::randn(&[hidden_dim, q_dim], device).map_err(|e| e.to_string())?,
            k_proj: Tensor::randn(&[hidden_dim, kv_dim], device).map_err(|e| e.to_string())?,
            v_proj: Tensor::randn(&[hidden_dim, kv_dim], device).map_err(|e| e.to_string())?,
            o_proj: Tensor::randn(&[q_dim, hidden_dim], device).map_err(|e| e.to_string())?,
            gate_proj: Tensor::randn(&[hidden_dim, sharded_intermediate], device)
                .map_err(|e| e.to_string())?,
            up_proj: Tensor::randn(&[hidden_dim, sharded_intermediate], device)
                .map_err(|e| e.to_string())?,
            down_proj: Tensor::randn(&[sharded_intermediate, hidden_dim], device)
                .map_err(|e| e.to_string())?,
        })
    }

    /// Total bytes occupied by this layer's weights, computed from actual
    /// tensor element counts and the tensor's dtype byte width (real
    /// allocation size, not an estimate).
    fn size_bytes(&self) -> usize {
        let dtype_bytes = self.q_proj.dtype().size_bytes();
        [
            &self.q_proj,
            &self.k_proj,
            &self.v_proj,
            &self.o_proj,
            &self.gate_proj,
            &self.up_proj,
            &self.down_proj,
        ]
        .iter()
        .map(|t| t.numel() * dtype_bytes)
        .sum()
    }
}

/// A shard of the model assigned to this node.
#[derive(Debug)]
pub struct ModelShard {
    /// Shard ID.
    pub id: usize,
    /// Layer range (start, end exclusive).
    pub layer_range: (usize, usize),
    /// Tensor parallel rank.
    pub tensor_rank: usize,
    /// Tensor parallel world size.
    pub tensor_world: usize,
    /// Real per-layer projection weight tensors allocated on load.
    pub layer_weights: Vec<LayerWeights>,
    /// Total weight size in MB, computed from actual tensor byte sizes.
    pub weights_mb: f64,
    /// Whether this shard is loaded.
    pub loaded: bool,
}

impl ModelShard {
    /// Create a new model shard.
    #[must_use]
    pub fn new(id: usize, layer_start: usize, layer_end: usize) -> Self {
        Self {
            id,
            layer_range: (layer_start, layer_end),
            tensor_rank: 0,
            tensor_world: 1,
            layer_weights: Vec::new(),
            weights_mb: 0.0,
            loaded: false,
        }
    }

    /// Get the number of layers in this shard.
    #[must_use]
    pub fn num_layers(&self) -> usize {
        self.layer_range.1 - self.layer_range.0
    }

    /// Allocate real weight tensors for every layer in this shard's range,
    /// sized (and tensor-parallel sharded) from `model_config`. `weights_mb`
    /// is then computed by summing the actual allocated tensor byte sizes.
    ///
    /// # Errors
    /// Returns an error string if tensor allocation fails (e.g. invalid
    /// shape derived from a degenerate model config).
    pub fn load(&mut self, model_config: &ModelConfig) -> Result<(), String> {
        let device = Device::cpu();
        let tensor_world = self.tensor_world.max(1);

        self.layer_weights = (0..self.num_layers())
            .map(|_| LayerWeights::allocate(model_config, tensor_world, &device))
            .collect::<Result<Vec<_>, _>>()?;

        let total_bytes: usize = self
            .layer_weights
            .iter()
            .map(LayerWeights::size_bytes)
            .sum();
        self.weights_mb = total_bytes as f64 / (1024.0 * 1024.0);
        self.loaded = true;
        Ok(())
    }
}

// ============================================================================
// Distributed Communicator
// ============================================================================

/// Communicator for distributed operations.
///
/// This is a **simulated network layer**: it performs no real inter-process
/// or inter-node communication. It models realistic transfer timing (ring
/// all-reduce cost, point-to-point send/recv cost) given the cluster's
/// configured bandwidth/latency (`ClusterConfig`) and the *real* byte size
/// of the data being "transferred" (now computed from actual tensor sizes
/// rather than invented numbers).
#[derive(Debug)]
pub struct DistributedCommunicator {
    /// This node's rank.
    pub rank: usize,
    /// Total world size.
    pub world_size: usize,
    /// Cluster configuration.
    pub cluster: ClusterConfig,
    /// Communication latency in microseconds.
    pub latency_us: f64,
    /// Bandwidth in GB/s.
    pub bandwidth_gbps: f64,
    /// Simulated message count.
    message_count: usize,
}

impl DistributedCommunicator {
    /// Create a new communicator.
    #[must_use]
    pub fn new(rank: usize, cluster: ClusterConfig) -> Self {
        let world_size = cluster.num_nodes();
        let bandwidth_gbps = if !cluster.nodes.is_empty() {
            cluster.nodes[0].network.bandwidth_gbps()
        } else {
            80.0 // Default to TB5
        };

        Self {
            rank,
            world_size,
            cluster,
            latency_us: 10.0, // 10us latency
            bandwidth_gbps,
            message_count: 0,
        }
    }

    /// Simulate an all-reduce operation.
    pub fn all_reduce(&mut self, data_size_mb: f64) -> f64 {
        self.message_count += 1;
        // Ring all-reduce: 2 * (n-1) / n * data_size
        let transfer_size =
            2.0 * (self.world_size - 1) as f64 / self.world_size as f64 * data_size_mb;
        let transfer_time_ms = transfer_size * 1000.0 / self.bandwidth_gbps;
        let latency_ms = self.latency_us / 1000.0 * 2.0 * (self.world_size - 1) as f64;
        transfer_time_ms + latency_ms
    }

    /// Simulate a point-to-point send.
    pub fn send(&mut self, _dest: usize, data_size_mb: f64) -> f64 {
        self.message_count += 1;
        let transfer_time_ms = data_size_mb * 1000.0 / self.bandwidth_gbps;
        transfer_time_ms + self.latency_us / 1000.0
    }

    /// Simulate a point-to-point receive.
    pub fn recv(&mut self, _source: usize, data_size_mb: f64) -> f64 {
        self.message_count += 1;
        let transfer_time_ms = data_size_mb * 1000.0 / self.bandwidth_gbps;
        transfer_time_ms + self.latency_us / 1000.0
    }

    /// Get total message count.
    #[must_use]
    pub fn message_count(&self) -> usize {
        self.message_count
    }
}

// ============================================================================
// Pipeline Scheduler
// ============================================================================

/// Scheduler for pipeline parallel execution.
#[derive(Debug)]
pub struct PipelineScheduler {
    /// Number of pipeline stages.
    pub num_stages: usize,
    /// Number of micro-batches.
    pub num_micro_batches: usize,
    /// Current micro-batch index.
    current_micro_batch: usize,
    /// Schedule type.
    pub schedule: distllm_shared::PipelineSchedule,
}

impl PipelineScheduler {
    /// Create a new pipeline scheduler.
    #[must_use]
    pub fn new(num_stages: usize, num_micro_batches: usize) -> Self {
        Self {
            num_stages,
            num_micro_batches,
            current_micro_batch: 0,
            schedule: distllm_shared::PipelineSchedule::GPipe,
        }
    }

    /// Get the next stage to execute for the current micro-batch.
    #[must_use]
    pub fn next_stage(&self) -> Option<usize> {
        if self.current_micro_batch < self.num_micro_batches {
            Some(self.current_micro_batch % self.num_stages)
        } else {
            None
        }
    }

    /// Advance to the next micro-batch.
    pub fn advance(&mut self) {
        self.current_micro_batch += 1;
    }

    /// Reset the scheduler.
    pub fn reset(&mut self) {
        self.current_micro_batch = 0;
    }

    /// Calculate pipeline bubble ratio.
    #[must_use]
    pub fn bubble_ratio(&self) -> f64 {
        let total_slots = self.num_stages * self.num_micro_batches;
        let bubble_slots = self.num_stages - 1;
        bubble_slots as f64 / total_slots as f64
    }
}

// ============================================================================
// Distributed LLM System
// ============================================================================

/// Main distributed LLM inference system.
#[derive(Debug)]
pub struct DistributedLLM {
    /// Model configuration.
    pub model_config: ModelConfig,
    /// Parallelism configuration.
    pub parallel_config: ParallelismConfig,
    /// Model shards on this node.
    pub shards: Vec<ModelShard>,
    /// Distributed communicator.
    pub communicator: DistributedCommunicator,
    /// Pipeline scheduler.
    pub scheduler: PipelineScheduler,
    /// Layer assignments.
    pub layer_assignments: Vec<LayerAssignment>,
    /// Distributed attention module.
    pub attention: DistributedAttention,
    /// KV cache.
    pub kv_cache: KVCache,
    /// Tensor parallel module.
    pub tensor_parallel: TensorParallel,
    /// Pipeline parallel module.
    pub pipeline_parallel: PipelineParallel,
    /// Whether model is loaded.
    pub loaded: bool,
    /// RNG state for sampling.
    rng_state: u64,
    /// Measured prefill latency (ms) from the most recent `generate()` call.
    last_prefill_latency_ms: f64,
    /// Measured average per-token decode latency (ms) from the most recent
    /// `generate()` call.
    last_decode_latency_ms: f64,
    /// Measured tokens/sec throughput from the most recent `generate()` call.
    last_tokens_per_second: f64,
    /// Accumulated real attention compute time (ms) across all layers/shards
    /// in the most recent `generate()` call.
    last_attention_compute_ms: f64,
    /// Accumulated simulated all-reduce time (ms) across all layers/shards
    /// in the most recent `generate()` call.
    last_allreduce_time_ms: f64,
    /// Total wall-clock time (ms) of the most recent `generate()` call.
    last_total_time_ms: f64,
}

impl DistributedLLM {
    /// Create a new distributed LLM system.
    #[must_use]
    pub fn new(
        model_config: ModelConfig,
        parallel_config: ParallelismConfig,
        cluster: ClusterConfig,
        rank: usize,
    ) -> Self {
        let communicator = DistributedCommunicator::new(rank, cluster);
        let scheduler = PipelineScheduler::new(
            parallel_config.pipeline_parallel,
            parallel_config.num_micro_batches,
        );

        let kv_cache_config = KVCacheConfig {
            max_seq_len: model_config.max_seq_len,
            num_layers: model_config.num_layers,
            num_kv_heads: model_config.num_kv_heads / parallel_config.tensor_parallel,
            head_dim: model_config.head_dim,
            dtype: model_config.dtype,
            ..Default::default()
        };

        let attention = DistributedAttention::new(
            model_config.num_heads / parallel_config.tensor_parallel,
            model_config.num_kv_heads / parallel_config.tensor_parallel,
            model_config.head_dim,
            parallel_config.tensor_parallel,
        );

        let kv_cache = KVCache::new(kv_cache_config);

        let tensor_parallel = TensorParallel::new(
            parallel_config.tensor_parallel,
            rank % parallel_config.tensor_parallel,
        );

        let pipeline_parallel = PipelineParallel::new(
            parallel_config.pipeline_parallel,
            rank / parallel_config.tensor_parallel,
        );

        Self {
            model_config,
            parallel_config,
            shards: Vec::new(),
            communicator,
            scheduler,
            layer_assignments: Vec::new(),
            attention,
            kv_cache,
            tensor_parallel,
            pipeline_parallel,
            loaded: false,
            rng_state: 42,
            last_prefill_latency_ms: 0.0,
            last_decode_latency_ms: 0.0,
            last_tokens_per_second: 0.0,
            last_attention_compute_ms: 0.0,
            last_allreduce_time_ms: 0.0,
            last_total_time_ms: 0.0,
        }
    }

    /// Load the model, distributing layers across nodes.
    pub fn load_model(&mut self) {
        // Use layer partitioner to assign layers
        let partitioner = LayerPartitioner::new(
            self.model_config.num_layers,
            self.parallel_config.world_size,
            self.parallel_config.tensor_parallel,
            self.parallel_config.pipeline_parallel,
        );

        self.layer_assignments = partitioner.partition();

        // Create shards for this node's layers
        let my_layers: Vec<_> = self
            .layer_assignments
            .iter()
            .filter(|a| a.node_id == self.communicator.rank)
            .collect();

        if !my_layers.is_empty() {
            let start = my_layers.first().map_or(0, |a| a.layer_id);
            let end = my_layers.last().map_or(0, |a| a.layer_id + 1);

            let mut shard = ModelShard::new(self.communicator.rank, start, end);
            shard.tensor_rank = self.tensor_parallel.rank;
            shard.tensor_world = self.tensor_parallel.world_size;
            shard
                .load(&self.model_config)
                .expect("real tensor weight allocation failed (CPU device should always succeed)");
            self.shards.push(shard);
        }

        self.loaded = true;
    }

    /// Generate text from a prompt.
    pub fn generate(
        &mut self,
        request: &InferenceRequest,
        config: &GenerationConfig,
    ) -> InferenceResult {
        let start = std::time::Instant::now();

        if !self.loaded {
            self.load_model();
        }

        // Tokenize prompt (simulated)
        let prompt_tokens = self.tokenize(&request.prompt);
        let num_prompt_tokens = prompt_tokens.len();

        // Reset per-call accumulators (real measurements, not constants).
        self.last_attention_compute_ms = 0.0;
        self.last_allreduce_time_ms = 0.0;

        // Prefill phase: runs a real matmul/softmax attention forward pass
        // per layer via `DistributedAttention::forward`.
        let prefill_start = std::time::Instant::now();
        self.prefill(&prompt_tokens);
        let prefill_time = prefill_start.elapsed().as_secs_f64() * 1000.0;
        self.last_prefill_latency_ms = prefill_time;

        // Decode phase
        let mut generated_tokens = Vec::new();
        let mut token_latencies = Vec::new();
        let max_tokens = config.max_new_tokens.min(request.max_tokens);

        let ttft = prefill_time; // Time to first token

        for _ in 0..max_tokens {
            let token_start = std::time::Instant::now();

            // Generate next token
            let (token, finished) = self.decode_step(&config.sampling);

            let token_time = token_start.elapsed().as_secs_f64() * 1000.0;
            token_latencies.push(token_time);

            generated_tokens.push(token);

            if finished {
                break;
            }
        }

        let total_time = start.elapsed().as_secs_f64() * 1000.0;
        let num_generated = generated_tokens.len();
        let tokens_per_second = num_generated as f64 / (total_time / 1000.0);

        self.last_tokens_per_second = tokens_per_second;
        self.last_total_time_ms = total_time;
        self.last_decode_latency_ms = if token_latencies.is_empty() {
            0.0
        } else {
            token_latencies.iter().sum::<f64>() / token_latencies.len() as f64
        };

        // Detokenize
        let text = self.detokenize(&generated_tokens);

        let finish_reason = if num_generated >= max_tokens {
            FinishReason::MaxTokens
        } else {
            FinishReason::EndOfSequence
        };

        InferenceResult {
            tokens: generated_tokens,
            text,
            latency_ms: total_time,
            tokens_per_second,
            time_to_first_token_ms: ttft,
            token_latencies_ms: token_latencies,
            prompt_tokens: num_prompt_tokens,
            completion_tokens: num_generated,
            finish_reason,
        }
    }

    /// Tokenize input text.
    ///
    /// Simplified simulation: a real subword tokenizer/vocabulary is out of
    /// scope for this demo, so tokens are assigned by whitespace-split word
    /// position rather than via a learned vocabulary.
    fn tokenize(&self, text: &str) -> Vec<u32> {
        text.split_whitespace()
            .enumerate()
            .map(|(i, _)| i as u32 + 1)
            .collect()
    }

    /// Detokenize tokens to text (simulated; see [`Self::tokenize`]).
    fn detokenize(&self, tokens: &[String]) -> String {
        tokens.join(" ")
    }

    /// Run prefill phase (process prompt).
    ///
    /// Runs a real attention forward pass (Q/K/V allocation, matmul,
    /// softmax) for every layer assigned to this shard, and accumulates the
    /// measured compute time. The tensor-parallel all-reduce remains a
    /// simulated network operation (see [`DistributedCommunicator`]), now
    /// sized from the real hidden-state byte count.
    fn prefill(&mut self, tokens: &[u32]) {
        let seq_len = tokens.len();

        for shard in &self.shards {
            for layer_id in shard.layer_range.0..shard.layer_range.1 {
                let (_output, elapsed) = self
                    .attention
                    .forward(seq_len, layer_id)
                    .expect("real attention forward pass failed");
                self.last_attention_compute_ms += elapsed.as_secs_f64() * 1000.0;

                self.kv_cache.update(layer_id, seq_len);
            }
        }

        // All-reduce for tensor parallel (simulated network transfer, real byte size).
        if self.tensor_parallel.world_size > 1 {
            let hidden_size_mb =
                (self.model_config.hidden_dim * seq_len * 2) as f64 / (1024.0 * 1024.0);
            self.last_allreduce_time_ms += self.communicator.all_reduce(hidden_size_mb);
        }
    }

    /// Run one decode step (generate one token).
    ///
    /// Runs a real single-token attention forward pass per layer, same as
    /// [`Self::prefill`] but with `seq_len = 1`.
    fn decode_step(&mut self, sampling: &distllm_shared::SamplingParams) -> (String, bool) {
        for shard in &self.shards {
            for layer_id in shard.layer_range.0..shard.layer_range.1 {
                let (_output, elapsed) = self
                    .attention
                    .forward(1, layer_id)
                    .expect("real attention forward pass failed");
                self.last_attention_compute_ms += elapsed.as_secs_f64() * 1000.0;

                self.kv_cache.update(layer_id, 1);
            }
        }

        // All-reduce for tensor parallel (simulated network transfer, real byte size).
        if self.tensor_parallel.world_size > 1 {
            let hidden_size_mb = (self.model_config.hidden_dim * 2) as f64 / (1024.0 * 1024.0);
            self.last_allreduce_time_ms += self.communicator.all_reduce(hidden_size_mb);
        }

        // Sample token (simplified simulation: word-list sampling, not a real
        // learned vocabulary distribution).
        let token = self.sample_token(sampling);

        // Check for end of sequence (simulated: fixed probability, not a
        // learned EOS distribution).
        let is_eos = self.random() < 0.02; // 2% chance of EOS

        (token, is_eos)
    }

    /// Sample a token based on sampling parameters.
    fn sample_token(&mut self, _sampling: &distllm_shared::SamplingParams) -> String {
        // Simulated token sampling
        let words = [
            "the", "a", "an", "is", "was", "are", "were", "be", "been", "being", "have", "has",
            "had", "do", "does", "did", "will", "would", "could", "should", "may", "might", "must",
            "shall", "can", "need", "and", "but", "or", "if", "when", "while", "as", "because",
            "although", "this", "that", "these", "those", "it", "they", "we", "you", "I", "which",
            "who", "what", "where", "how", "why", "when",
        ];

        let idx = (self.random() * words.len() as f64) as usize;
        words[idx.min(words.len() - 1)].to_string()
    }

    /// Get performance metrics.
    ///
    /// `memory_utilization` is derived from actual allocated tensor byte
    /// sizes (`ModelShard::weights_mb`). Latency/throughput/utilization
    /// figures are derived from real `Instant`-measured timings accumulated
    /// during the most recent `generate()` call (zero before the first
    /// call). `compute_utilization` and `network_utilization` are the
    /// measured attention-compute and all-reduce shares of total wall-clock
    /// time, respectively — real ratios, not fixed constants.
    #[must_use]
    pub fn get_metrics(&self) -> PerformanceMetrics {
        let memory_used: f64 = self.shards.iter().map(|s| s.weights_mb).sum();
        let total_memory = self.communicator.cluster.total_memory() * 1024.0; // Convert to MB

        let (compute_utilization, network_utilization) = if self.last_total_time_ms > 0.0 {
            (
                (self.last_attention_compute_ms / self.last_total_time_ms).min(1.0),
                (self.last_allreduce_time_ms / self.last_total_time_ms).min(1.0),
            )
        } else {
            (0.0, 0.0)
        };

        PerformanceMetrics {
            prefill_latency_ms: self.last_prefill_latency_ms,
            decode_latency_ms: self.last_decode_latency_ms,
            tokens_per_second: self.last_tokens_per_second,
            memory_utilization: memory_used / total_memory,
            compute_utilization,
            network_utilization,
            bubble_ratio: self.scheduler.bubble_ratio(),
            allreduce_time_ms: self.last_allreduce_time_ms,
        }
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

// ============================================================================
// Run Demo
// ============================================================================

/// Run the distributed LLM demo.
///
/// Uses [`sample_data::tiny_realcompute_config`] rather than a real 70B/405B
/// config: this demo now performs *real* `Tensor::randn` weight allocation
/// and real attention matmul/softmax compute (see module docs), and a
/// genuine 70B-parameter model's weights cannot be allocated on typical
/// demo/CI hardware. The cluster topology and parallelism strategy are
/// still the full 4-node Thunderbolt 5 configuration.
#[must_use]
pub fn run_demo() -> InferenceResult {
    let model_config = sample_data::tiny_realcompute_config();
    let parallel_config = sample_data::four_node_parallel_config();
    let cluster = sample_data::four_node_cluster();

    let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
    llm.load_model();

    let request = sample_data::chat_inference_request();
    let gen_config = GenerationConfig {
        max_new_tokens: 128,
        do_sample: true,
        ..Default::default()
    };

    llm.generate(&request, &gen_config)
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_model_shard_creation() {
        let shard = ModelShard::new(0, 0, 20);
        assert_eq!(shard.num_layers(), 20);
        assert!(!shard.loaded);
    }

    #[test]
    fn test_model_shard_load() {
        let mut shard = ModelShard::new(0, 0, 4);
        let model_config = sample_data::tiny_realcompute_config();
        shard
            .load(&model_config)
            .expect("weight allocation should succeed");
        assert!(shard.loaded);
        assert_eq!(shard.layer_weights.len(), 4);
        // Real tensor-derived byte size must be positive.
        assert!(shard.weights_mb > 0.0);
    }

    #[test]
    fn test_communicator_creation() {
        let cluster = sample_data::four_node_cluster();
        let comm = DistributedCommunicator::new(0, cluster);
        assert_eq!(comm.rank, 0);
        assert_eq!(comm.world_size, 4);
    }

    #[test]
    fn test_all_reduce_time() {
        let cluster = sample_data::four_node_cluster();
        let mut comm = DistributedCommunicator::new(0, cluster);
        let time = comm.all_reduce(100.0);
        assert!(time > 0.0);
        assert_eq!(comm.message_count(), 1);
    }

    #[test]
    fn test_pipeline_scheduler() {
        let mut scheduler = PipelineScheduler::new(4, 8);
        assert!(scheduler.next_stage().is_some());
        scheduler.advance();
        assert_eq!(scheduler.current_micro_batch, 1);
    }

    #[test]
    fn test_bubble_ratio() {
        let scheduler = PipelineScheduler::new(4, 16);
        let ratio = scheduler.bubble_ratio();
        // 3 bubble slots out of 64 total = 0.046875
        assert!(ratio < 0.1);
    }

    #[test]
    fn test_distributed_llm_creation() {
        let model_config = sample_data::llama_70b_config();
        let parallel_config = sample_data::four_node_parallel_config();
        let cluster = sample_data::four_node_cluster();

        let llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
        assert!(!llm.loaded);
    }

    #[test]
    fn test_load_model() {
        // Real weight allocation happens here, so use the tiny compute
        // config rather than a genuine 70B-scale model (see
        // `sample_data::tiny_realcompute_config` docs).
        let model_config = sample_data::tiny_realcompute_config();
        let parallel_config = sample_data::four_node_parallel_config();
        let cluster = sample_data::four_node_cluster();

        let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
        llm.load_model();

        assert!(llm.loaded);
        assert!(!llm.layer_assignments.is_empty());
    }

    #[test]
    fn test_generate() {
        let model_config = sample_data::tiny_realcompute_config();
        let parallel_config = ParallelismConfig {
            tensor_parallel: 1,
            pipeline_parallel: 1,
            world_size: 1,
            ..Default::default()
        };
        let cluster = distllm_shared::ClusterConfig::default();

        let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);

        let request = InferenceRequest {
            prompt: "Hello world".to_string(),
            max_tokens: 10,
            ..Default::default()
        };

        let config = GenerationConfig {
            max_new_tokens: 10,
            ..Default::default()
        };

        let result = llm.generate(&request, &config);

        assert!(!result.tokens.is_empty());
        assert!(result.latency_ms > 0.0);
        assert!(result.tokens_per_second > 0.0);
    }

    #[test]
    fn test_metrics() {
        let model_config = sample_data::tiny_realcompute_config();
        let parallel_config = ParallelismConfig {
            tensor_parallel: 1,
            pipeline_parallel: 1,
            world_size: 1,
            ..Default::default()
        };
        let cluster = distllm_shared::ClusterConfig::default();

        let mut llm = DistributedLLM::new(model_config, parallel_config, cluster, 0);
        llm.load_model();

        let request = InferenceRequest {
            prompt: "Hello world".to_string(),
            max_tokens: 4,
            ..Default::default()
        };
        let config = GenerationConfig {
            max_new_tokens: 4,
            ..Default::default()
        };
        let _ = llm.generate(&request, &config);

        let metrics = llm.get_metrics();
        assert!(metrics.tokens_per_second > 0.0);
        assert!(metrics.memory_utilization > 0.0);
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(!result.tokens.is_empty());
        assert!(result.completion_tokens > 0);
    }
}