//! Grammar-constrained sampling with CFG parsing
//!
//! Provides comprehensive grammar-guided generation including:
//! - Context-free grammar (CFG) parsing and enforcement
//! - BNF/EBNF grammar specification support  
//! - Grammar-guided beam search with constraint propagation
//! - Syntax error recovery and repair mechanisms
//! - Grammar conflict resolution and ambiguity handling
//! - Custom grammar definition and validation

use anyhow::{Result, anyhow};
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet, VecDeque},
    fmt,
};

/// Grammar rule types
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum GrammarSymbol {
    Terminal(String),
    NonTerminal(String),
    Epsilon, // Empty production
}

impl fmt::Display for GrammarSymbol {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            GrammarSymbol::Terminal(s) => write!(f, "'{s}'"),
            GrammarSymbol::NonTerminal(s) => write!(f, "{s}"),
            GrammarSymbol::Epsilon => write!(f, "ε"),
        }
    }
}

/// Grammar production rule
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ProductionRule {
    pub left_hand_side: String,
    pub right_hand_side: Vec<GrammarSymbol>,
    pub weight: f64, // For probabilistic grammars
}

impl ProductionRule {
    /// Create new production rule
    #[must_use]
    pub fn new(lhs: String, rhs: Vec<GrammarSymbol>) -> Self {
        Self {
            left_hand_side: lhs,
            right_hand_side: rhs,
            weight: 1.0,
        }
    }

    /// Create weighted production rule
    #[must_use]
    pub fn with_weight(lhs: String, rhs: Vec<GrammarSymbol>, weight: f64) -> Self {
        Self {
            left_hand_side: lhs,
            right_hand_side: rhs,
            weight,
        }
    }

    /// Check if rule produces epsilon
    #[must_use]
    pub fn is_epsilon(&self) -> bool {
        self.right_hand_side.len() == 1 && self.right_hand_side[0] == GrammarSymbol::Epsilon
    }
}

/// Context-free grammar
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextFreeGrammar {
    pub start_symbol: String,
    pub rules: Vec<ProductionRule>,
    pub terminals: HashSet<String>,
    pub non_terminals: HashSet<String>,
}

impl ContextFreeGrammar {
    /// Create new grammar
    #[must_use]
    pub fn new(start_symbol: String) -> Self {
        Self {
            start_symbol: start_symbol.clone(),
            rules: Vec::new(),
            terminals: HashSet::new(),
            non_terminals: vec![start_symbol].into_iter().collect(),
        }
    }

    /// Add production rule
    pub fn add_rule(&mut self, rule: ProductionRule) {
        self.non_terminals.insert(rule.left_hand_side.clone());

        // Note: non-terminals referenced only on the right-hand side are
        // intentionally NOT added here -- `non_terminals` tracks symbols
        // that have a defining rule, so referencing an undefined symbol
        // can be detected by `validate()`.
        for symbol in &rule.right_hand_side {
            if let GrammarSymbol::Terminal(term) = symbol {
                self.terminals.insert(term.clone());
            }
        }

        self.rules.push(rule);
    }

    /// Get rules for a non-terminal
    #[must_use]
    pub fn get_rules(&self, non_terminal: &str) -> Vec<&ProductionRule> {
        self.rules
            .iter()
            .filter(|rule| rule.left_hand_side == non_terminal)
            .collect()
    }

    /// Check if symbol can derive epsilon
    #[must_use]
    pub fn nullable(&self, symbol: &str) -> bool {
        let mut nullable_set = HashSet::new();
        let mut changed = true;

        while changed {
            changed = false;

            for rule in &self.rules {
                if rule.is_epsilon() && !nullable_set.contains(&rule.left_hand_side) {
                    nullable_set.insert(rule.left_hand_side.clone());
                    changed = true;
                } else if rule.right_hand_side.iter().all(|sym| match sym {
                    GrammarSymbol::NonTerminal(nt) => nullable_set.contains(nt),
                    GrammarSymbol::Terminal(_) => false,
                    GrammarSymbol::Epsilon => true,
                }) && !nullable_set.contains(&rule.left_hand_side)
                {
                    nullable_set.insert(rule.left_hand_side.clone());
                    changed = true;
                }
            }
        }

        nullable_set.contains(symbol)
    }

