//! Lookahead Decoding (Fu et al. 2024, arXiv:2402.02057)
//!
//! Accelerates autoregressive LLM decoding via Jacobi iteration: instead of
//! generating one token per forward pass the decoder speculatively generates a
//! "lookahead window" of W tokens in parallel using an n-gram cache, then
//! verifies the best continuation against a scoring function that simulates a
//! single model forward pass.  Accepted tokens advance the sequence without
//! additional forward passes.
//!
//! # Quick start
//!
//! ```
//! use rtx_inference::lookahead::{LookaheadDecoder, LookaheadConfig};
//!
//! let config = LookaheadConfig::default();
//! let mut decoder = LookaheadDecoder::new(config);
//!
//! // Warm the cache with some known text (token ids)
//! decoder.warm_cache(&[1, 2, 3, 4, 5, 6]);
//!
//! // score_fn: given the current token sequence, return the single most-likely
//! // next token (greedy — simulates one model forward pass).
//! let score_fn = |tokens: &[u32]| -> u32 {
//!     tokens.last().copied().unwrap_or(0).wrapping_add(1)
//! };
//!
//! let generated = decoder.decode(&[1, 2, 3], 10, score_fn);
//! assert!(!generated.is_empty());
//! ```

use std::collections::{HashMap, VecDeque};

// ---------------------------------------------------------------------------
// NGram
// ---------------------------------------------------------------------------

/// A single n-gram entry: a prefix of `n-1` tokens paired with the observed
/// continuation token.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct NGram {
    /// The `n-1` token prefix (context).
    pub prefix: Vec<u32>,
    /// The observed next token.
    pub next: u32,
}

// ---------------------------------------------------------------------------
// NGramCache
// ---------------------------------------------------------------------------

/// Cache of n-gram statistics accumulated during or before generation.
///
/// Internally maps each observed prefix to a frequency table of continuation
/// tokens.  When the total number of distinct prefixes exceeds `max_size` the
/// oldest prefix (FIFO) is evicted.
#[derive(Debug, Clone)]
pub struct NGramCache {
    /// The *n* in n-gram: prefix length is `n-1`, context window is `n`.
    n: usize,
    /// prefix → { next_token → count }
    table: HashMap<Vec<u32>, HashMap<u32, usize>>,
    /// Maximum number of distinct prefixes to store before eviction.
    max_size: usize,
    /// FIFO insertion order for eviction.
    insertion_order: VecDeque<Vec<u32>>,
    /// Running total of all individual (prefix, next) observations.
    total_entries: usize,
}

impl NGramCache {
    /// Create a new empty cache.
    ///
    /// # Arguments
    /// * `n` — n-gram order (must be ≥ 2; `n == 1` means prefix length 0,
    ///   which degenerates to a unigram and is not useful here).
    /// * `max_size` — maximum number of distinct prefixes before FIFO eviction.
    #[must_use]
    pub fn new(n: usize, max_size: usize) -> Self {
        let n = n.max(2); // guard: n must be at least 2 so prefix length ≥ 1
        Self {
            n,
            table: HashMap::new(),
            max_size: max_size.max(1),
            insertion_order: VecDeque::new(),
            total_entries: 0,
        }
    }

    /// Slide a window of size `n` over `tokens`, recording every (prefix, next)
    /// pair.  Does nothing if `tokens.len() < n`.
    pub fn observe(&mut self, tokens: &[u32]) {
        let prefix_len = self.n - 1;
        if tokens.len() < self.n {
            return;
        }
        for i in 0..=(tokens.len() - self.n) {
            let prefix: Vec<u32> = tokens[i..i + prefix_len].to_vec();
            let next = tokens[i + prefix_len];
            self.insert(prefix, next);
        }
    }

    /// Return the top-`k` candidate next tokens for `prefix`, sorted by
    /// observed count descending.  Returns an empty `Vec` if the prefix is
    /// not in the cache or `prefix.len() != n-1`.
    #[must_use]
    pub fn top_candidates(&self, prefix: &[u32], k: usize) -> Vec<u32> {
        if k == 0 {
            return Vec::new();
        }
        let counts = match self.table.get(prefix) {
            Some(m) => m,
            None => return Vec::new(),
        };
        let mut pairs: Vec<(u32, usize)> = counts.iter().map(|(&t, &c)| (t, c)).collect();
        // Stable sort descending by count, ties broken by token id (determinism).
        pairs.sort_unstable_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
        pairs.into_iter().take(k).map(|(t, _)| t).collect()
    }

