//! Speculative Decoding
//!
//! Per spec §8.3: Speculative decoding with draft model for up to 3x speedup.
//! Reference: [6] Zheng et al. (2024) "SGLang: Efficient Execution of Structured LM Programs"
//!
//! ## Algorithm
//!
//! 1. Generate K speculative tokens with fast draft model
//! 2. Verify all K tokens with single forward pass of target model
//! 3. Accept tokens matching probability distribution
//! 4. Resample at first rejection point
//!
//! ## Benefits
//!
//! - Up to 3x speedup for greedy decoding
//! - Maintains output quality (mathematically equivalent)
//! - Reduces memory-bound bottleneck of autoregressive decoding

// Module-level clippy allows
#![allow(clippy::must_use_candidate)]
#![allow(clippy::return_self_not_must_use)]
#![allow(clippy::missing_errors_doc)]

use serde::{Deserialize, Serialize};
use std::time::Instant;
use thiserror::Error;

/// Error type for speculative decoding
#[derive(Debug, Error)]
pub enum SpeculativeError {
    /// Draft model error
    #[error("Draft model error: {0}")]
    DraftModelError(String),

    /// Target model error
    #[error("Target model error: {0}")]
    TargetModelError(String),

    /// Invalid speculation length
    #[error("Invalid speculation length: {0}")]
    InvalidSpecLength(usize),

    /// Verification failed
    #[error("Verification failed at position {position}")]
    VerificationFailed {
        /// Position where verification failed
        position: usize,
    },
}

/// Speculative decoding statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SpeculativeStats {
    /// Total speculative iterations
    pub iterations: u64,
    /// Total tokens speculated
    pub tokens_speculated: u64,
    /// Total tokens accepted
    pub tokens_accepted: u64,
    /// Average acceptance rate
    pub acceptance_rate: f32,
    /// Average speculation length
    pub avg_spec_length: f32,
    /// Total time saved (estimated)
    pub time_saved_ms: f64,
    /// Draft model time (ms)
    pub draft_time_ms: f64,
    /// Target model time (ms)
    pub target_time_ms: f64,
}

impl SpeculativeStats {
    /// Update stats after an iteration
    pub fn record_iteration(
        &mut self,
        speculated: usize,
        accepted: usize,
        draft_ms: f64,
        target_ms: f64,
    ) {
        self.iterations += 1;
        self.tokens_speculated += speculated as u64;
        self.tokens_accepted += accepted as u64;
        self.draft_time_ms += draft_ms;
        self.target_time_ms += target_ms;

        // Update running averages
        if self.tokens_speculated > 0 {
            self.acceptance_rate = self.tokens_accepted as f32 / self.tokens_speculated as f32;
        }
        if self.iterations > 0 {
            self.avg_spec_length = self.tokens_speculated as f32 / self.iterations as f32;
        }

        // Estimate time saved (accepted tokens would have required sequential target calls)
        // Each accepted token saves ~target_time/batch_size
        let time_per_token = target_ms / speculated.max(1) as f64;
        self.time_saved_ms += (accepted.saturating_sub(1)) as f64 * time_per_token;
    }

    /// Get speedup ratio
    pub fn speedup(&self) -> f32 {
        if self.tokens_accepted == 0 {
            return 1.0;
        }
        // Speedup = (tokens * time_per_token) / actual_time
        // Assuming draft is ~10x faster than target
        let draft_tokens_equivalent = self.tokens_speculated as f64 * 0.1;
        let baseline_time = self.tokens_accepted as f64;
        let actual_time = draft_tokens_equivalent + self.iterations as f64;

        if actual_time > 0.0 {
            (baseline_time / actual_time) as f32
        } else {
            1.0
        }
    }
}

/// Token with probability
#[derive(Debug, Clone)]
pub struct TokenProb {
    /// Token ID
    pub token: u32,
    /// Log probability
    pub log_prob: f32,
}

impl TokenProb {
    /// Create a new token with probability
    pub fn new(token: u32, log_prob: f32) -> Self {
        Self { token, log_prob }
    }

    /// Get probability (exp of log_prob)
    pub fn prob(&self) -> f32 {
        self.log_prob.exp()
    }
}

/// Speculative decoding result
#[derive(Debug, Clone)]
pub struct SpeculativeResult {
    /// Accepted tokens
    pub accepted_tokens: Vec<u32>,
    /// Number of tokens speculated
    pub num_speculated: usize,
    /// Number of tokens accepted
    pub num_accepted: usize,
    /// Token that was resampled (if any)
    pub resampled_token: Option<u32>,
    /// Time taken for draft model (ms)
    pub draft_time_ms: f64,
    /// Time taken for target model (ms)
    pub target_time_ms: f64,
}

impl SpeculativeResult {
    /// Get acceptance rate for this iteration
    pub fn acceptance_rate(&self) -> f32 {
        if self.num_speculated == 0 {
            return 0.0;
        }
        self.num_accepted as f32 / self.num_speculated as f32
    }

    /// Check if all tokens were accepted
    pub fn all_accepted(&self) -> bool {
        self.num_accepted == self.num_speculated
    }
}

/// Trait for models that can be used in speculative decoding
pub trait SpeculativeModel {
    /// Generate next token logits
    fn forward(&self, input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError>;

    /// Sample from logits
    fn sample(&self, logits: &[f32]) -> Result<TokenProb, SpeculativeError>;

    /// Get vocabulary size
    fn vocab_size(&self) -> usize;

    /// Get EOS token
    fn eos_token(&self) -> u32;
}

/// Speculative decoder with draft and target models
pub struct SpeculativeDecoder<D: SpeculativeModel, T: SpeculativeModel> {
    /// Draft model (smaller, faster)
    draft: D,
    /// Target model (larger, slower)
    target: T,
    /// Number of tokens to speculate per iteration
    spec_length: usize,
    /// Statistics
    stats: SpeculativeStats,
}

impl<D: SpeculativeModel, T: SpeculativeModel> SpeculativeDecoder<D, T> {
    /// Create a new speculative decoder
    pub fn new(draft: D, target: T, spec_length: usize) -> Result<Self, SpeculativeError> {
        if spec_length == 0 || spec_length > 32 {
            return Err(SpeculativeError::InvalidSpecLength(spec_length));
        }

        Ok(Self {
            draft,
            target,
            spec_length,
            stats: SpeculativeStats::default(),
        })
    }

    /// Get speculation length
    pub fn spec_length(&self) -> usize {
        self.spec_length
    }

    /// Set speculation length
    pub fn set_spec_length(&mut self, spec_length: usize) -> Result<(), SpeculativeError> {
        if spec_length == 0 || spec_length > 32 {
            return Err(SpeculativeError::InvalidSpecLength(spec_length));
        }
        self.spec_length = spec_length;
        Ok(())
    }