    /// Compute FIRST set for a symbol
    #[must_use]
    pub fn first_set(&self, symbol: &str) -> HashSet<String> {
        let mut first_sets: HashMap<String, HashSet<String>> = HashMap::new();
        let mut changed = true;

        // Initialize FIRST sets
        for terminal in &self.terminals {
            first_sets.insert(
                terminal.clone(),
                vec![terminal.clone()].into_iter().collect(),
            );
        }

        for non_terminal in &self.non_terminals {
            first_sets.insert(non_terminal.clone(), HashSet::new());
        }

        while changed {
            changed = false;

            for rule in &self.rules {
                let lhs = &rule.left_hand_side;
                let old_size = first_sets
                    .get(lhs)
                    .map_or(0, std::collections::HashSet::len);

                if rule.is_epsilon() {
                    first_sets
                        .entry(lhs.clone())
                        .or_default()
                        .insert("ε".to_string());
                } else {
                    for symbol in &rule.right_hand_side {
                        match symbol {
                            GrammarSymbol::Terminal(term) => {
                                first_sets
                                    .entry(lhs.clone())
                                    .or_default()
                                    .insert(term.clone());
                                break;
                            }
                            GrammarSymbol::NonTerminal(nt) => {
                                // Clone the set to avoid borrow checker issues
                                let nt_first_clone = first_sets.get(nt).cloned();
                                if let Some(nt_first) = nt_first_clone {
                                    let mut added_epsilon = false;
                                    for first_symbol in &nt_first {
                                        if first_symbol == "ε" {
                                            added_epsilon = true;
                                        } else {
                                            first_sets
                                                .entry(lhs.clone())
                                                .or_default()
                                                .insert(first_symbol.clone());
                                        }
                                    }
                                    if !added_epsilon {
                                        break;
                                    }
                                }
                            }
                            GrammarSymbol::Epsilon => {
                                first_sets
                                    .entry(lhs.clone())
                                    .or_default()
                                    .insert("ε".to_string());
                            }
                        }
                    }
                }

                let new_size = first_sets
                    .get(lhs)
                    .map_or(0, std::collections::HashSet::len);
                if new_size > old_size {
                    changed = true;
                }
            }
        }

        first_sets.get(symbol).cloned().unwrap_or_default()
    }

    /// Validate grammar for consistency
    pub fn validate(&self) -> Result<GrammarValidationResult> {
        let mut errors = Vec::new();
        let mut warnings = Vec::new();

        // Check if start symbol exists
        if !self.non_terminals.contains(&self.start_symbol) {
            errors.push(GrammarError {
                message: format!(
                    "Start symbol '{}' not found in non-terminals",
                    self.start_symbol
                ),
                error_type: GrammarErrorType::UndefinedSymbol,
                location: None,
            });
        }

        // Check for undefined non-terminals
        for rule in &self.rules {
            for symbol in &rule.right_hand_side {
                if let GrammarSymbol::NonTerminal(nt) = symbol
                    && !self.non_terminals.contains(nt)
                {
                    errors.push(GrammarError {
                        message: format!("Undefined non-terminal '{nt}' in rule"),
                        error_type: GrammarErrorType::UndefinedSymbol,
                        location: Some(rule.left_hand_side.clone()),
                    });
                }
            }
        }

        // Check for unreachable symbols
        let reachable = self.compute_reachable_symbols();
        for nt in &self.non_terminals {
            if !reachable.contains(nt) {
                warnings.push(GrammarWarning {
                    message: format!("Non-terminal '{nt}' is unreachable"),
                    warning_type: GrammarWarningType::UnreachableSymbol,
                });
            }
        }

        // Check for useless productions (don't derive terminals)
        let productive = self.compute_productive_symbols();
        for rule in &self.rules {
            if !productive.contains(&rule.left_hand_side) {
                warnings.push(GrammarWarning {
                    message: format!("Rule for '{}' is non-productive", rule.left_hand_side),
                    warning_type: GrammarWarningType::NonProductiveRule,
                });
            }
        }

        Ok(GrammarValidationResult {
            valid: errors.is_empty(),
            errors,
            warnings,
        })
    }

