//! Shared types for the DistributedLLM trillion-parameter inference demo.
//!
//! This crate provides IPC types for distributed large language model inference
//! across Thunderbolt 5 Mac clusters. Supports tensor parallelism, pipeline
//! parallelism, and hybrid strategies for models up to trillion parameters.

use serde::{Deserialize, Serialize};

// ============================================================================
// Model Configuration
// ============================================================================

/// Configuration for a large language model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Model name (e.g., "llama-70b", "llama-405b").
    pub name: String,
    /// Total number of parameters (in billions).
    pub num_params: f64,
    /// Number of transformer layers.
    pub num_layers: usize,
    /// Hidden dimension size.
    pub hidden_dim: usize,
    /// Number of attention heads.
    pub num_heads: usize,
    /// Number of key-value heads (for GQA).
    pub num_kv_heads: usize,
    /// Intermediate dimension (FFN).
    pub intermediate_dim: usize,
    /// Vocabulary size.
    pub vocab_size: usize,
    /// Maximum sequence length.
    pub max_seq_len: usize,
    /// Head dimension.
    pub head_dim: usize,
    /// RoPE theta for positional encoding.
    pub rope_theta: f64,
    /// Data type for weights.
    pub dtype: DataType,
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            name: "llama-7b".to_string(),
            num_params: 7.0,
            num_layers: 32,
            hidden_dim: 4096,
            num_heads: 32,
            num_kv_heads: 32,
            intermediate_dim: 11008,
            vocab_size: 32000,
            max_seq_len: 4096,
            head_dim: 128,
            rope_theta: 10000.0,
            dtype: DataType::BFloat16,
        }
    }
}

impl ModelConfig {
    /// Calculate estimated memory required in GB.
    #[must_use]
    pub fn estimated_memory_gb(&self) -> f64 {
        let bytes_per_param = match self.dtype {
            DataType::Float32 => 4.0,
            DataType::Float16 | DataType::BFloat16 => 2.0,
            DataType::Int8 | DataType::FP8 => 1.0,
            DataType::Int4 => 0.5,
        };
        self.num_params * bytes_per_param
    }

    /// Calculate memory per layer in MB.
    #[must_use]
    pub fn memory_per_layer_mb(&self) -> f64 {
        self.estimated_memory_gb() * 1024.0 / self.num_layers as f64
    }
}

/// Data types for model weights.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum DataType {
    Float32,
    Float16,
    #[default]
    BFloat16,
    Int8,
    Int4,
    FP8,
}

// ============================================================================
// Parallelism Configuration
// ============================================================================

/// Configuration for distributed parallelism.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ParallelismConfig {
    /// Tensor parallelism degree (splits attention heads/FFN).
    pub tensor_parallel: usize,
    /// Pipeline parallelism degree (splits layers).
    pub pipeline_parallel: usize,
    /// Total world size (number of nodes).
    pub world_size: usize,
    /// Number of micro-batches for pipeline.
    pub num_micro_batches: usize,
    /// Pipeline schedule type.
    pub pipeline_schedule: PipelineSchedule,
    /// Communication backend.
    pub comm_backend: CommBackend,
}

impl Default for ParallelismConfig {
    fn default() -> Self {
        Self {
            tensor_parallel: 1,
            pipeline_parallel: 1,
            world_size: 1,
            num_micro_batches: 1,
            pipeline_schedule: PipelineSchedule::GPipe,
            comm_backend: CommBackend::Thunderbolt5,
        }
    }
}

impl ParallelismConfig {
    /// Calculate data parallelism degree.
    #[must_use]
    pub fn data_parallel(&self) -> usize {
        self.world_size / (self.tensor_parallel * self.pipeline_parallel)
    }

    /// Validate parallelism configuration.
    #[must_use]
    pub fn is_valid(&self) -> bool {
        self.world_size > 0
            && self.tensor_parallel > 0
            && self.pipeline_parallel > 0
            && self
                .world_size
                .is_multiple_of(self.tensor_parallel * self.pipeline_parallel)
    }
}