    /// Generate one iteration of speculative decoding
    pub fn decode_iteration(
        &mut self,
        context: &[u32],
    ) -> Result<SpeculativeResult, SpeculativeError> {
        let mut accepted_tokens = Vec::new();
        let mut draft_tokens = Vec::new();
        let mut draft_probs = Vec::new();

        // 1. Generate speculative tokens with draft model
        let draft_start = Instant::now();
        let mut current_context = context.to_vec();

        for _ in 0..self.spec_length {
            let logits = self
                .draft
                .forward(&current_context)
                .map_err(|e| SpeculativeError::DraftModelError(e.to_string()))?;
            let token_prob = self
                .draft
                .sample(&logits)
                .map_err(|e| SpeculativeError::DraftModelError(e.to_string()))?;

            let token = token_prob.token;
            draft_tokens.push(token);
            draft_probs.push(token_prob);
            current_context.push(token);

            // Stop if EOS
            if token == self.draft.eos_token() {
                break;
            }
        }
        let draft_time = draft_start.elapsed();

        // 2. Verify with target model (single forward pass for all positions)
        let target_start = Instant::now();
        let mut verify_context = context.to_vec();
        verify_context.extend(&draft_tokens);

        // Get target logits for verification
        let target_logits = self
            .target
            .forward(&verify_context)
            .map_err(|e| SpeculativeError::TargetModelError(e.to_string()))?;

        let target_time = target_start.elapsed();

        // 3. Accept/reject based on probability matching
        let mut resampled_token = None;

        for draft_prob in &draft_probs {
            // Simple acceptance: compare draft and target probabilities
            // In practice, use more sophisticated rejection sampling
            let target_token = self
                .target
                .sample(&target_logits)
                .map_err(|e| SpeculativeError::TargetModelError(e.to_string()))?;

            // Accept if draft token matches target distribution
            // Simplified: accept if tokens match or random acceptance
            if self.should_accept(draft_prob, &target_token) {
                accepted_tokens.push(draft_prob.token);
            } else {
                // Resample from target distribution
                resampled_token = Some(target_token.token);
                accepted_tokens.push(target_token.token);
                break;
            }
        }

        let num_speculated = draft_tokens.len();
        let num_accepted = accepted_tokens.len();

        // Update stats
        self.stats.record_iteration(
            num_speculated,
            num_accepted,
            draft_time.as_secs_f64() * 1000.0,
            target_time.as_secs_f64() * 1000.0,
        );

        Ok(SpeculativeResult {
            accepted_tokens,
            num_speculated,
            num_accepted,
            resampled_token,
            draft_time_ms: draft_time.as_secs_f64() * 1000.0,
            target_time_ms: target_time.as_secs_f64() * 1000.0,
        })
    }

    /// Acceptance criterion for speculative decoding
    #[allow(clippy::unused_self)] // Will use self for config-based thresholds
    fn should_accept(&self, draft: &TokenProb, target: &TokenProb) -> bool {
        // Simple acceptance: tokens must match
        // More sophisticated: use probability ratio for rejection sampling
        if draft.token == target.token {
            return true;
        }

        // Probabilistic acceptance based on ratio
        let ratio = target.prob() / draft.prob().max(1e-10);
        ratio >= 1.0 || ratio > 0.5 // Simplified threshold
    }

    /// Get statistics
    pub fn stats(&self) -> &SpeculativeStats {
        &self.stats
    }

    /// Reset statistics
    pub fn reset_stats(&mut self) {
        self.stats = SpeculativeStats::default();
    }
}

/// Configuration for speculative decoding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpeculativeConfig {
    /// Number of tokens to speculate
    pub spec_length: usize,
    /// Minimum acceptance rate before adapting spec_length
    pub min_acceptance_rate: f32,
    /// Enable adaptive speculation length
    pub adaptive: bool,
    /// Maximum speculation length for adaptive mode
    pub max_spec_length: usize,
}

impl Default for SpeculativeConfig {
    fn default() -> Self {
        Self {
            spec_length: 4,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        }
    }
}

impl SpeculativeConfig {
    /// Create a new config with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set speculation length
    pub fn with_spec_length(mut self, spec_length: usize) -> Self {
        self.spec_length = spec_length;
        self
    }

    /// Enable/disable adaptive mode
    pub fn with_adaptive(mut self, adaptive: bool) -> Self {
        self.adaptive = adaptive;
        self
    }