    /// Compute reachable symbols from start symbol
    fn compute_reachable_symbols(&self) -> HashSet<String> {
        let mut reachable = HashSet::new();
        let mut queue = VecDeque::new();

        queue.push_back(self.start_symbol.clone());
        reachable.insert(self.start_symbol.clone());

        while let Some(symbol) = queue.pop_front() {
            for rule in self.get_rules(&symbol) {
                for rhs_symbol in &rule.right_hand_side {
                    if let GrammarSymbol::NonTerminal(nt) = rhs_symbol
                        && !reachable.contains(nt)
                    {
                        reachable.insert(nt.clone());
                        queue.push_back(nt.clone());
                    }
                }
            }
        }

        reachable
    }

    /// Compute productive symbols (can derive terminal strings)
    fn compute_productive_symbols(&self) -> HashSet<String> {
        let mut productive = HashSet::new();
        let mut changed = true;

        while changed {
            changed = false;

            for rule in &self.rules {
                if productive.contains(&rule.left_hand_side) {
                    continue;
                }

                let all_productive = rule.right_hand_side.iter().all(|symbol| match symbol {
                    GrammarSymbol::Terminal(_) => true,
                    GrammarSymbol::NonTerminal(nt) => productive.contains(nt),
                    GrammarSymbol::Epsilon => true,
                });

                if all_productive {
                    productive.insert(rule.left_hand_side.clone());
                    changed = true;
                }
            }
        }

        productive
    }
}

/// Parse state for grammar-guided generation
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ParseState {
    pub stack: Vec<GrammarSymbol>,
    pub input_position: usize,
    pub derivation: Vec<String>, // Track derivation steps
}

impl ParseState {
    /// Create new parse state
    #[must_use]
    pub fn new(start_symbol: String) -> Self {
        Self {
            stack: vec![GrammarSymbol::NonTerminal(start_symbol)],
            input_position: 0,
            derivation: Vec::new(),
        }
    }

    /// Check if parsing is complete
    #[must_use]
    pub fn is_complete(&self) -> bool {
        self.stack.is_empty()
    }

    /// Get next expected terminals
    #[must_use]
    pub fn get_expected_terminals(&self, grammar: &ContextFreeGrammar) -> HashSet<String> {
        let mut expected = HashSet::new();

        if let Some(top) = self.stack.last() {
            match top {
                GrammarSymbol::Terminal(term) => {
                    expected.insert(term.clone());
                }
                GrammarSymbol::NonTerminal(nt) => {
                    let first_set = grammar.first_set(nt);
                    for symbol in first_set {
                        if symbol != "ε" {
                            expected.insert(symbol);
                        }
                    }
                }
                GrammarSymbol::Epsilon => {}
            }
        }

        expected
    }
}

/// Grammar-guided sampler
pub struct GrammarGuidedSampler {
    grammar: ContextFreeGrammar,
    beam_width: usize,
    max_depth: usize,
}

impl GrammarGuidedSampler {
    /// Create new grammar-guided sampler
    #[must_use]
    pub fn new(grammar: ContextFreeGrammar, beam_width: usize, max_depth: usize) -> Self {
        Self {
            grammar,
            beam_width,
            max_depth,
        }
    }

    /// Sample next tokens given current parse state
    pub fn sample_next_tokens(
        &self,
        state: &ParseState,
        token_probabilities: &HashMap<String, f64>,
    ) -> Result<Vec<(String, f64, ParseState)>> {
        let mut candidates = Vec::new();
        let expected_terminals = state.get_expected_terminals(&self.grammar);

        if expected_terminals.is_empty() {
            return Ok(candidates);
        }

        // Filter token probabilities by expected terminals
        for (token, prob) in token_probabilities {
            if expected_terminals.contains(token)
                && let Some(new_state) = self.advance_state(state, token)?
            {
                candidates.push((token.clone(), *prob, new_state));
            }
        }

        // Sort by probability and take top candidates
        candidates.sort_by(|a, b| b.1.total_cmp(&a.1));
        candidates.truncate(self.beam_width);

        Ok(candidates)
    }