/// Pipeline parallelism schedule.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum PipelineSchedule {
    /// GPipe: all-forward then all-backward.
    #[default]
    GPipe,
    /// 1F1B: interleaved forward-backward.
    OneForwardOneBackward,
    /// Interleaved 1F1B with virtual stages.
    InterleavedOneForwardOneBackward,
    /// Zero bubble schedule.
    ZeroBubble,
}

/// Communication backend for distributed inference.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum CommBackend {
    /// Thunderbolt 5 (80 Gbps bidirectional).
    #[default]
    Thunderbolt5,
    /// Thunderbolt 4 (40 Gbps bidirectional).
    Thunderbolt4,
    /// Ethernet (10/25/100 GbE).
    Ethernet,
    /// InfiniBand.
    InfiniBand,
    /// PCIe direct.
    PCIe,
    /// NVLink (for GPU clusters).
    NVLink,
}

impl CommBackend {
    /// Get theoretical bandwidth in GB/s.
    #[must_use]
    pub fn bandwidth_gbps(&self) -> f64 {
        match self {
            Self::Thunderbolt5 => 80.0,
            Self::Thunderbolt4 => 40.0,
            Self::Ethernet => 100.0,
            Self::InfiniBand => 400.0,
            Self::PCIe => 64.0,
            Self::NVLink => 900.0,
        }
    }
}

// ============================================================================
// Inference Request/Response
// ============================================================================

/// Request for text generation inference.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceRequest {
    /// Input prompt text.
    pub prompt: String,
    /// Maximum tokens to generate.
    pub max_tokens: usize,
    /// Sampling temperature.
    pub temperature: f32,
    /// Top-p (nucleus) sampling parameter.
    pub top_p: f32,
    /// Top-k sampling parameter.
    pub top_k: usize,
    /// Frequency penalty.
    pub frequency_penalty: f32,
    /// Presence penalty.
    pub presence_penalty: f32,
    /// Stop sequences.
    pub stop_sequences: Vec<String>,
    /// Random seed for reproducibility.
    pub seed: Option<u64>,
    /// Whether to stream tokens.
    pub stream: bool,
}

impl Default for InferenceRequest {
    fn default() -> Self {
        Self {
            prompt: String::new(),
            max_tokens: 256,
            temperature: 0.7,
            top_p: 0.9,
            top_k: 50,
            frequency_penalty: 0.0,
            presence_penalty: 0.0,
            stop_sequences: vec![],
            seed: None,
            stream: false,
        }
    }
}

/// Result from text generation inference.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceResult {
    /// Generated tokens (as strings).
    pub tokens: Vec<String>,
    /// Generated text.
    pub text: String,
    /// Total latency in milliseconds.
    pub latency_ms: f64,
    /// Tokens per second throughput.
    pub tokens_per_second: f64,
    /// Time to first token in milliseconds.
    pub time_to_first_token_ms: f64,
    /// Per-token latencies.
    pub token_latencies_ms: Vec<f64>,
    /// Number of prompt tokens.
    pub prompt_tokens: usize,
    /// Number of generated tokens.
    pub completion_tokens: usize,
    /// Finish reason.
    pub finish_reason: FinishReason,
}

impl Default for InferenceResult {
    fn default() -> Self {
        Self {
            tokens: vec![],
            text: String::new(),
            latency_ms: 0.0,
            tokens_per_second: 0.0,
            time_to_first_token_ms: 0.0,
            token_latencies_ms: vec![],
            prompt_tokens: 0,
            completion_tokens: 0,
            finish_reason: FinishReason::EndOfSequence,
        }
    }
}

/// Reason for finishing generation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum FinishReason {
    /// Reached end of sequence token.
    #[default]
    EndOfSequence,
    /// Reached maximum token limit.
    MaxTokens,
    /// Hit a stop sequence.
    StopSequence,
    /// Generation was cancelled.
    Cancelled,
}