    /// Number of distinct prefixes currently stored.
    #[must_use]
    pub fn num_prefixes(&self) -> usize {
        self.table.len()
    }

    /// Total number of individual (prefix, next_token) observations recorded
    /// (sum of all counts across all entries).
    #[must_use]
    pub fn total_observations(&self) -> usize {
        self.total_entries
    }

    // ------------------------------------------------------------------
    // Private helpers
    // ------------------------------------------------------------------

    fn insert(&mut self, prefix: Vec<u32>, next: u32) {
        // If prefix is new and we're at capacity, evict the oldest.
        if !self.table.contains_key(&prefix) {
            if self.table.len() >= self.max_size {
                self.evict_oldest();
            }
            self.table.insert(prefix.clone(), HashMap::new());
            self.insertion_order.push_back(prefix.clone());
        }
        // Increment count.
        let count = self
            .table
            .get_mut(&prefix)
            .expect("just inserted above; prefix must exist")
            .entry(next)
            .or_insert(0);
        *count += 1;
        self.total_entries += 1;
    }

    fn evict_oldest(&mut self) {
        if let Some(oldest) = self.insertion_order.pop_front() {
            if let Some(inner) = self.table.remove(&oldest) {
                // Deduct evicted counts from the running total.
                let removed: usize = inner.values().sum();
                self.total_entries = self.total_entries.saturating_sub(removed);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// LookaheadConfig
// ---------------------------------------------------------------------------

/// Configuration for [`LookaheadDecoder`].
#[derive(Debug, Clone)]
pub struct LookaheadConfig {
    /// `W`: number of speculative tokens to draft per step.
    pub window_size: usize,
    /// `n` in n-gram; prefix length is `n-1`.
    pub ngram_order: usize,
    /// Maximum number of distinct prefixes in the n-gram cache.
    pub cache_size: usize,
    /// Top-k n-gram candidates to consider at each draft position.
    pub max_candidates_per_step: usize,
    /// Token id that signals end-of-sequence.
    pub eos_token_id: u32,
    /// Hard cap on the total output length (initial + generated).
    pub max_length: usize,
}

impl Default for LookaheadConfig {
    fn default() -> Self {
        Self {
            window_size: 5,
            ngram_order: 3,
            cache_size: 10_000,
            max_candidates_per_step: 3,
            eos_token_id: 2,
            max_length: 256,
        }
    }
}

// ---------------------------------------------------------------------------
// LookaheadStats
// ---------------------------------------------------------------------------

/// Cumulative statistics collected during [`LookaheadDecoder::decode`].
#[derive(Debug, Clone, Default)]
pub struct LookaheadStats {
    /// Total tokens appended to the output sequence.
    pub total_tokens_generated: usize,
    /// Number of lookahead iterations (decode steps).
    pub total_steps: usize,
    /// Tokens accepted directly from n-gram cache predictions.
    pub total_accepted_from_ngram: usize,
    /// Tokens accepted from the model (either as correction or no-draft path).
    pub total_accepted_from_model: usize,
    /// Average tokens accepted per step (`total_tokens_generated / total_steps`).
    pub avg_tokens_per_step: f32,
}

// ---------------------------------------------------------------------------
// LookaheadDecoder
// ---------------------------------------------------------------------------

/// CPU-reference lookahead decoder implementing the Jacobi-iteration algorithm
/// from Fu et al. 2024 (arXiv:2402.02057).
///
/// This implementation runs entirely on the CPU and uses a caller-supplied
/// `score_fn` to simulate a single model forward pass (greedy: returns the
/// single most-likely next token).  It is suitable for unit testing, benchmarking
/// the acceptance logic, and as a reference implementation for GPU ports.
///
/// # Algorithm overview (per step)
///
/// 1. **Lookahead phase** — draft up to `window_size` tokens from the n-gram
///    cache by repeatedly looking up the most-likely next token for the current
///    prefix.
/// 2. **Verification phase** — run `score_fn` greedily over the draft,
///    accepting tokens while they match and stopping at the first mismatch
///    (accepting the model's correction).
/// 3. **Cache update** — observe the newly appended tokens to warm the cache.
/// 4. **EOS check** — stop if the accepted tail contains the EOS token id.
pub struct LookaheadDecoder {
    config: LookaheadConfig,
    cache: NGramCache,
    stats: LookaheadStats,
}

impl LookaheadDecoder {
    /// Create a new decoder with the given configuration.
    #[must_use]
    pub fn new(config: LookaheadConfig) -> Self {
        let cache = NGramCache::new(config.ngram_order, config.cache_size);
        Self {
            config,
            cache,
            stats: LookaheadStats::default(),
        }
    }

    /// Create a decoder with default configuration.
    #[must_use]
    pub fn with_defaults() -> Self {
        Self::new(LookaheadConfig::default())
    }

    /// Pre-load n-grams from a reference corpus to warm-start the cache.
    ///
    /// Call this before [`decode`](Self::decode) when you have representative
    /// text available (e.g., a system prompt or previous conversation turns).
    pub fn warm_cache(&mut self, corpus: &[u32]) {
        self.cache.observe(corpus);
    }

    /// Run lookahead decoding.
    ///
    /// # Arguments
    /// * `initial_tokens` — the prompt token ids (not included in the return value).
    /// * `max_new_tokens` — maximum number of *new* tokens to generate.
    /// * `score_fn` — closure that accepts the full current token sequence and
    ///   returns the single greedy-best next token.  This simulates one model
    ///   forward pass.
    ///
    /// # Returns
    /// The newly generated tokens (does not include `initial_tokens`).
    pub fn decode<F>(
        &mut self,
        initial_tokens: &[u32],
        max_new_tokens: usize,
        score_fn: F,
    ) -> Vec<u32>
    where
        F: Fn(&[u32]) -> u32,
    {
        let mut tokens: Vec<u32> = initial_tokens.to_vec();
        let mut generated: Vec<u32> = Vec::new();

        while generated.len() < max_new_tokens {
            let remaining = max_new_tokens - generated.len();

            // ----------------------------------------------------------
            // 1. Lookahead phase: draft up to window_size tokens from cache
            // ----------------------------------------------------------
            let mut draft: Vec<u32> = Vec::with_capacity(self.config.window_size);
            for _ in 0..self.config.window_size.min(remaining) {
                // Build the prefix from the tail of (tokens + draft so far)
                let combined_len = tokens.len() + draft.len();
                let prefix_len = self.config.ngram_order - 1; // n-1

                let prefix: Vec<u32> = if combined_len >= prefix_len {
                    // Take last prefix_len tokens from the combined sequence.
                    let start = combined_len - prefix_len;
                    if draft.len() >= prefix_len {
                        // The prefix falls entirely within draft.
                        draft[draft.len() - prefix_len..].to_vec()
                    } else {
                        // The prefix spans the end of tokens and the start of draft.
                        let from_tokens = prefix_len - draft.len();
                        let token_start = tokens.len().saturating_sub(from_tokens);
                        tokens[token_start..]
                            .iter()
                            .chain(draft.iter())
                            .cloned()
                            .collect()
                    }
                } else {
                    // Not enough history yet — build what we can.
                    tokens.iter().chain(draft.iter()).cloned().collect()
                };
                let _ = combined_len; // silence unused warning

                let candidates = self
                    .cache
                    .top_candidates(&prefix, self.config.max_candidates_per_step);
                if candidates.is_empty() {
                    break; // no n-gram available — stop drafting
                }
                draft.push(candidates[0]);
            }

            // ----------------------------------------------------------
            // 2. Verification phase
            // ----------------------------------------------------------
            let mut accepted_this_step: Vec<u32> = Vec::new();

            if !draft.is_empty() {
                let mut verify_ctx = tokens.clone();
                for &draft_tok in &draft {
                    let model_tok = score_fn(&verify_ctx);
                    if model_tok == draft_tok {
                        accepted_this_step.push(draft_tok);
                        verify_ctx.push(draft_tok);
                        self.stats.total_accepted_from_ngram += 1;
                    } else {
                        // Accept model's correction and stop verifying.
                        accepted_this_step.push(model_tok);
                        self.stats.total_accepted_from_model += 1;
                        break;
                    }
                    // Don't exceed max_new_tokens in a single step either.
                    if accepted_this_step.len() >= remaining {
                        break;
                    }
                }
            } else {
                // No draft available: just run the model once.
                let model_tok = score_fn(&tokens);
                accepted_this_step.push(model_tok);
                self.stats.total_accepted_from_model += 1;
            }

            // ----------------------------------------------------------
            // 3. Update n-gram cache with newly accepted tokens
            // ----------------------------------------------------------
            {
                // Build a small window: last (n-1) tokens from the current
                // sequence followed by the accepted tokens.
                let prefix_len = self.config.ngram_order.saturating_sub(1);
                let context_start = tokens.len().saturating_sub(prefix_len);
                let new_tokens: Vec<u32> = tokens[context_start..]
                    .iter()
                    .chain(accepted_this_step.iter())
                    .cloned()
                    .collect();
                self.cache.observe(&new_tokens);
            }

            // ----------------------------------------------------------
            // 4. Extend sequence
            // ----------------------------------------------------------
            let to_add = accepted_this_step.len().min(remaining);
            let accepted_slice = &accepted_this_step[..to_add];

            tokens.extend_from_slice(accepted_slice);
            generated.extend_from_slice(accepted_slice);
            self.stats.total_tokens_generated += accepted_slice.len();
            self.stats.total_steps += 1;

            // ----------------------------------------------------------
            // 5. EOS check
            // ----------------------------------------------------------
            if accepted_slice.last() == Some(&self.config.eos_token_id) {
                break;
            }
        }

        // Update derived statistic.
        self.stats.avg_tokens_per_step =
            self.stats.total_tokens_generated as f32 / self.stats.total_steps.max(1) as f32;

        generated
    }

    /// Access accumulated decoding statistics.
    #[must_use]
    pub fn stats(&self) -> &LookaheadStats {
        &self.stats
    }

    /// Reset all statistics counters to zero (cache is retained).
    pub fn reset_stats(&mut self) {
        self.stats = LookaheadStats::default();
    }

    /// Access the underlying n-gram cache.
    #[must_use]
    pub fn cache(&self) -> &NGramCache {
        &self.cache
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // Helper: score_fn that always predicts last_token + 1 (greedy incrementer).
    fn incrementer(tokens: &[u32]) -> u32 {
        tokens.last().copied().unwrap_or(0).wrapping_add(1)
    }

    // ------------------------------------------------------------------
    // NGramCache unit tests
    // ------------------------------------------------------------------

    #[test]
    fn test_ngram_cache_observe_single() {
        // observe [1,2,3] with n=3: prefix=[1,2], next=3
        let mut cache = NGramCache::new(3, 100);
        cache.observe(&[1, 2, 3]);
        let candidates = cache.top_candidates(&[1, 2], 5);
        assert_eq!(candidates, vec![3]);
    }

    #[test]
    fn test_ngram_cache_top_candidates_most_frequent_first() {
        // observe token 42 three times and token 7 once for prefix [1,2]
        let mut cache = NGramCache::new(3, 100);
        cache.observe(&[1, 2, 42]);
        cache.observe(&[1, 2, 42]);
        cache.observe(&[1, 2, 42]);
        cache.observe(&[1, 2, 7]);
        let candidates = cache.top_candidates(&[1, 2], 2);
        assert_eq!(candidates[0], 42, "42 should rank first (count=3)");
        assert_eq!(candidates[1], 7, "7 should rank second (count=1)");
    }

    #[test]
    fn test_ngram_cache_empty_prefix_returns_empty() {
        let cache = NGramCache::new(3, 100);
        let candidates = cache.top_candidates(&[99, 100], 5);
        assert!(
            candidates.is_empty(),
            "unknown prefix must return empty vec"
        );
    }

    #[test]
    fn test_ngram_cache_observe_window_slides() {
        // observe [1,2,3,4] with n=3: should store (1,2)→3 and (2,3)→4
        let mut cache = NGramCache::new(3, 100);
        cache.observe(&[1, 2, 3, 4]);
        assert_eq!(cache.top_candidates(&[1, 2], 1), vec![3]);
        assert_eq!(cache.top_candidates(&[2, 3], 1), vec![4]);
    }

    #[test]
    fn test_ngram_cache_top_k_limit() {
        let mut cache = NGramCache::new(3, 100);
        // Insert 5 different next tokens for the same prefix.
        for next in 10u32..15 {
            cache.observe(&[1, 2, next]);
        }
        let candidates = cache.top_candidates(&[1, 2], 3);
        assert_eq!(candidates.len(), 3, "should return at most k=3 candidates");
    }

    #[test]
    fn test_ngram_cache_total_observations() {
        let mut cache = NGramCache::new(3, 100);
        // [1,2,3] → 1 observation, [2,3,4] → 1 observation from one call
        cache.observe(&[1, 2, 3, 4]);
        assert_eq!(cache.total_observations(), 2);
        // Add two more.
        cache.observe(&[5, 6, 7]);
        assert_eq!(cache.total_observations(), 3);
    }

    #[test]
    fn test_ngram_cache_num_prefixes() {
        let mut cache = NGramCache::new(3, 100);
        cache.observe(&[1, 2, 3]);
        assert_eq!(cache.num_prefixes(), 1);
        cache.observe(&[4, 5, 6]);
        assert_eq!(cache.num_prefixes(), 2);
        // Observing the same prefix again must not increase count.
        cache.observe(&[1, 2, 9]);
        assert_eq!(cache.num_prefixes(), 2);
    }

    #[test]
    fn test_warm_cache_populates_candidates() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        decoder.warm_cache(&[10, 20, 30, 40, 50]);
        // After warming, the prefix [10,20] should yield at least one candidate.
        let candidates = decoder.cache().top_candidates(&[10, 20], 3);
        assert!(
            !candidates.is_empty(),
            "warm_cache should populate the cache"
        );
    }

    // ------------------------------------------------------------------
    // Decoder integration tests
    // ------------------------------------------------------------------

    #[test]
    fn test_decode_empty_cache_falls_back_to_model() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig {
            window_size: 5,
            ..LookaheadConfig::default()
        });
        // No warm_cache call — cache is empty, so every step must call score_fn.
        let initial = vec![1u32, 2, 3];
        let generated = decoder.decode(&initial, 5, incrementer);
        assert_eq!(generated.len(), 5);
        // All tokens should come from the model (none from ngram).
        assert_eq!(decoder.stats().total_accepted_from_ngram, 0);
        assert!(decoder.stats().total_accepted_from_model > 0);
    }

    #[test]
    fn test_decode_returns_correct_length() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        let generated = decoder.decode(&[1, 2, 3], 8, incrementer);
        assert_eq!(generated.len(), 8, "should generate exactly max_new_tokens");
    }

    #[test]
    fn test_decode_stops_on_eos() {
        let eos: u32 = 99;
        let config = LookaheadConfig {
            eos_token_id: eos,
            ..LookaheadConfig::default()
        };
        let mut decoder = LookaheadDecoder::new(config);
        // score_fn always returns the EOS token.
        let generated = decoder.decode(&[1, 2, 3], 100, |_| eos);
        // Should stop after the first token (which is EOS).
        assert!(!generated.is_empty());
        assert_eq!(*generated.last().unwrap(), eos);
        assert!(generated.len() < 100, "should stop early on EOS");
    }

    #[test]
    fn test_decode_with_matching_cache() {
        // Warm the cache so it exactly predicts what the incrementer will generate.
        // incrementer produces: 4,5,6,7,8,... starting from initial=[1,2,3].
        let corpus: Vec<u32> = (1u32..=20).collect();
        let config = LookaheadConfig {
            window_size: 4,
            ngram_order: 3,
            eos_token_id: 255,
            ..LookaheadConfig::default()
        };
        let mut decoder = LookaheadDecoder::new(config);
        decoder.warm_cache(&corpus);

        let generated = decoder.decode(&[1, 2, 3], 10, incrementer);
        assert_eq!(generated.len(), 10);
        // Cache matches the model so we should see some ngram acceptances.
        assert!(
            decoder.stats().total_accepted_from_ngram > 0,
            "warm matching cache should yield ngram acceptances"
        );
    }

    #[test]
    fn test_decode_with_mismatched_cache() {
        // Cache predicts token 99 for every position; model (incrementer) disagrees.
        let mut decoder = LookaheadDecoder::new(LookaheadConfig {
            window_size: 4,
            ngram_order: 2,
            eos_token_id: 255,
            ..LookaheadConfig::default()
        });
        // Observe [1, 99] → cache predicts 99 after any prefix ending in 1.
        decoder.warm_cache(&[1, 99, 99, 99, 99, 99]);

        // incrementer starting from [1] will produce 2,3,4,... never 99.
        let generated = decoder.decode(&[1], 5, incrementer);
        assert_eq!(generated.len(), 5);
        // Model should have overridden the cache prediction many times.
        assert!(decoder.stats().total_accepted_from_model > 0);
    }

    // ------------------------------------------------------------------
    // Statistics tests
    // ------------------------------------------------------------------

    #[test]
    fn test_stats_total_tokens_matches_generated_length() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        let generated = decoder.decode(&[1, 2, 3], 12, incrementer);
        assert_eq!(
            decoder.stats().total_tokens_generated,
            generated.len(),
            "stats counter must equal returned vec length"
        );
    }