    /// Advance parse state with given token
    pub fn advance_state(&self, state: &ParseState, token: &str) -> Result<Option<ParseState>> {
        if state.stack.is_empty() {
            return Ok(None);
        }

        let mut new_state = state.clone();

        // Check if top of stack matches token
        if let Some(top) = new_state.stack.pop() {
            match top {
                GrammarSymbol::Terminal(term) => {
                    if term == token {
                        new_state.input_position += 1;
                        new_state.derivation.push(format!("consume '{token}'"));
                        Ok(Some(new_state))
                    } else {
                        Ok(None) // Token doesn't match expected terminal
                    }
                }
                GrammarSymbol::NonTerminal(nt) => {
                    // Find applicable rules
                    let rules = self.grammar.get_rules(&nt);
                    let mut possible_states = Vec::new();

                    for rule in rules {
                        let mut candidate_state = ParseState {
                            stack: new_state.stack.clone(),
                            input_position: new_state.input_position,
                            derivation: new_state.derivation.clone(),
                        };

                        // Add rule's RHS to stack in reverse order
                        for symbol in rule.right_hand_side.iter().rev() {
                            if *symbol != GrammarSymbol::Epsilon {
                                candidate_state.stack.push(symbol.clone());
                            }
                        }

                        candidate_state
                            .derivation
                            .push(format!("apply rule: {} -> {:?}", nt, rule.right_hand_side));

                        // Check if this rule can eventually consume the token
                        if self.can_consume_token(&candidate_state, token)? {
                            possible_states.push((candidate_state, rule.weight));
                        }
                    }

                    if possible_states.is_empty() {
                        Ok(None)
                    } else {
                        // Select rule based on weights (for now, pick the first valid one)
                        possible_states.sort_by(|a, b| b.1.total_cmp(&a.1));

                        // Recursively advance with the selected rule
                        self.advance_state(&possible_states[0].0, token)
                    }
                }
                GrammarSymbol::Epsilon => {
                    // Should not happen as epsilon is handled during rule application
                    Ok(None)
                }
            }
        } else {
            Ok(None)
        }
    }

    /// Check if state can eventually consume the given token
    fn can_consume_token(&self, state: &ParseState, token: &str) -> Result<bool> {
        if state.stack.is_empty() {
            return Ok(false);
        }

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(state.clone());

        let mut depth = 0;
        while let Some(current_state) = queue.pop_front() {
            depth += 1;
            if depth > self.max_depth {
                break;
            }

            let state_key = (current_state.stack.clone(), current_state.input_position);
            if visited.contains(&state_key) {
                continue;
            }
            visited.insert(state_key);

            if let Some(top) = current_state.stack.last() {
                match top {
                    GrammarSymbol::Terminal(term) => {
                        if term == token {
                            return Ok(true);
                        }
                    }
                    GrammarSymbol::NonTerminal(nt) => {
                        let rules = self.grammar.get_rules(nt);
                        for rule in rules {
                            let mut new_state = current_state.clone();
                            new_state.stack.pop(); // Remove the non-terminal

                            // Add rule's RHS to stack in reverse order
                            for symbol in rule.right_hand_side.iter().rev() {
                                if *symbol != GrammarSymbol::Epsilon {
                                    new_state.stack.push(symbol.clone());
                                }
                            }

                            queue.push_back(new_state);
                        }
                    }
                    GrammarSymbol::Epsilon => {}
                }
            }
        }

        Ok(false)
    }