// ============================================================================
// Layer Assignment
// ============================================================================

/// Assignment of a layer to a node.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LayerAssignment {
    /// Layer index.
    pub layer_id: usize,
    /// Node ID where layer is placed.
    pub node_id: usize,
    /// Memory required in MB.
    pub memory_mb: f64,
    /// Pipeline stage index.
    pub pipeline_stage: usize,
    /// Tensor parallel rank within stage.
    pub tensor_rank: usize,
}

impl LayerAssignment {
    /// Create a new layer assignment.
    #[must_use]
    pub fn new(layer_id: usize, node_id: usize, memory_mb: f64) -> Self {
        Self {
            layer_id,
            node_id,
            memory_mb,
            pipeline_stage: 0,
            tensor_rank: 0,
        }
    }
}

// ============================================================================
// Generation Configuration
// ============================================================================

/// Configuration for text generation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationConfig {
    /// Maximum new tokens to generate.
    pub max_new_tokens: usize,
    /// Minimum new tokens to generate.
    pub min_new_tokens: usize,
    /// Whether to use sampling (vs greedy).
    pub do_sample: bool,
    /// Sampling parameters.
    pub sampling: SamplingParams,
    /// Number of beams for beam search.
    pub num_beams: usize,
    /// Early stopping for beam search.
    pub early_stopping: bool,
    /// Repetition penalty.
    pub repetition_penalty: f32,
    /// Length penalty for beam search.
    pub length_penalty: f32,
    /// Number of return sequences.
    pub num_return_sequences: usize,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_new_tokens: 256,
            min_new_tokens: 1,
            do_sample: true,
            sampling: SamplingParams::default(),
            num_beams: 1,
            early_stopping: false,
            repetition_penalty: 1.0,
            length_penalty: 1.0,
            num_return_sequences: 1,
        }
    }
}

/// Sampling parameters for generation.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct SamplingParams {
    /// Temperature for softmax.
    pub temperature: f32,
    /// Top-p (nucleus) sampling.
    pub top_p: f32,
    /// Top-k sampling.
    pub top_k: usize,
    /// Typical-p sampling.
    pub typical_p: f32,
    /// Eta cutoff for entropy-based sampling.
    pub eta_cutoff: f32,
    /// Epsilon cutoff for probability-based sampling.
    pub epsilon_cutoff: f32,
}

impl Default for SamplingParams {
    fn default() -> Self {
        Self {
            temperature: 0.7,
            top_p: 0.9,
            top_k: 50,
            typical_p: 1.0,
            eta_cutoff: 0.0,
            epsilon_cutoff: 0.0,
        }
    }
}

// ============================================================================
// Cluster Configuration
// ============================================================================

/// Configuration for a compute cluster.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterConfig {
    /// Cluster name.
    pub name: String,
    /// List of nodes in the cluster.
    pub nodes: Vec<NodeConfig>,
    /// Network topology.
    pub topology: NetworkTopology,
    /// Total cluster memory in GB.
    pub total_memory_gb: f64,
}

impl Default for ClusterConfig {
    fn default() -> Self {
        Self {
            name: "default-cluster".to_string(),
            nodes: vec![NodeConfig::default()],
            topology: NetworkTopology::Ring,
            total_memory_gb: 192.0,
        }
    }
}

impl ClusterConfig {
    /// Get total number of nodes.
    #[must_use]
    pub fn num_nodes(&self) -> usize {
        self.nodes.len()
    }

    /// Get total memory across all nodes.
    #[must_use]
    pub fn total_memory(&self) -> f64 {
        self.nodes.iter().map(|n| n.memory_gb).sum()
    }
}