    #[test]
    fn test_stats_steps_counted() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        decoder.decode(&[1, 2, 3], 6, incrementer);
        assert!(
            decoder.stats().total_steps > 0,
            "total_steps must be incremented each iteration"
        );
    }

    #[test]
    fn test_stats_avg_tokens_per_step_at_least_one() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        decoder.decode(&[1, 2, 3], 10, incrementer);
        assert!(
            decoder.stats().avg_tokens_per_step >= 1.0,
            "avg_tokens_per_step must be >= 1.0"
        );
    }

    #[test]
    fn test_stats_ngram_accepted_positive_when_cache_matches() {
        let corpus: Vec<u32> = (1u32..=30).collect();
        let config = LookaheadConfig {
            window_size: 4,
            ngram_order: 3,
            eos_token_id: 200,
            ..LookaheadConfig::default()
        };
        let mut decoder = LookaheadDecoder::new(config);
        decoder.warm_cache(&corpus);
        decoder.decode(&[1, 2, 3], 15, incrementer);
        assert!(
            decoder.stats().total_accepted_from_ngram > 0,
            "matching cache should produce ngram acceptances"
        );
    }

    #[test]
    fn test_reset_stats_clears_all_counters() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        decoder.decode(&[1, 2, 3], 5, incrementer);
        // Sanity: stats were actually incremented.
        assert!(decoder.stats().total_steps > 0);
        decoder.reset_stats();
        let s = decoder.stats();
        assert_eq!(s.total_tokens_generated, 0);
        assert_eq!(s.total_steps, 0);
        assert_eq!(s.total_accepted_from_ngram, 0);
        assert_eq!(s.total_accepted_from_model, 0);
        assert_eq!(s.avg_tokens_per_step, 0.0);
    }

    #[test]
    fn test_decode_longer_acceptance_when_cache_warm_and_model_agrees() {
        // Build a corpus that perfectly matches what incrementer will generate.
        let corpus: Vec<u32> = (1u32..=50).collect();
        let config = LookaheadConfig {
            window_size: 5,
            ngram_order: 3,
            cache_size: 10_000,
            max_candidates_per_step: 3,
            eos_token_id: 200,
            max_length: 256,
        };
        let mut decoder = LookaheadDecoder::new(config);
        decoder.warm_cache(&corpus);

        decoder.decode(&[1, 2, 3], 20, incrementer);
        assert!(
            decoder.stats().avg_tokens_per_step > 1.0,
            "avg_tokens_per_step should exceed 1.0 when cache is warm and matches the model"
        );
    }

    // Extra coverage for edge conditions and k=0 guard.
    #[test]
    fn test_ngram_cache_top_k_zero_returns_empty() {
        let mut cache = NGramCache::new(3, 100);
        cache.observe(&[1, 2, 3]);
        assert!(cache.top_candidates(&[1, 2], 0).is_empty());
    }

    #[test]
    fn test_decode_zero_new_tokens_returns_empty() {
        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
        let generated = decoder.decode(&[1, 2, 3], 0, incrementer);
        assert!(
            generated.is_empty(),
            "max_new_tokens=0 should return empty vec"
        );
    }
}