    /// Perform beam search with grammar constraints
    pub fn beam_search(
        &self,
        initial_state: ParseState,
        max_length: usize,
        get_token_probs: impl Fn(&ParseState) -> Result<HashMap<String, f64>>,
    ) -> Result<Vec<(Vec<String>, f64, ParseState)>> {
        let mut beam = vec![(Vec::new(), 1.0, initial_state)];
        let mut results = Vec::new();

        for _step in 0..max_length {
            let mut candidates = Vec::new();

            // Use drain to consume beam while allowing reuse
            for (sequence, prob, state) in beam.drain(..) {
                if state.is_complete() {
                    results.push((sequence, prob, state));
                    continue;
                }

                let token_probs = get_token_probs(&state)?;
                let next_candidates = self.sample_next_tokens(&state, &token_probs)?;

                for (token, token_prob, new_state) in next_candidates {
                    let mut new_sequence = sequence.clone();
                    new_sequence.push(token);
                    let new_prob = prob * token_prob;

                    candidates.push((new_sequence, new_prob, new_state));
                }
            }

            if candidates.is_empty() {
                break;
            }

            // Select top candidates for next beam
            candidates.sort_by(|a, b| b.1.total_cmp(&a.1));
            candidates.truncate(self.beam_width);

            beam = candidates;
        }

        // Add remaining beam states to results
        results.extend(beam);

        // Sort results by probability
        results.sort_by(|a, b| b.1.total_cmp(&a.1));

        Ok(results)
    }
}

/// BNF grammar parser
pub struct BnfParser;

impl BnfParser {
    /// Parse BNF grammar from string
    pub fn parse(input: &str) -> Result<ContextFreeGrammar> {
        let lines: Vec<&str> = input
            .lines()
            .map(str::trim)
            .filter(|line| !line.is_empty() && !line.starts_with('#'))
            .collect();

        if lines.is_empty() {
            return Err(anyhow!("Empty grammar"));
        }

        let mut grammar = ContextFreeGrammar::new(String::new());
        let mut first_rule = true;

        for line in lines {
            if let Some((lhs, rhs)) = line.split_once("::=").or_else(|| line.split_once("->")) {
                let lhs = lhs.trim();

                // Set start symbol from first rule
                if first_rule {
                    grammar.start_symbol = lhs.to_string();
                    first_rule = false;
                }

                let productions: Vec<&str> = rhs.split('|').collect();

                for production in productions {
                    let production = production.trim();
                    let symbols = Self::parse_production(production)?;
                    let rule = ProductionRule::new(lhs.to_string(), symbols);
                    grammar.add_rule(rule);
                }
            } else {
                return Err(anyhow!("Invalid rule format: {line}"));
            }
        }

        Ok(grammar)
    }

    /// Parse production right-hand side
    fn parse_production(production: &str) -> Result<Vec<GrammarSymbol>> {
        let mut symbols = Vec::new();

        if production.is_empty() || production == "ε" || production == "epsilon" {
            symbols.push(GrammarSymbol::Epsilon);
            return Ok(symbols);
        }

        // Simple tokenization - could be improved for complex grammars
        let tokens: Vec<&str> = production.split_whitespace().collect();

        for token in tokens {
            if token.starts_with('"') && token.ends_with('"') {
                // Terminal symbol
                let terminal = token[1..token.len() - 1].to_string();
                symbols.push(GrammarSymbol::Terminal(terminal));
            } else if token.starts_with('\'') && token.ends_with('\'') {
                // Terminal symbol
                let terminal = token[1..token.len() - 1].to_string();
                symbols.push(GrammarSymbol::Terminal(terminal));
            } else if token.starts_with('<') && token.ends_with('>') {
                // Non-terminal symbol
                let non_terminal = token[1..token.len() - 1].to_string();
                symbols.push(GrammarSymbol::NonTerminal(non_terminal));
            } else {
                // Assume non-terminal if not quoted
                symbols.push(GrammarSymbol::NonTerminal(token.to_string()));
            }
        }

        Ok(symbols)
    }
}

/// Grammar validation result
#[derive(Debug, Serialize, Deserialize)]
pub struct GrammarValidationResult {
    pub valid: bool,
    pub errors: Vec<GrammarError>,
    pub warnings: Vec<GrammarWarning>,
}

/// Grammar error
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GrammarError {
    pub message: String,
    pub error_type: GrammarErrorType,
    pub location: Option<String>,
}