/// Configuration for a single compute node.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeConfig {
    /// Node ID.
    pub id: usize,
    /// Node hostname.
    pub hostname: String,
    /// Memory available in GB.
    pub memory_gb: f64,
    /// Number of CPU cores.
    pub cpu_cores: usize,
    /// GPU type if available.
    pub gpu_type: Option<String>,
    /// GPU memory in GB if available.
    pub gpu_memory_gb: Option<f64>,
    /// Network interface.
    pub network: CommBackend,
}

impl Default for NodeConfig {
    fn default() -> Self {
        Self {
            id: 0,
            hostname: "localhost".to_string(),
            memory_gb: 192.0,
            cpu_cores: 12,
            gpu_type: None,
            gpu_memory_gb: None,
            network: CommBackend::Thunderbolt5,
        }
    }
}

/// Network topology for cluster.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum NetworkTopology {
    /// Ring topology.
    #[default]
    Ring,
    /// Fully connected mesh.
    Mesh,
    /// Star topology with central switch.
    Star,
    /// Tree/hierarchical topology.
    Tree,
    /// Daisy chain.
    DaisyChain,
}

// ============================================================================
// KV Cache Configuration
// ============================================================================

/// Configuration for KV cache.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KVCacheConfig {
    /// Maximum sequence length to cache.
    pub max_seq_len: usize,
    /// Maximum batch size.
    pub max_batch_size: usize,
    /// Number of layers.
    pub num_layers: usize,
    /// Number of KV heads.
    pub num_kv_heads: usize,
    /// Head dimension.
    pub head_dim: usize,
    /// Data type for cache.
    pub dtype: DataType,
    /// Whether to use paged attention.
    pub paged_attention: bool,
    /// Block size for paged attention.
    pub block_size: usize,
}

impl Default for KVCacheConfig {
    fn default() -> Self {
        Self {
            max_seq_len: 4096,
            max_batch_size: 1,
            num_layers: 32,
            num_kv_heads: 32,
            head_dim: 128,
            dtype: DataType::BFloat16,
            paged_attention: true,
            block_size: 16,
        }
    }
}

impl KVCacheConfig {
    /// Calculate total cache memory in GB.
    #[must_use]
    pub fn total_memory_gb(&self) -> f64 {
        let bytes_per_elem = match self.dtype {
            DataType::Float32 => 4,
            DataType::Float16 | DataType::BFloat16 => 2,
            DataType::Int8 | DataType::FP8 => 1,
            DataType::Int4 => 1, // Rounded up
        };
        let total_elements = self.max_batch_size
            * self.max_seq_len
            * self.num_layers
            * self.num_kv_heads
            * self.head_dim
            * 2; // K and V
        (total_elements * bytes_per_elem) as f64 / (1024.0 * 1024.0 * 1024.0)
    }
}

// ============================================================================
// Performance Metrics
// ============================================================================

/// Performance metrics for inference.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PerformanceMetrics {
    /// Prefill latency in ms.
    pub prefill_latency_ms: f64,
    /// Decode latency per token in ms.
    pub decode_latency_ms: f64,
    /// Tokens per second.
    pub tokens_per_second: f64,
    /// Memory utilization (0-1).
    pub memory_utilization: f64,
    /// Compute utilization (0-1).
    pub compute_utilization: f64,
    /// Network bandwidth utilization (0-1).
    pub network_utilization: f64,
    /// Pipeline bubble ratio.
    pub bubble_ratio: f64,
    /// All-reduce time in ms.
    pub allreduce_time_ms: f64,
}

// ============================================================================
// Sample Data Functions
// ============================================================================

/// Create a LLaMA 70B model configuration.
#[must_use]
pub fn llama_70b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-70b".to_string(),
        num_params: 70.0,
        num_layers: 80,
        hidden_dim: 8192,
        num_heads: 64,
        num_kv_heads: 8,
        intermediate_dim: 28672,
        vocab_size: 32000,
        max_seq_len: 4096,
        head_dim: 128,
        rope_theta: 10000.0,
        dtype: DataType::BFloat16,
    }
}