    /// Adapt speculation length based on acceptance rate
    pub fn adapt_spec_length(&mut self, acceptance_rate: f32) {
        if !self.adaptive {
            return;
        }

        if acceptance_rate > 0.8 && self.spec_length < self.max_spec_length {
            // High acceptance: increase speculation
            self.spec_length = (self.spec_length + 1).min(self.max_spec_length);
        } else if acceptance_rate < self.min_acceptance_rate && self.spec_length > 1 {
            // Low acceptance: decrease speculation
            self.spec_length = (self.spec_length - 1).max(1);
        }
    }
}

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

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

    /// Mock model for testing
    struct MockModel {
        vocab_size: usize,
        eos_token: u32,
        /// Fixed token to return
        fixed_token: u32,
    }

    impl MockModel {
        fn new(vocab_size: usize, fixed_token: u32) -> Self {
            Self {
                vocab_size,
                eos_token: 0,
                fixed_token,
            }
        }
    }

    impl SpeculativeModel for MockModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            // Return uniform logits
            Ok(vec![0.0; self.vocab_size])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Ok(TokenProb::new(self.fixed_token, -1.0))
        }

        fn vocab_size(&self) -> usize {
            self.vocab_size
        }

        fn eos_token(&self) -> u32 {
            self.eos_token
        }
    }

    // === TokenProb Tests ===

    #[test]
    fn test_token_prob_new() {
        let tp = TokenProb::new(42, -1.0);
        assert_eq!(tp.token, 42);
        assert_eq!(tp.log_prob, -1.0);
    }

    #[test]
    fn test_token_prob_prob() {
        let tp = TokenProb::new(42, 0.0);
        assert!((tp.prob() - 1.0).abs() < 0.001);

        let tp2 = TokenProb::new(42, -1.0);
        assert!((tp2.prob() - 0.368).abs() < 0.01);
    }

    // === SpeculativeStats Tests ===

    #[test]
    fn test_speculative_stats_default() {
        let stats = SpeculativeStats::default();
        assert_eq!(stats.iterations, 0);
        assert_eq!(stats.tokens_speculated, 0);
        assert_eq!(stats.acceptance_rate, 0.0);
    }

    #[test]
    fn test_speculative_stats_record() {
        let mut stats = SpeculativeStats::default();
        stats.record_iteration(4, 3, 1.0, 10.0);

        assert_eq!(stats.iterations, 1);
        assert_eq!(stats.tokens_speculated, 4);
        assert_eq!(stats.tokens_accepted, 3);
        assert_eq!(stats.acceptance_rate, 0.75);
    }

    #[test]
    fn test_speculative_stats_speedup() {
        let mut stats = SpeculativeStats::default();
        stats.record_iteration(4, 4, 1.0, 10.0);

        let speedup = stats.speedup();
        assert!(speedup > 1.0);
    }

    #[test]
    fn test_speculative_stats_serialization() {
        let stats = SpeculativeStats {
            iterations: 10,
            tokens_speculated: 40,
            tokens_accepted: 30,
            acceptance_rate: 0.75,
            avg_spec_length: 4.0,
            time_saved_ms: 100.0,
            draft_time_ms: 10.0,
            target_time_ms: 100.0,
        };

        let json = serde_json::to_string(&stats).expect("test");
        let parsed: SpeculativeStats = serde_json::from_str(&json).expect("test");

        assert_eq!(parsed.iterations, stats.iterations);
        assert_eq!(parsed.acceptance_rate, stats.acceptance_rate);
    }

    // === SpeculativeResult Tests ===

    #[test]
    fn test_speculative_result_acceptance_rate() {
        let result = SpeculativeResult {
            accepted_tokens: vec![1, 2, 3],
            num_speculated: 4,
            num_accepted: 3,
            resampled_token: Some(4),
            draft_time_ms: 1.0,
            target_time_ms: 10.0,
        };

        assert_eq!(result.acceptance_rate(), 0.75);
        assert!(!result.all_accepted());
    }

    #[test]
    fn test_speculative_result_all_accepted() {
        let result = SpeculativeResult {
            accepted_tokens: vec![1, 2, 3, 4],
            num_speculated: 4,
            num_accepted: 4,
            resampled_token: None,
            draft_time_ms: 1.0,
            target_time_ms: 10.0,
        };

        assert!(result.all_accepted());
    }

    // === SpeculativeConfig Tests ===

    #[test]
    fn test_speculative_config_default() {
        let config = SpeculativeConfig::default();
        assert_eq!(config.spec_length, 4);
        assert!(config.adaptive);
    }

    #[test]
    fn test_speculative_config_builder() {
        let config = SpeculativeConfig::new()
            .with_spec_length(6)
            .with_adaptive(false);

        assert_eq!(config.spec_length, 6);
        assert!(!config.adaptive);
    }

    #[test]
    fn test_speculative_config_adapt_increase() {
        let mut config = SpeculativeConfig::new().with_spec_length(4);
        config.adapt_spec_length(0.9); // High acceptance

        assert!(config.spec_length >= 4);
    }

    #[test]
    fn test_speculative_config_adapt_decrease() {
        let mut config = SpeculativeConfig::new().with_spec_length(4);
        config.adapt_spec_length(0.3); // Low acceptance

        assert!(config.spec_length <= 4);
    }

    #[test]
    fn test_speculative_config_no_adapt_when_disabled() {
        let mut config = SpeculativeConfig::new()
            .with_spec_length(4)
            .with_adaptive(false);

        config.adapt_spec_length(0.1);
        assert_eq!(config.spec_length, 4); // Should not change
    }

    // === SpeculativeDecoder Tests ===

    #[test]
    fn test_speculative_decoder_new() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("test");
        assert_eq!(decoder.spec_length(), 4);
    }

    #[test]
    fn test_speculative_decoder_invalid_spec_length() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let result = SpeculativeDecoder::new(draft, target, 0);
        assert!(matches!(
            result,
            Err(SpeculativeError::InvalidSpecLength(0))
        ));
    }

    #[test]
    fn test_speculative_decoder_decode_iteration() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test");
        let result = decoder.decode_iteration(&[10, 20, 30]).expect("test");

        assert!(!result.accepted_tokens.is_empty());
        assert!(result.num_speculated > 0);
    }

    #[test]
    fn test_speculative_decoder_set_spec_length() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test");
        decoder.set_spec_length(8).expect("test");
        assert_eq!(decoder.spec_length(), 8);

        let err = decoder.set_spec_length(0);
        assert!(err.is_err());
    }

    #[test]
    fn test_speculative_decoder_stats() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test");
        let _ = decoder.decode_iteration(&[10]).expect("test");

        let stats = decoder.stats();
        assert_eq!(stats.iterations, 1);
    }

    #[test]
    fn test_speculative_decoder_reset_stats() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test");
        let _ = decoder.decode_iteration(&[10]).expect("test");

        decoder.reset_stats();
        assert_eq!(decoder.stats().iterations, 0);
    }

    // === Error Tests ===

    #[test]
    fn test_speculative_error_display() {
        let err = SpeculativeError::DraftModelError("test".to_string());
        assert!(err.to_string().contains("Draft"));

        let err = SpeculativeError::TargetModelError("test".to_string());
        assert!(err.to_string().contains("Target"));

        let err = SpeculativeError::InvalidSpecLength(0);
        assert!(err.to_string().contains("0"));

        let err = SpeculativeError::VerificationFailed { position: 3 };
        assert!(err.to_string().contains("3"));
    }

    // ============================================================================
    // Additional Coverage Tests
    // ============================================================================

    // === TokenProb Extended Tests ===

    #[test]
    fn test_token_prob_clone() {
        let tp = TokenProb::new(42, -2.0);
        let tp_clone = tp.clone();
        assert_eq!(tp.token, tp_clone.token);
        assert_eq!(tp.log_prob, tp_clone.log_prob);
    }

    #[test]
    fn test_token_prob_debug() {
        let tp = TokenProb::new(42, -1.5);
        let debug_str = format!("{:?}", tp);
        assert!(debug_str.contains("42"));
        assert!(debug_str.contains("-1.5"));
    }

    #[test]
    fn test_token_prob_extreme_values() {
        // Very negative log prob (near-zero probability)
        let tp = TokenProb::new(1, -100.0);
        assert!(tp.prob() < 1e-40);

        // Positive log prob (probability > 1, edge case)
        let tp2 = TokenProb::new(1, 1.0);
        assert!((tp2.prob() - std::f32::consts::E).abs() < 0.01);
    }

    // === SpeculativeStats Extended Tests ===

    #[test]
    fn test_speculative_stats_speedup_zero_accepted() {
        let stats = SpeculativeStats::default();
        // tokens_accepted is 0, should return 1.0
        assert_eq!(stats.speedup(), 1.0);
    }

    #[test]
    fn test_speculative_stats_speedup_with_many_iterations() {
        let mut stats = SpeculativeStats::default();
        // Simulate many iterations
        for _ in 0..100 {
            stats.record_iteration(4, 3, 1.0, 10.0);
        }
        let speedup = stats.speedup();
        // With 400 speculated, 300 accepted, 100 iterations
        // baseline_time = 300, draft_equiv = 400 * 0.1 = 40, actual = 40 + 100 = 140
        // speedup = 300 / 140 ≈ 2.14
        assert!(speedup > 1.5);
        assert!(speedup < 3.0);
    }

    #[test]
    fn test_speculative_stats_record_zero_speculated() {
        let mut stats = SpeculativeStats::default();
        // Edge case: zero speculated tokens
        stats.record_iteration(0, 0, 0.5, 5.0);

        assert_eq!(stats.iterations, 1);
        assert_eq!(stats.tokens_speculated, 0);
        // acceptance_rate remains 0 when tokens_speculated is 0
        assert_eq!(stats.acceptance_rate, 0.0);
        // avg_spec_length = 0 / 1 = 0
        assert_eq!(stats.avg_spec_length, 0.0);
    }

    #[test]
    fn test_speculative_stats_time_saved_calculation() {
        let mut stats = SpeculativeStats::default();
        // 4 speculated, 4 accepted, target took 40ms
        // time_per_token = 40 / 4 = 10ms
        // time_saved = (4-1) * 10 = 30ms
        stats.record_iteration(4, 4, 1.0, 40.0);
        assert!((stats.time_saved_ms - 30.0).abs() < 0.01);
    }

    #[test]
    fn test_speculative_stats_time_saved_one_accepted() {
        let mut stats = SpeculativeStats::default();
        // Only 1 accepted: saturating_sub(1) = 0, no time saved
        stats.record_iteration(4, 1, 1.0, 40.0);
        assert_eq!(stats.time_saved_ms, 0.0);
    }

    #[test]
    fn test_speculative_stats_clone() {
        let stats = SpeculativeStats {
            iterations: 5,
            tokens_speculated: 20,
            tokens_accepted: 15,
            acceptance_rate: 0.75,
            avg_spec_length: 4.0,
            time_saved_ms: 50.0,
            draft_time_ms: 5.0,
            target_time_ms: 50.0,
        };
        let cloned = stats.clone();
        assert_eq!(stats.iterations, cloned.iterations);
        assert_eq!(stats.acceptance_rate, cloned.acceptance_rate);
    }

    // === SpeculativeResult Extended Tests ===

    #[test]
    fn test_speculative_result_acceptance_rate_zero_speculated() {
        let result = SpeculativeResult {
            accepted_tokens: vec![],
            num_speculated: 0,
            num_accepted: 0,
            resampled_token: None,
            draft_time_ms: 0.0,
            target_time_ms: 0.0,
        };
        // Should return 0.0 when num_speculated is 0
        assert_eq!(result.acceptance_rate(), 0.0);
    }

    #[test]
    fn test_speculative_result_clone() {
        let result = SpeculativeResult {
            accepted_tokens: vec![1, 2, 3],
            num_speculated: 4,
            num_accepted: 3,
            resampled_token: Some(5),
            draft_time_ms: 2.5,
            target_time_ms: 25.0,
        };
        let cloned = result.clone();
        assert_eq!(result.accepted_tokens, cloned.accepted_tokens);
        assert_eq!(result.resampled_token, cloned.resampled_token);
    }

    #[test]
    fn test_speculative_result_debug() {
        let result = SpeculativeResult {
            accepted_tokens: vec![1, 2],
            num_speculated: 3,
            num_accepted: 2,
            resampled_token: Some(4),
            draft_time_ms: 1.0,
            target_time_ms: 10.0,
        };
        let debug_str = format!("{:?}", result);
        assert!(debug_str.contains("accepted_tokens"));
        assert!(debug_str.contains("resampled_token"));
    }

    // === SpeculativeConfig Extended Tests ===

    #[test]
    fn test_speculative_config_adapt_at_max() {
        let mut config = SpeculativeConfig {
            spec_length: 8, // Already at max
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };
        config.adapt_spec_length(0.95); // Very high acceptance
                                        // Should not exceed max
        assert_eq!(config.spec_length, 8);
    }

    #[test]
    fn test_speculative_config_adapt_at_min() {
        let mut config = SpeculativeConfig {
            spec_length: 1, // Already at min
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };
        config.adapt_spec_length(0.1); // Very low acceptance
                                       // Should not go below 1
        assert_eq!(config.spec_length, 1);
    }

    #[test]
    fn test_speculative_config_adapt_medium_rate_no_change() {
        let mut config = SpeculativeConfig::new().with_spec_length(4);
        // 0.6 is above min_acceptance_rate (0.5) but below 0.8
        config.adapt_spec_length(0.6);
        // Should not change
        assert_eq!(config.spec_length, 4);
    }

    #[test]
    fn test_speculative_config_serialization() {
        let config = SpeculativeConfig {
            spec_length: 6,
            min_acceptance_rate: 0.4,
            adaptive: false,
            max_spec_length: 12,
        };
        let json = serde_json::to_string(&config).expect("serialize");
        let parsed: SpeculativeConfig = serde_json::from_str(&json).expect("deserialize");
        assert_eq!(parsed.spec_length, 6);
        assert_eq!(parsed.max_spec_length, 12);
        assert!(!parsed.adaptive);
    }

    #[test]
    fn test_speculative_config_clone() {
        let config = SpeculativeConfig::new()
            .with_spec_length(5)
            .with_adaptive(false);
        let cloned = config.clone();
        assert_eq!(config.spec_length, cloned.spec_length);
        assert_eq!(config.adaptive, cloned.adaptive);
    }

    #[test]
    fn test_speculative_config_debug() {
        let config = SpeculativeConfig::default();
        let debug_str = format!("{:?}", config);
        assert!(debug_str.contains("spec_length"));
        assert!(debug_str.contains("adaptive"));
    }

    // === SpeculativeDecoder Extended Tests ===

    #[test]
    fn test_speculative_decoder_spec_length_too_large() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let result = SpeculativeDecoder::new(draft, target, 33);
        assert!(matches!(
            result,
            Err(SpeculativeError::InvalidSpecLength(33))
        ));
    }

    #[test]
    fn test_speculative_decoder_set_spec_length_too_large() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.set_spec_length(33);
        assert!(matches!(
            result,
            Err(SpeculativeError::InvalidSpecLength(33))
        ));
        // Original value unchanged
        assert_eq!(decoder.spec_length(), 4);
    }

    #[test]
    fn test_speculative_decoder_boundary_spec_length() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        // Exactly 32 should work
        let decoder = SpeculativeDecoder::new(draft, target, 32);
        assert!(decoder.is_ok());
        assert_eq!(decoder.unwrap().spec_length(), 32);
    }

    #[test]
    fn test_speculative_decoder_boundary_spec_length_one() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        // Exactly 1 should work
        let decoder = SpeculativeDecoder::new(draft, target, 1);
        assert!(decoder.is_ok());
        assert_eq!(decoder.unwrap().spec_length(), 1);
    }

    /// Mock model that returns EOS token
    struct EosModel {
        vocab_size: usize,
        eos_token: u32,
    }

    impl EosModel {
        fn new(vocab_size: usize, eos_token: u32) -> Self {
            Self {
                vocab_size,
                eos_token,
            }
        }
    }

    impl SpeculativeModel for EosModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; self.vocab_size])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            // Always return EOS
            Ok(TokenProb::new(self.eos_token, -0.5))
        }

        fn vocab_size(&self) -> usize {
            self.vocab_size
        }

        fn eos_token(&self) -> u32 {
            self.eos_token
        }
    }

    #[test]
    fn test_speculative_decoder_eos_stops_draft() {
        let draft = EosModel::new(100, 0); // Returns EOS immediately
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[10]).expect("decode");

        // Should stop after first token (EOS)
        assert_eq!(result.num_speculated, 1);
    }

    /// Mock model that fails on forward
    struct FailingForwardModel;

    impl SpeculativeModel for FailingForwardModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Err(SpeculativeError::DraftModelError(
                "forward failed".to_string(),
            ))
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Ok(TokenProb::new(1, -1.0))
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_draft_forward_error() {
        let draft = FailingForwardModel;
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[10]);

        assert!(matches!(result, Err(SpeculativeError::DraftModelError(_))));
    }

    /// Mock model that fails on sample
    struct FailingSampleModel;

    impl SpeculativeModel for FailingSampleModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; 100])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Err(SpeculativeError::DraftModelError(
                "sample failed".to_string(),
            ))
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_draft_sample_error() {
        let draft = FailingSampleModel;
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[10]);

        assert!(matches!(result, Err(SpeculativeError::DraftModelError(_))));
    }

    /// Mock model that fails target forward
    struct FailingTargetModel;

    impl SpeculativeModel for FailingTargetModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Err(SpeculativeError::TargetModelError(
                "target forward failed".to_string(),
            ))
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Ok(TokenProb::new(1, -1.0))
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_target_forward_error() {
        let draft = MockModel::new(100, 1);
        let target = FailingTargetModel;

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[10]);

        assert!(matches!(result, Err(SpeculativeError::TargetModelError(_))));
    }

    /// Mock model that fails on target sample (after successful forward)
    struct FailingTargetSampleModel {
        sample_fail: bool,
    }

    impl SpeculativeModel for FailingTargetSampleModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; 100])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            if self.sample_fail {
                Err(SpeculativeError::TargetModelError(
                    "target sample failed".to_string(),
                ))
            } else {
                Ok(TokenProb::new(1, -1.0))
            }
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_target_sample_error() {
        let draft = MockModel::new(100, 1);
        let target = FailingTargetSampleModel { sample_fail: true };

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[10]);

        assert!(matches!(result, Err(SpeculativeError::TargetModelError(_))));
    }

    /// Mock model that returns different tokens (for rejection testing)
    struct DifferentTokenModel {
        token: u32,
        prob: f32,
    }

    impl DifferentTokenModel {
        fn new(token: u32, prob: f32) -> Self {
            Self { token, prob }
        }
    }

    impl SpeculativeModel for DifferentTokenModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; 100])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Ok(TokenProb::new(self.token, self.prob))
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_rejection_resamples() {
        // Draft returns token 5 with low probability
        let draft = DifferentTokenModel::new(5, -10.0); // Very low prob
                                                        // Target returns token 10 with high probability
        let target = DifferentTokenModel::new(10, -0.1); // High prob

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[1]).expect("decode");

        // Should have resampled token since tokens differ and ratio check may fail
        // The first token triggers the should_accept check
        assert!(!result.accepted_tokens.is_empty());
    }

    #[test]
    fn test_speculative_decoder_multiple_iterations_stats() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 2).expect("create");

        // Run multiple iterations
        for _ in 0..5 {
            let _ = decoder.decode_iteration(&[1, 2, 3]).expect("decode");
        }

        let stats = decoder.stats();
        assert_eq!(stats.iterations, 5);
        assert!(stats.tokens_speculated >= 5);
    }

    #[test]
    fn test_speculative_decoder_empty_context() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[]).expect("decode");

        // Should still work with empty context
        assert!(result.num_speculated > 0);
    }

    // === should_accept Edge Cases ===

    #[test]
    fn test_should_accept_same_token() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        let draft_prob = TokenProb::new(42, -1.0);
        let target_prob = TokenProb::new(42, -2.0); // Same token, different prob

        // Same tokens should always be accepted
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_high_target_prob() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Different tokens, but target has much higher probability
        let draft_prob = TokenProb::new(5, -10.0); // Very low
        let target_prob = TokenProb::new(10, -0.1); // High

        // Ratio = exp(-0.1) / exp(-10.0) ≈ 0.9 / 0.00005 >> 1.0
        // Should accept because ratio >= 1.0
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_similar_probs() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Different tokens with similar probabilities
        let draft_prob = TokenProb::new(5, -1.0);
        let target_prob = TokenProb::new(10, -1.0);

        // ratio = 1.0, should accept (ratio >= 1.0 is true)
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_moderate_ratio() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Different tokens with ratio between 0.5 and 1.0
        let draft_prob = TokenProb::new(5, -1.0); // prob ≈ 0.368
        let target_prob = TokenProb::new(10, -1.2); // prob ≈ 0.301

        // ratio = 0.301 / 0.368 ≈ 0.82, > 0.5, should accept
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_very_low_target_prob() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Target has much lower probability
        let draft_prob = TokenProb::new(5, -0.1); // High prob ≈ 0.9
        let target_prob = TokenProb::new(10, -5.0); // Low prob ≈ 0.0067

        // ratio = 0.0067 / 0.9 ≈ 0.007, < 0.5, should reject
        assert!(!decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_draft_near_zero() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Draft prob very close to zero (edge case for max(1e-10))
        let draft_prob = TokenProb::new(5, -100.0); // Extremely low
        let target_prob = TokenProb::new(10, -1.0);

        // ratio = 0.368 / max(tiny, 1e-10) = huge, should accept
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    // === Error Type Coverage ===

    #[test]
    fn test_speculative_error_debug() {
        let err = SpeculativeError::DraftModelError("test error".to_string());
        let debug_str = format!("{:?}", err);
        assert!(debug_str.contains("DraftModelError"));
        assert!(debug_str.contains("test error"));

        let err2 = SpeculativeError::VerificationFailed { position: 7 };
        let debug_str2 = format!("{:?}", err2);
        assert!(debug_str2.contains("VerificationFailed"));
        assert!(debug_str2.contains("7"));
    }

    // === MockModel Extended Tests ===

    #[test]
    fn test_mock_model_vocab_size() {
        let model = MockModel::new(256, 42);
        assert_eq!(model.vocab_size(), 256);
    }

    #[test]
    fn test_mock_model_eos_token() {
        let model = MockModel::new(100, 5);
        assert_eq!(model.eos_token(), 0);
    }

    #[test]
    fn test_mock_model_forward_returns_correct_size() {
        let model = MockModel::new(50, 1);
        let logits = model.forward(&[1, 2, 3]).expect("forward");
        assert_eq!(logits.len(), 50);
    }

    // === SpeculativeResult all_accepted Edge Cases ===

    #[test]
    fn test_speculative_result_all_accepted_zero() {
        let result = SpeculativeResult {
            accepted_tokens: vec![],
            num_speculated: 0,
            num_accepted: 0,
            resampled_token: None,
            draft_time_ms: 0.0,
            target_time_ms: 0.0,
        };
        // 0 == 0, so all_accepted is true
        assert!(result.all_accepted());
    }

    // === Config adapt_spec_length boundary tests ===

    #[test]
    fn test_speculative_config_adapt_exactly_at_threshold() {
        let mut config = SpeculativeConfig {
            spec_length: 4,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };

        // Exactly at min_acceptance_rate - should not decrease
        config.adapt_spec_length(0.5);
        assert_eq!(config.spec_length, 4);

        // Exactly at 0.8 - should not increase (> 0.8 is required)
        config.adapt_spec_length(0.8);
        assert_eq!(config.spec_length, 4);
    }

    #[test]
    fn test_speculative_config_adapt_just_above_high_threshold() {
        let mut config = SpeculativeConfig {
            spec_length: 4,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };

        // Just above 0.8 - should increase
        config.adapt_spec_length(0.81);
        assert_eq!(config.spec_length, 5);
    }

    #[test]
    fn test_speculative_config_adapt_just_below_low_threshold() {
        let mut config = SpeculativeConfig {
            spec_length: 4,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };

        // Just below min_acceptance_rate - should decrease
        config.adapt_spec_length(0.49);
        assert_eq!(config.spec_length, 3);
    }

    // ============================================================================
    // Part 02: Additional Coverage Tests for Edge Cases
    // ============================================================================

    // === SpeculativeStats Debug Trait ===

    #[test]
    fn test_speculative_stats_debug() {
        let stats = SpeculativeStats {
            iterations: 10,
            tokens_speculated: 40,
            tokens_accepted: 32,
            acceptance_rate: 0.8,
            avg_spec_length: 4.0,
            time_saved_ms: 120.5,
            draft_time_ms: 15.0,
            target_time_ms: 150.0,
        };
        let debug_str = format!("{:?}", stats);
        assert!(debug_str.contains("iterations"));
        assert!(debug_str.contains("tokens_speculated"));
        assert!(debug_str.contains("acceptance_rate"));
        assert!(debug_str.contains("time_saved_ms"));
    }

    // === SpeculativeStats Edge Cases ===

    #[test]
    fn test_speculative_stats_speedup_zero_iterations_zero_speculated() {
        // Edge case: manually created stats with no activity
        let stats = SpeculativeStats {
            iterations: 0,
            tokens_speculated: 0,
            tokens_accepted: 0,
            acceptance_rate: 0.0,
            avg_spec_length: 0.0,
            time_saved_ms: 0.0,
            draft_time_ms: 0.0,
            target_time_ms: 0.0,
        };
        // tokens_accepted == 0 means early return with 1.0
        assert_eq!(stats.speedup(), 1.0);
    }

    #[test]
    fn test_speculative_stats_speedup_actual_time_calculation() {
        // Test the actual speedup calculation path
        let mut stats = SpeculativeStats::default();
        // 10 speculated, 8 accepted, 1 iteration
        // draft_tokens_equivalent = 10 * 0.1 = 1.0
        // baseline_time = 8
        // actual_time = 1.0 + 1 = 2.0
        // speedup = 8 / 2 = 4.0
        stats.record_iteration(10, 8, 1.0, 10.0);
        let speedup = stats.speedup();
        assert!((speedup - 4.0).abs() < 0.1);
    }

    #[test]
    fn test_speculative_stats_multiple_records_accumulate() {
        let mut stats = SpeculativeStats::default();
        stats.record_iteration(4, 3, 1.0, 10.0);
        stats.record_iteration(4, 4, 1.5, 12.0);
        stats.record_iteration(4, 2, 0.8, 8.0);

        assert_eq!(stats.iterations, 3);
        assert_eq!(stats.tokens_speculated, 12);
        assert_eq!(stats.tokens_accepted, 9);
        assert!((stats.draft_time_ms - 3.3).abs() < 0.01);
        assert!((stats.target_time_ms - 30.0).abs() < 0.01);
    }

    #[test]
    fn test_speculative_stats_acceptance_rate_precision() {
        let mut stats = SpeculativeStats::default();
        stats.record_iteration(3, 2, 1.0, 10.0);
        // 2/3 = 0.6666...
        assert!((stats.acceptance_rate - 0.6666667).abs() < 0.0001);
    }

    // === SpeculativeResult Edge Cases ===

    #[test]
    fn test_speculative_result_partial_acceptance() {
        let result = SpeculativeResult {
            accepted_tokens: vec![100, 200],
            num_speculated: 5,
            num_accepted: 2,
            resampled_token: Some(300),
            draft_time_ms: 2.5,
            target_time_ms: 25.0,
        };
        // 2/5 = 0.4
        assert!((result.acceptance_rate() - 0.4).abs() < 0.001);
        assert!(!result.all_accepted());
    }

    #[test]
    fn test_speculative_result_100_percent_acceptance() {
        let result = SpeculativeResult {
            accepted_tokens: vec![10, 20, 30, 40, 50],
            num_speculated: 5,
            num_accepted: 5,
            resampled_token: None,
            draft_time_ms: 1.0,
            target_time_ms: 10.0,
        };
        assert!((result.acceptance_rate() - 1.0).abs() < 0.001);
        assert!(result.all_accepted());
    }

    // === TokenProb Edge Cases ===

    #[test]
    fn test_token_prob_zero_log_prob() {
        // log_prob = 0 means prob = exp(0) = 1.0
        let tp = TokenProb::new(999, 0.0);
        assert!((tp.prob() - 1.0).abs() < 0.0001);
    }

    #[test]
    fn test_token_prob_very_positive_log_prob() {
        // Edge case: positive log_prob (prob > 1, mathematically unusual but handled)
        let tp = TokenProb::new(1, 2.0);
        let expected = std::f32::consts::E.powi(2);
        assert!((tp.prob() - expected).abs() < 0.1);
    }

    #[test]
    fn test_token_prob_negative_infinity_log_prob() {
        // Very negative log_prob approximates zero probability
        let tp = TokenProb::new(1, f32::NEG_INFINITY);
        assert_eq!(tp.prob(), 0.0);
    }

    // === SpeculativeConfig Edge Cases ===

    #[test]
    fn test_speculative_config_adapt_spec_length_two_decreases() {
        let mut config = SpeculativeConfig {
            spec_length: 3,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };
        // First decrease
        config.adapt_spec_length(0.3);
        assert_eq!(config.spec_length, 2);
        // Second decrease
        config.adapt_spec_length(0.3);
        assert_eq!(config.spec_length, 1);
        // Third attempt - should stay at 1 (minimum)
        config.adapt_spec_length(0.3);
        assert_eq!(config.spec_length, 1);
    }

    #[test]
    fn test_speculative_config_adapt_spec_length_two_increases() {
        let mut config = SpeculativeConfig {
            spec_length: 6,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 8,
        };
        // First increase
        config.adapt_spec_length(0.9);
        assert_eq!(config.spec_length, 7);
        // Second increase
        config.adapt_spec_length(0.9);
        assert_eq!(config.spec_length, 8);
        // Third attempt - should stay at 8 (maximum)
        config.adapt_spec_length(0.9);
        assert_eq!(config.spec_length, 8);
    }

    #[test]
    fn test_speculative_config_custom_min_acceptance_rate() {
        let mut config = SpeculativeConfig {
            spec_length: 4,
            min_acceptance_rate: 0.7, // Custom higher threshold
            adaptive: true,
            max_spec_length: 8,
        };
        // 0.65 is below 0.7, should decrease
        config.adapt_spec_length(0.65);
        assert_eq!(config.spec_length, 3);
    }

    // === SpeculativeDecoder Edge Cases ===

    /// Model that generates multiple tokens then hits EOS
    struct DelayedEosModel {
        vocab_size: usize,
        eos_token: u32,
        call_count: std::cell::RefCell<usize>,
        eos_after: usize,
    }

    impl DelayedEosModel {
        fn new(vocab_size: usize, eos_token: u32, eos_after: usize) -> Self {
            Self {
                vocab_size,
                eos_token,
                call_count: std::cell::RefCell::new(0),
                eos_after,
            }
        }
    }

    impl SpeculativeModel for DelayedEosModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; self.vocab_size])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            let mut count = self.call_count.borrow_mut();
            *count += 1;
            if *count > self.eos_after {
                Ok(TokenProb::new(self.eos_token, -0.5))
            } else {
                Ok(TokenProb::new(42, -1.0)) // Non-EOS token
            }
        }

        fn vocab_size(&self) -> usize {
            self.vocab_size
        }

        fn eos_token(&self) -> u32 {
            self.eos_token
        }
    }

    #[test]
    fn test_speculative_decoder_eos_after_two_tokens() {
        let draft = DelayedEosModel::new(100, 0, 2); // EOS after 2 tokens
        let target = MockModel::new(100, 42);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[1]).expect("decode");

        // Should have 3 tokens: 2 non-EOS + 1 EOS
        assert_eq!(result.num_speculated, 3);
    }

    #[test]
    fn test_speculative_decoder_large_context() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 2).expect("create");
        // Large context of 1000 tokens
        let context: Vec<u32> = (0..1000).collect();
        let result = decoder.decode_iteration(&context).expect("decode");

        // Should still work with large context
        assert!(result.num_speculated > 0);
    }

    #[test]
    fn test_speculative_decoder_single_token_spec_length() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 1).expect("create");
        let result = decoder.decode_iteration(&[10, 20]).expect("decode");

        // With spec_length=1, should only speculate 1 token
        assert_eq!(result.num_speculated, 1);
    }

    #[test]
    fn test_speculative_decoder_max_spec_length_32() {
        // Model that doesn't hit EOS
        struct NoEosModel {
            vocab_size: usize,
        }

        impl SpeculativeModel for NoEosModel {
            fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
                Ok(vec![0.0; self.vocab_size])
            }

            fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
                Ok(TokenProb::new(999, -1.0)) // Never EOS (EOS is 0)
            }

            fn vocab_size(&self) -> usize {
                self.vocab_size
            }

            fn eos_token(&self) -> u32 {
                0
            }
        }

        let draft = NoEosModel { vocab_size: 100 };
        let target = NoEosModel { vocab_size: 100 };

        let mut decoder = SpeculativeDecoder::new(draft, target, 32).expect("create");
        let result = decoder.decode_iteration(&[1]).expect("decode");

        // Should speculate exactly 32 tokens (max spec_length)
        assert_eq!(result.num_speculated, 32);
    }

    // === should_accept Comprehensive Tests ===

    #[test]
    fn test_should_accept_ratio_exactly_one() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Same probability, different tokens
        let draft_prob = TokenProb::new(5, -2.0);
        let target_prob = TokenProb::new(10, -2.0);

        // ratio = exp(-2) / exp(-2) = 1.0, should accept (ratio >= 1.0)
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_ratio_slightly_above_half() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Calculate log probs to get ratio just above 0.5
        // draft_prob = 0.6, target_prob = 0.31 => ratio ≈ 0.517
        let _draft_prob = TokenProb::new(5, (-0.6_f32).ln());
        let _target_prob = TokenProb::new(10, (-0.31_f32).ln());

        // This is a negative log prob scenario - let's use cleaner values
        // draft prob = exp(-1) ≈ 0.368, target prob = exp(-1.3) ≈ 0.273
        // ratio = 0.273 / 0.368 ≈ 0.74 > 0.5, should accept
        let draft_prob = TokenProb::new(5, -1.0);
        let target_prob = TokenProb::new(10, -1.3);
        assert!(decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_ratio_exactly_half() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // We need ratio = 0.5 exactly
        // target_prob / draft_prob = 0.5
        // exp(target_log) / exp(draft_log) = 0.5
        // exp(target_log - draft_log) = 0.5
        // target_log - draft_log = ln(0.5) ≈ -0.693
        // If draft_log = -1.0, target_log = -1.693
        let draft_prob = TokenProb::new(5, -1.0);
        let target_prob = TokenProb::new(10, -1.0 + (-0.5_f32).ln());

        // ratio = 0.5, condition is ratio > 0.5, so should NOT accept
        // But ratio >= 1.0 OR ratio > 0.5, so 0.5 exactly should NOT match > 0.5
        assert!(!decoder.should_accept(&draft_prob, &target_prob));
    }

    #[test]
    fn test_should_accept_ratio_below_half() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);
        let decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // ratio needs to be < 0.5
        // target_prob / draft_prob < 0.5
        // If draft is high and target is low
        let draft_prob = TokenProb::new(5, 0.0); // prob = 1.0
        let target_prob = TokenProb::new(10, -2.0); // prob ≈ 0.135

        // ratio = 0.135 / 1.0 = 0.135 < 0.5, should NOT accept
        assert!(!decoder.should_accept(&draft_prob, &target_prob));
    }

    // === Error Handling Edge Cases ===

    #[test]
    fn test_speculative_error_verification_failed_fields() {
        let err = SpeculativeError::VerificationFailed { position: 42 };
        if let SpeculativeError::VerificationFailed { position } = err {
            assert_eq!(position, 42);
        } else {
            panic!("Expected VerificationFailed variant");
        }
    }

    #[test]
    fn test_speculative_error_draft_model_error_message() {
        let err = SpeculativeError::DraftModelError("custom error message".to_string());
        let display = format!("{}", err);
        assert!(display.contains("Draft model error"));
        assert!(display.contains("custom error message"));
    }

    #[test]
    fn test_speculative_error_target_model_error_message() {
        let err = SpeculativeError::TargetModelError("target failed".to_string());
        let display = format!("{}", err);
        assert!(display.contains("Target model error"));
        assert!(display.contains("target failed"));
    }

    #[test]
    fn test_speculative_error_invalid_spec_length_message() {
        let err = SpeculativeError::InvalidSpecLength(100);
        let display = format!("{}", err);
        assert!(display.contains("Invalid speculation length"));
        assert!(display.contains("100"));
    }

    // === Decoder with Rejection Scenarios ===

    /// Model where draft and target always disagree with very low target probability
    struct AlwaysRejectModel {
        token: u32,
        log_prob: f32,
    }

    impl SpeculativeModel for AlwaysRejectModel {
        fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> {
            Ok(vec![0.0; 100])
        }

        fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> {
            Ok(TokenProb::new(self.token, self.log_prob))
        }

        fn vocab_size(&self) -> usize {
            100
        }

        fn eos_token(&self) -> u32 {
            0
        }
    }

    #[test]
    fn test_speculative_decoder_immediate_rejection() {
        // Draft has token 5 with high prob, target has token 10 with very low prob
        let draft = AlwaysRejectModel {
            token: 5,
            log_prob: 0.0, // prob = 1.0
        };
        let target = AlwaysRejectModel {
            token: 10,
            log_prob: -10.0, // prob ≈ 0.00005
        };

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");
        let result = decoder.decode_iteration(&[1]).expect("decode");

        // First token should be rejected and resampled
        assert!(result.resampled_token.is_some());
        // Only 1 token accepted (the resampled one)
        assert_eq!(result.num_accepted, 1);
    }

    // === Stats Accumulation Verification ===

    #[test]
    fn test_speculative_stats_time_accumulation() {
        let mut stats = SpeculativeStats::default();

        stats.record_iteration(4, 3, 1.5, 15.0);
        assert!((stats.draft_time_ms - 1.5).abs() < 0.001);
        assert!((stats.target_time_ms - 15.0).abs() < 0.001);

        stats.record_iteration(4, 4, 2.0, 20.0);
        assert!((stats.draft_time_ms - 3.5).abs() < 0.001);
        assert!((stats.target_time_ms - 35.0).abs() < 0.001);
    }

    #[test]
    fn test_speculative_stats_avg_spec_length_calculation() {
        let mut stats = SpeculativeStats::default();

        stats.record_iteration(2, 2, 1.0, 10.0);
        assert!((stats.avg_spec_length - 2.0).abs() < 0.001);

        stats.record_iteration(6, 4, 1.0, 10.0);
        // Total speculated = 8, iterations = 2, avg = 4.0
        assert!((stats.avg_spec_length - 4.0).abs() < 0.001);

        stats.record_iteration(4, 3, 1.0, 10.0);
        // Total speculated = 12, iterations = 3, avg = 4.0
        assert!((stats.avg_spec_length - 4.0).abs() < 0.001);
    }

    // === Config with Custom max_spec_length ===

    #[test]
    fn test_speculative_config_custom_max_spec_length() {
        let mut config = SpeculativeConfig {
            spec_length: 10,
            min_acceptance_rate: 0.5,
            adaptive: true,
            max_spec_length: 12,
        };

        config.adapt_spec_length(0.95);
        assert_eq!(config.spec_length, 11);

        config.adapt_spec_length(0.95);
        assert_eq!(config.spec_length, 12);

        config.adapt_spec_length(0.95);
        assert_eq!(config.spec_length, 12); // Capped at max
    }

    // === Deserialization Tests ===

    #[test]
    fn test_speculative_stats_deserialize_from_json() {
        let json = r#"{
            "iterations": 5,
            "tokens_speculated": 20,
            "tokens_accepted": 15,
            "acceptance_rate": 0.75,
            "avg_spec_length": 4.0,
            "time_saved_ms": 50.0,
            "draft_time_ms": 5.0,
            "target_time_ms": 50.0
        }"#;

        let stats: SpeculativeStats = serde_json::from_str(json).expect("deserialize");
        assert_eq!(stats.iterations, 5);
        assert_eq!(stats.tokens_speculated, 20);
        assert_eq!(stats.tokens_accepted, 15);
        assert!((stats.acceptance_rate - 0.75).abs() < 0.001);
    }

    #[test]
    fn test_speculative_config_deserialize_from_json() {
        let json = r#"{
            "spec_length": 8,
            "min_acceptance_rate": 0.6,
            "adaptive": false,
            "max_spec_length": 16
        }"#;

        let config: SpeculativeConfig = serde_json::from_str(json).expect("deserialize");
        assert_eq!(config.spec_length, 8);
        assert!((config.min_acceptance_rate - 0.6).abs() < 0.001);
        assert!(!config.adaptive);
        assert_eq!(config.max_spec_length, 16);
    }

    // === Decoder Stats Integration ===

    #[test]
    fn test_speculative_decoder_stats_after_multiple_iterations() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 3).expect("create");

        for _ in 0..10 {
            let _ = decoder.decode_iteration(&[1, 2, 3]).expect("decode");
        }

        let stats = decoder.stats();
        assert_eq!(stats.iterations, 10);
        assert!(stats.tokens_speculated >= 10);
        assert!(stats.draft_time_ms >= 0.0);
        assert!(stats.target_time_ms >= 0.0);
    }

    #[test]
    fn test_speculative_decoder_reset_clears_all_stats() {
        let draft = MockModel::new(100, 1);
        let target = MockModel::new(100, 1);

        let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("create");

        // Run some iterations
        for _ in 0..5 {
            let _ = decoder.decode_iteration(&[1]).expect("decode");
        }

        // Reset
        decoder.reset_stats();

        let stats = decoder.stats();
        assert_eq!(stats.iterations, 0);
        assert_eq!(stats.tokens_speculated, 0);
        assert_eq!(stats.tokens_accepted, 0);
        assert_eq!(stats.acceptance_rate, 0.0);
        assert_eq!(stats.avg_spec_length, 0.0);
        assert_eq!(stats.time_saved_ms, 0.0);
        assert_eq!(stats.draft_time_ms, 0.0);
        assert_eq!(stats.target_time_ms, 0.0);
    }
}