/// Grammar error types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GrammarErrorType {
    UndefinedSymbol,
    CyclicRule,
    InvalidSyntax,
}

/// Grammar warning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GrammarWarning {
    pub message: String,
    pub warning_type: GrammarWarningType,
}

/// Grammar warning types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GrammarWarningType {
    UnreachableSymbol,
    NonProductiveRule,
    AmbiguousGrammar,
}

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

    #[test]
    fn test_production_rule_creation() {
        let rule = ProductionRule::new(
            "S".to_string(),
            vec![
                GrammarSymbol::Terminal("a".to_string()),
                GrammarSymbol::NonTerminal("B".to_string()),
            ],
        );

        assert_eq!(rule.left_hand_side, "S");
        assert_eq!(rule.right_hand_side.len(), 2);
        assert!(!rule.is_epsilon());
    }

    #[test]
    fn test_epsilon_rule() {
        let rule = ProductionRule::new("S".to_string(), vec![GrammarSymbol::Epsilon]);

        assert!(rule.is_epsilon());
    }

    #[test]
    fn test_grammar_creation() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![GrammarSymbol::Terminal("a".to_string())],
        ));

        assert_eq!(grammar.start_symbol, "S");
        assert_eq!(grammar.rules.len(), 1);
        assert!(grammar.terminals.contains("a"));
        assert!(grammar.non_terminals.contains("S"));
    }

    #[test]
    fn test_nullable_computation() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        // S -> ε
        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![GrammarSymbol::Epsilon],
        ));

        assert!(grammar.nullable("S"));
    }

    #[test]
    fn test_first_set_computation() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        // S -> "a" B
        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![
                GrammarSymbol::Terminal("a".to_string()),
                GrammarSymbol::NonTerminal("B".to_string()),
            ],
        ));

        let first_set = grammar.first_set("S");
        assert!(first_set.contains("a"));
    }

    #[test]
    fn test_parse_state() {
        let state = ParseState::new("S".to_string());

        assert_eq!(state.stack.len(), 1);
        assert_eq!(state.input_position, 0);
        assert!(!state.is_complete());

        let mut complete_state = ParseState::new("S".to_string());
        complete_state.stack.clear();
        assert!(complete_state.is_complete());
    }

    #[test]
    fn test_bnf_parser() {
        let bnf_grammar = r#"
            S ::= "a" B | "b"
            B ::= "c" | ε
        "#;

        let grammar = BnfParser::parse(bnf_grammar).unwrap();
        assert_eq!(grammar.start_symbol, "S");
        assert_eq!(grammar.rules.len(), 4); // S -> "a" B, S -> "b", B -> "c", B -> ε
    }

    #[test]
    fn test_grammar_validation() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        // Valid grammar: S -> "a"
        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![GrammarSymbol::Terminal("a".to_string())],
        ));

        let result = grammar.validate().unwrap();
        assert!(result.valid);
        assert!(result.errors.is_empty());
    }

    #[test]
    fn test_invalid_grammar() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        // Invalid grammar: S -> B (B is undefined)
        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![GrammarSymbol::NonTerminal("B".to_string())],
        ));

        let result = grammar.validate().unwrap();
        assert!(!result.valid);
        assert!(!result.errors.is_empty());
    }

    #[test]
    fn test_grammar_guided_sampler() {
        let mut grammar = ContextFreeGrammar::new("S".to_string());

        grammar.add_rule(ProductionRule::new(
            "S".to_string(),
            vec![GrammarSymbol::Terminal("hello".to_string())],
        ));

        let sampler = GrammarGuidedSampler::new(grammar, 5, 10);
        let state = ParseState::new("S".to_string());

        let mut token_probs = HashMap::new();
        token_probs.insert("hello".to_string(), 0.8);
        token_probs.insert("world".to_string(), 0.2);

        let candidates = sampler.sample_next_tokens(&state, &token_probs).unwrap();
        assert!(!candidates.is_empty());
        assert_eq!(candidates[0].0, "hello");
    }
}