/// Create a LLaMA 405B model configuration.
#[must_use]
pub fn llama_405b_config() -> ModelConfig {
    ModelConfig {
        name: "llama-405b".to_string(),
        num_params: 405.0,
        num_layers: 126,
        hidden_dim: 16384,
        num_heads: 128,
        num_kv_heads: 8,
        intermediate_dim: 53248,
        vocab_size: 128256,
        max_seq_len: 131072,
        head_dim: 128,
        rope_theta: 500000.0,
        dtype: DataType::BFloat16,
    }
}

/// Create a 4-node Thunderbolt 5 cluster configuration.
#[must_use]
pub fn four_node_cluster_config() -> ClusterConfig {
    ClusterConfig {
        name: "tb5-4node-cluster".to_string(),
        nodes: (0..4)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("mac-studio-{}", i),
                memory_gb: 192.0,
                cpu_cores: 24,
                gpu_type: None,
                gpu_memory_gb: None,
                network: CommBackend::Thunderbolt5,
            })
            .collect(),
        topology: NetworkTopology::Ring,
        total_memory_gb: 768.0,
    }
}

/// Create an 8-node Thunderbolt 5 cluster configuration.
#[must_use]
pub fn eight_node_cluster_config() -> ClusterConfig {
    ClusterConfig {
        name: "tb5-8node-cluster".to_string(),
        nodes: (0..8)
            .map(|i| NodeConfig {
                id: i,
                hostname: format!("mac-studio-{}", i),
                memory_gb: 192.0,
                cpu_cores: 24,
                gpu_type: None,
                gpu_memory_gb: None,
                network: CommBackend::Thunderbolt5,
            })
            .collect(),
        topology: NetworkTopology::Mesh,
        total_memory_gb: 1536.0,
    }
}

/// Create a sample chat inference request.
#[must_use]
pub fn chat_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "You are a helpful AI assistant.\n\nUser: Explain quantum computing in simple terms.\n\nAssistant:".to_string(),
        max_tokens: 512,
        temperature: 0.7,
        top_p: 0.9,
        top_k: 50,
        frequency_penalty: 0.0,
        presence_penalty: 0.0,
        stop_sequences: vec!["User:".to_string()],
        seed: Some(42),
        stream: true,
    }
}

/// Create a sample completion inference request.
#[must_use]
pub fn completion_request() -> InferenceRequest {
    InferenceRequest {
        prompt: "The Rust programming language is known for".to_string(),
        max_tokens: 256,
        temperature: 0.8,
        top_p: 0.95,
        top_k: 40,
        frequency_penalty: 0.1,
        presence_penalty: 0.1,
        stop_sequences: vec![],
        seed: None,
        stream: false,
    }
}

/// Create sample layer assignments for a model across nodes.
#[must_use]
pub fn sample_layer_assignments(num_layers: usize, num_nodes: usize) -> Vec<LayerAssignment> {
    let layers_per_node = num_layers.div_ceil(num_nodes);
    (0..num_layers)
        .map(|layer_id| {
            let node_id = layer_id / layers_per_node;
            LayerAssignment {
                layer_id,
                node_id: node_id.min(num_nodes - 1),
                memory_mb: 1024.0,
                pipeline_stage: node_id.min(num_nodes - 1),
                tensor_rank: 0,
            }
        })
        .collect()
}

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

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

    #[test]
    fn test_model_config_default() {
        let config = ModelConfig::default();
        assert_eq!(config.name, "llama-7b");
        assert_eq!(config.num_layers, 32);
    }

    #[test]
    fn test_model_memory_estimation() {
        let config = llama_70b_config();
        let memory = config.estimated_memory_gb();
        // 70B params * 2 bytes (BF16) = 140 GB
        assert!((memory - 140.0).abs() < 1.0);
    }

    #[test]
    fn test_parallelism_config() {
        let config = ParallelismConfig {
            tensor_parallel: 4,
            pipeline_parallel: 2,
            world_size: 8,
            ..Default::default()
        };
        assert!(config.is_valid());
        assert_eq!(config.data_parallel(), 1);
    }

    #[test]
    fn test_parallelism_invalid() {
        let config = ParallelismConfig {
            tensor_parallel: 3,
            pipeline_parallel: 2,
            world_size: 8,
            ..Default::default()
        };
        assert!(!config.is_valid());
    }

    #[test]
    fn test_inference_request_default() {
        let request = InferenceRequest::default();
        assert_eq!(request.max_tokens, 256);
        assert!((request.temperature - 0.7).abs() < 0.01);
    }

    #[test]
    fn test_layer_assignment() {
        let assignment = LayerAssignment::new(5, 1, 512.0);
        assert_eq!(assignment.layer_id, 5);
        assert_eq!(assignment.node_id, 1);
    }

    #[test]
    fn test_cluster_config() {
        let cluster = four_node_cluster_config();
        assert_eq!(cluster.num_nodes(), 4);
        assert!((cluster.total_memory() - 768.0).abs() < 1.0);
    }

    #[test]
    fn test_kv_cache_memory() {
        let config = KVCacheConfig {
            max_seq_len: 4096,
            max_batch_size: 1,
            num_layers: 80,
            num_kv_heads: 8,
            head_dim: 128,
            dtype: DataType::BFloat16,
            ..Default::default()
        };
        let memory = config.total_memory_gb();
        // Should be reasonable for LLaMA 70B
        assert!(memory > 0.0 && memory < 100.0);
    }

    #[test]
    fn test_llama_70b_config() {
        let config = llama_70b_config();
        assert_eq!(config.num_layers, 80);
        assert_eq!(config.num_heads, 64);
        assert_eq!(config.num_kv_heads, 8);
    }

    #[test]
    fn test_llama_405b_config() {
        let config = llama_405b_config();
        assert_eq!(config.num_layers, 126);
        assert!((config.num_params - 405.0).abs() < 1.0);
    }

    #[test]
    fn test_sample_layer_assignments() {
        let assignments = sample_layer_assignments(80, 4);
        assert_eq!(assignments.len(), 80);
        // First 20 layers should be on node 0
        assert_eq!(assignments[0].node_id, 0);
        assert_eq!(assignments[19].node_id, 0);
        // Next 20 on node 1
        assert_eq!(assignments[20].node_id, 1);
    }

    #[test]
    fn test_chat_request() {
        let request = chat_request();
        assert!(request.prompt.contains("User:"));
        assert!(request.stream);
    }

    #[test]
    fn test_completion_request() {
        let request = completion_request();
        assert!(request.prompt.contains("Rust"));
        assert!(!request.stream);
    }

    #[test]
    fn test_comm_backend_bandwidth() {
        assert_eq!(CommBackend::Thunderbolt5.bandwidth_gbps(), 80.0);
        assert_eq!(CommBackend::NVLink.bandwidth_gbps(), 900.0);
    }

    #[test]
    fn test_serialization() {
        let config = llama_70b_config();
        let json = serde_json::to_string(&config).unwrap();
        let parsed: ModelConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(config, parsed);
    }

    #[test]
    fn test_generation_config() {
        let config = GenerationConfig::default();
        assert!(config.do_sample);
        assert_eq!(config.num_beams, 1);
    }

    #[test]
    fn test_sampling_params() {
        let params = SamplingParams::default();
        assert!((params.temperature - 0.7).abs() < 0.01);
        assert_eq!(params.top_k, 50);
    }

    #[test]
    fn test_performance_metrics() {
        let metrics = PerformanceMetrics {
            tokens_per_second: 50.0,
            memory_utilization: 0.85,
            ..Default::default()
        };
        assert!(metrics.tokens_per_second > 0.0);
    }
}