glossia-cli 0.2.0

// Proof-of-concept: tiny CFG sentence generator that embeds a POS-tagged payload
// (e.g., BIP39 words) in-order, inserting only SFW cover words.
//
// Cargo.toml:
// [dependencies]
// rand = "0.8"

use rand::{seq::SliceRandom, Rng, SeedableRng};
use rand::rngs::StdRng;
use std::collections::{HashMap, HashSet};
use std::env;
use glossia::types::Pos;
use glossia::grammar::Grammar;
use glossia::generator::{
    PayloadTok, Lexicon, GenerationMode, SentenceLengthMode,
    generate_text_with_original_payload, generate_text_best_of,
};
use glossia::generator::utils::{
    normalize_token_for_bip39, wrap_payload_with_bars, wrap_payload_with_color, word_wrap,
};
use glossia::generator::cache::print_sentence_kinds_once;
#[cfg(test)]
use glossia::generator::SequenceCache;

/// Check if a language has embedded files (packaged with the binary).
/// The language index is generated by glossia's build script and re-exported
/// from the library crate, so the CLI just delegates to it.
fn has_embedded_files(language: &str) -> bool {
    glossia::has_embedded_files(language)
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HighlightMode {
    None,
    Bars,
    Color(u8),  // ANSI color code (30-37 for foreground colors)
    Madlib,
}

/// Parse color name or ANSI code to color value
fn parse_color(color_str: &str) -> Result<u8, String> {
    match color_str.to_lowercase().as_str() {
        "black" => Ok(30),
        "red" => Ok(31),
        "green" => Ok(32),
        "yellow" => Ok(33),
        "blue" => Ok(34),
        "magenta" => Ok(35),
        "cyan" => Ok(36),
        "white" => Ok(37),
        _ => {
            // Try parsing as a number
            color_str.parse::<u8>()
                .map_err(|_| format!("Invalid color: {}. Use a color name (red, green, blue, etc.) or ANSI code (30-37)", color_str))
        }
    }
}



// Lexicon implementation is now in the library

// All core generation functions are now in the library (generator::core)
// All data loading functions are now in the library (generator::data)
// All utility functions are now in the library (generator::utils)
// Keeping only CLI-specific code below: HighlightMode, highlighting functions, argument parsing, main

/// Apply highlighting to generated text based on highlight_mode.
/// Preserves embedded newlines (e.g., from CS grammar's Conj "\n" cover words)
/// by processing each line independently.
fn apply_highlighting(
    text: &str,
    payload_set: &HashSet<String>,
    payload: &[PayloadTok],
    highlight_mode: HighlightMode,
    merkle_set: Option<&HashSet<String>>,
    merkle_highlight_mode: Option<HighlightMode>,
) -> String {
    // Process each line independently to preserve embedded newlines
    // (e.g., ASCII armor header/body/footer separation via Conj "\n" cover words)
    text.split('\n')
        .map(|line| {
            let words: Vec<&str> = line.split_whitespace().collect();
            if words.is_empty() {
                return String::new();
            }

            if highlight_mode == HighlightMode::Madlib {
                // For madlib mode, replace payload words with [POS] placeholders
                words.iter().map(|word| {
                    let word_clean = normalize_token_for_bip39(word);
                    let is_payload = !word_clean.is_empty() && payload_set.contains(&word_clean);
                    let is_merkle = merkle_set.map(|s| !word_clean.is_empty() && s.contains(&word_clean)).unwrap_or(false);

                    if is_payload || is_merkle {
                        // Find the payload token to get its POS tag
                        if let Some(payload_tok) = payload.iter().find(|t| t.word.to_lowercase() == word_clean) {
                            // Use the first allowed POS tag for madlib
                            if let Some(&first_pos) = payload_tok.allowed.iter().next() {
                                let pos_str = first_pos.as_str();
                                // Preserve punctuation
                                let punct: String = word.chars().filter(|c| !c.is_alphabetic()).collect();
                                format!("[{}]{}", pos_str, punct)
                            } else {
                                word.to_string()
                            }
                        } else {
                            word.to_string()
                        }
                    } else {
                        word.to_string()
                    }
                }).collect::<Vec<String>>().join(" ")
            } else {
                words.iter().map(|word| {
                    let word_clean = normalize_token_for_bip39(word);
                    let is_payload = !word_clean.is_empty() && payload_set.contains(&word_clean);
                    let is_merkle = merkle_set.map(|s| !word_clean.is_empty() && s.contains(&word_clean)).unwrap_or(false);

                    if is_payload {
                        match highlight_mode {
                            HighlightMode::None => word.to_string(),
                            HighlightMode::Bars => wrap_payload_with_bars(word),
                            HighlightMode::Color(color) => wrap_payload_with_color(word, color),
                            HighlightMode::Madlib => unreachable!(), // Handled above
                        }
                    } else if is_merkle {
                        // Use merkle highlight mode if specified, otherwise don't highlight
                        match merkle_highlight_mode {
                            Some(HighlightMode::None) => word.to_string(),
                            Some(HighlightMode::Bars) => wrap_payload_with_bars(word),
                            Some(HighlightMode::Color(color)) => wrap_payload_with_color(word, color),
                            Some(HighlightMode::Madlib) => word.to_string(), // Madlib doesn't apply to merkle words
                            None => word.to_string(), // No merkle highlighting specified
                        }
                    } else {
                        word.to_string()
                    }
                }).collect::<Vec<String>>().join(" ")
            }
        })
        .collect::<Vec<String>>()
        .join("\n")
}


/// Compute k candidates based on the length mode.
/// Returns a vector of k values to try in order.
#[cfg(test)]
fn compute_k_candidates<R: Rng>(
    rng: &mut R,
    cache: &SequenceCache,
    start_symbol: &str,
    k_min: usize,
    k_max: usize,
    length_mode: SentenceLengthMode,
    require_prefix: bool,
) -> Vec<usize> {
    match length_mode {
        SentenceLengthMode::Compact => {
            // Compact mode: try k from k_min to k_max, shortest first
            let k_start = if require_prefix { k_min + 1 } else { k_min };
            (k_start..=k_max).collect()
        }
        SentenceLengthMode::Natural => {
            // Natural mode: sample k from grammar's length distribution
            // Respect k_min as a floor, but sample naturally above it
            let natural_k_start = if require_prefix { k_min.max(2) } else { k_min.max(1) };
            
            // Compute weights for each k >= k_min
            let mut k_weights: Vec<(usize, f64)> = Vec::new();
            for k in natural_k_start..=k_max {
                if let Some(sequences) = cache.get(start_symbol, k) {
                    let weight: f64 = if require_prefix {
                        sequences.iter()
                            .filter(|seq_prob| !seq_prob.sequence.is_empty() && seq_prob.sequence[0] == Pos::Prefix)
                            .map(|seq_prob| seq_prob.probability)
                            .sum()
                    } else {
                        sequences.iter()
                            .map(|seq_prob| seq_prob.probability)
                            .sum()
                    };
                    if weight > 0.0 {
                        k_weights.push((k, weight));
                    }
                }
            }
            
            if k_weights.is_empty() {
                // Fallback to compact mode if no weights found
                let k_start = if require_prefix { k_min + 1 } else { k_min };
                return (k_start..=k_max).collect();
            }
            
            // Sample one k from the distribution
            let total_weight: f64 = k_weights.iter().map(|(_, w)| w).sum();
            if total_weight <= 0.0 {
                // Fallback to compact mode if total weight is zero
                let k_start = if require_prefix { k_min + 1 } else { k_min };
                return (k_start..=k_max).collect();
            }
            
            let mut r = rng.gen::<f64>() * total_weight;
            let mut sampled_k = None;
            for (k, weight) in &k_weights {
                if r <= *weight {
                    sampled_k = Some(*k);
                    break;
                }
                r -= weight;
            }
            let sampled_k = sampled_k.unwrap_or_else(|| k_weights[0].0);
            
            // Build candidate list:
            // 1. Sampled k first
            // 2. Remaining k's in descending weight order
            // 3. Compact fallback (k_min..=k_max) for robustness
            
            let mut candidates = vec![sampled_k];
            
            // Add remaining k's in descending weight order (excluding sampled_k)
            let mut remaining: Vec<(usize, f64)> = k_weights.iter()
                .filter(|(k, _)| *k != sampled_k)
                .cloned()
                .collect();
            remaining.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            candidates.extend(remaining.into_iter().map(|(k, _)| k));
            
            // Add compact fallback for robustness (only k's not already in candidates)
            // Still respect k_min for the fallback to ensure we don't miss any valid k values
            let candidates_set: std::collections::HashSet<usize> = candidates.iter().cloned().collect();
            let k_start = if require_prefix { k_min + 1 } else { k_min };
            for k in k_start..=k_max {
                if !candidates_set.contains(&k) {
                    candidates.push(k);
                }
            }
            
            candidates
        }
    }
}

/// Parse an endpoint string into an Endpoint value.
///
/// Supports formats: "hex", "base64", "ascii7", "bytes", "auto"
/// And languages: "english", "latin", "english/bip39", "english/bip39/body"
///
/// When `wordlist_override` is "default", resolves to the language's actual
/// default wordlist (e.g., "bip39" for English) via `default_wordlist()`.
fn parse_endpoint(s: &str, wordlist_override: &str) -> glossia::Endpoint {
    match s.to_lowercase().as_str() {
        "hex" => return glossia::Endpoint::Format(glossia::DataMode::Hex),
        "base64" => return glossia::Endpoint::Format(glossia::DataMode::Base64),
        "ascii7" | "ascii" => return glossia::Endpoint::Format(glossia::DataMode::Ascii7),
        "bytes" | "bytes8" => return glossia::Endpoint::Format(glossia::DataMode::Bytes8),
        "auto" => return glossia::Endpoint::Auto,
        _ => {}
    }

    // Resolve effective wordlist for a language name.
    let resolve_wordlist = |lang: &str, explicit_wl: &str| -> String {
        if explicit_wl != "default" {
            explicit_wl.to_string()
        } else {
            let dw = glossia::generator::default_wordlist(lang);
            if dw == "default" { "default".to_string() } else { dw.to_string() }
        }
    };

    // language/wordlist/dialect syntax
    let parts: Vec<&str> = s.split('/').collect();
    match parts.len() {
        3 => glossia::Endpoint::language_full(parts[0], parts[1], parts[2]),
        2 => glossia::Endpoint::Language {
            language: parts[0].to_string(),
            wordlist: parts[1].to_string(),
            dialect: "body".to_string(),
        },
        1 => {
            let lang = parts[0];
            let wl = resolve_wordlist(lang, wordlist_override);
            glossia::Endpoint::Language {
                language: lang.to_string(),
                wordlist: wl,
                dialect: "body".to_string(),
            }
        }
        _ => glossia::Endpoint::language(s),
    }
}

/// Encode text to wordlist words using base-N conversion.
fn encode_ascii_to_words(ascii_text: &str, language: &str, wordlist: &str, dialect: &str) -> Result<(Vec<String>, glossia::DataMode), String> {
    let all_words = glossia::generator::load_payload_words_for_wordlist(language, wordlist)?;
    let tree = glossia::WordlistTree::new(all_words);

    let codec = glossia::grammar::Grammar::from_language_dialect(language, dialect)
        .map(|g| g.codec().to_string())
        .unwrap_or_else(|_| "bitpack".to_string());
    glossia::codec::encode_str_base_n(ascii_text, &tree, &codec).map_err(|e| e.to_string())
}

/// Calculate input bits based on data mode and input text
fn calculate_input_bits(input_text: &str, mode: glossia::DataMode) -> usize {
    use glossia::DataMode;
    match mode {
        DataMode::Hex => {
            // Hex: 2 chars = 1 byte = 8 bits
            (input_text.len() / 2) * 8
        }
        DataMode::Base64 => {
            // Base64: 4 chars = 3 bytes = 24 bits
            (input_text.len() / 4) * 3 * 8
        }
        DataMode::Ascii7 => {
            // ASCII-7: 1 char = 7 bits
            input_text.len() * 7
        }
        DataMode::Bytes8 => {
            // UTF-8: count actual bytes
            input_text.as_bytes().len() * 8
        }
    }
}

/// Calculate output data bits based on wordlist size and word count.
///
/// For power-of-2 wordlists (bitpack mode), the first word is a padding word,
/// so data bits = (word_count - 1) * bits_per_word (minus padding, but we
/// report the gross bits here since the padding word is metadata).
///
/// The `codec` parameter controls whether bitpack accounting applies:
/// - `"bitpack"`: first word is padding metadata, rest are data words
/// - `"base_n"`: all words are data words (base-N encoding)
fn calculate_output_bits(word_count: usize, wordlist_size: usize, codec: &str) -> Option<usize> {
    if wordlist_size == 0 || !wordlist_size.is_power_of_two() {
        return None;
    }
    let bits_per_word = wordlist_size.trailing_zeros() as usize;
    if codec == "bitpack" && wordlist_size > 1 && word_count > 0 {
        // Bitpack mode: first word is padding metadata, rest are data words
        let data_words = word_count.saturating_sub(1);
        Some(data_words * bits_per_word)
    } else {
        Some(word_count * bits_per_word)
    }
}



fn print_usage(program_name: &str) {
    eprintln!("Usage: {} [OPTIONS] [<word1> <word2> ... <wordN>]", program_name);
    eprintln!();
    eprintln!("Generate natural sentences embedding payload words in-order.");
    eprintln!();
    eprintln!("Arguments:");
    eprintln!("  <word1> <word2> ...    Words to embed (positional, optional if --random or --from-ascii used)");
    eprintln!();
    eprintln!("Options:");
    eprintln!("  --random <N>            Generate sentences from N random payload words");
    eprintln!("  --from-ascii <text>     Encode ASCII plaintext to wordlist words");
    eprintln!("                          Use '-' to read from stdin");
    eprintln!("  --dialect <dialect>       Dialect: 'body' (default), 'subject', or 'payload_only'");
    eprintln!("                          Can include language: 'latin-body', 'english-subject', 'cs-nip04'");
    eprintln!("                          subject: Short sentences, may include prefixes (Re:, Fwd:, etc.)");
    eprintln!("                          body: Longer sentences, no prefixes");
    eprintln!("                          payload_only: Use only payload words, no cover words");
    eprintln!("  --width <N>              Fixed width for line wrapping (default: 80)");
    eprintln!("  --delimiter <str>        Delimiter between words in payload_only mode (default: \" \")");
    eprintln!("  --highlight-payload <mode>      Highlight payload words: 'none' (default), 'bars' (| |), or color name");
    eprintln!("                          Colors: black, red, green, yellow, blue, magenta, cyan, white");
    eprintln!("                          Or ANSI codes: 30-37");
    eprintln!("  --highlight-merkle <mode>      Highlight Merkle words: 'none', 'bars' (| |), or color name");
    eprintln!("                          (only works with --merkle)");
    eprintln!("                          Same color options as --highlight-payload");
    eprintln!("  --merkle                Wrap payload words in Merkle proof (expands N words to 2N-1)");
    eprintln!("  --demerkle              Parse mode: extract original payload from merkleized sequence");
    eprintln!("  --decode                Decode payload words back to bytes (inverse of --from-ascii)");
    eprintln!("                          Provide words as positional args or pipe via stdin");
    eprintln!("                          Auto-detects language/wordlist if --language not specified");
    eprintln!();
    eprintln!("Pipeline mode (translate between languages):");
    eprintln!("  The verb is authoritative — it decides how the input is read:");
    eprintln!("  --encode                Input is raw data; encode it (requires --to).");
    eprintln!("                          A bare mnemonic is encoded as text and round-trips.");
    eprintln!("  --transcode             Input is an existing Glossia encoding; detect its");
    eprintln!("                          source and re-encode (requires --to). With no --from");
    eprintln!("                          the source language is auto-detected.");
    eprintln!("  --decode                Decode an encoding back to data (with --from/--to);");
    eprintln!("                          on its own with words, the legacy word->bytes mode.");
    eprintln!("  --to / --into <target>  Target language/format (e.g. english, latin,");
    eprintln!("                          latin/default/prose, hex, base64).");
    eprintln!("  --from <source>         Source language/format (e.g. english, english/bip39,");
    eprintln!("                          english/bip39/raw, hex). Optional under --transcode.");
    eprintln!("  --meta <instruction>    Run a full meta-language instruction directly,");
    eprintln!("                          e.g. \"transcode from english into latin\".");
    eprintln!("                          Input comes from --from-ascii, positional args, or stdin");
    eprintln!("  --madlib                Replace payload words with [POS] placeholders");
    eprintln!("  --seed <N>              Seed for deterministic random generation");
    eprintln!("  --variations <N>         Generate N variations and select the most compact (default: 1)");
    eprintln!("  --best-of <N>            Sample N candidates per output and keep the densest / most");
    eprintln!("                           semantically coherent one (default: 1). English only.");
    eprintln!("  --language, -l <lang>    Language for wordlist: 'english' (default), 'french', 'german', 'meta', etc.");
    eprintln!("  --k-min <N>              Minimum sentence length in POS slots including Dot (default: 3)");
    eprintln!("  --k-max <N>              Maximum sentence length in POS slots including Dot (default: 20)");
    eprintln!("  --length-mode <mode>      Sentence length selection: 'compact' or 'natural'");
    eprintln!("                          default: subject -> compact, body -> natural");
    eprintln!("                          compact: Try k from k_min to k_max, shortest first");
    eprintln!("                          natural: Sample k from grammar's length distribution");
    eprintln!("  --wordlist <name>        Wordlist to use: 'bip39' (default), 'ngram', 'hp'");
    eprintln!("  --show-grammar           Display the grammar rules (then continue execution)");
    eprintln!("  --verbose, -v           Show detailed debugging information");
    eprintln!("  --help                  Show this help message");
    eprintln!();
    eprintln!("Examples:");
    eprintln!("  {} abandon ability able about above absent", program_name);
    eprintln!("  {} --random 10", program_name);
    eprintln!("  {} --from-ascii \"Hello World\"", program_name);
    eprintln!("  {} --from-ascii - < input.txt", program_name);
    eprintln!("  {} --decode word1 word2 word3", program_name);
    eprintln!("  {} --random 5 --dialect subject --highlight-payload none", program_name);
    eprintln!("  {} --random 5 --dialect latin-body --highlight-payload bars", program_name);
    eprintln!();
    eprintln!("Pipeline examples:");
    eprintln!("  {} --encode --to english --from-ascii \"Hello World\"", program_name);
    eprintln!("  {} --transcode --from english --to latin < english_prose.txt", program_name);
    eprintln!("  {} --transcode --to latin < seed_words.txt   (auto-detect source)", program_name);
    eprintln!("  {} --decode --from english < prose.txt   (decode to original data)", program_name);
    eprintln!("  {} --meta \"transcode from english into latin\" < prose.txt", program_name);
    eprintln!();
    eprintln!("Image rendering:");
    eprintln!("  -o <file.svg|file.png>      Output image file (with -l image)");
    eprintln!("  --render-image <file|->      Render existing text notation to image");
    eprintln!("                              Output format determined by -o extension (.svg default)");
    eprintln!();
    eprintln!("  {} -l image --dialect voronoi --random 20 -o out.svg", program_name);
    eprintln!("  {} -l image --dialect grid --random 16 -o grid.svg", program_name);
    eprintln!("  {} -l image --dialect constellation --random 12 -o stars.svg", program_name);
    eprintln!("  {} -l image --random 20 | {} --render-image - -o out.svg", program_name, program_name);
}

fn parse_args() -> Result<(Vec<String>, Option<usize>, Option<String>, bool, Option<u64>, usize, usize, HighlightMode, GenerationMode, String, String, bool, bool, usize, usize, bool, bool, SentenceLengthMode, usize, String, bool, bool, bool, Option<HighlightMode>, String, Option<String>, Option<String>, Option<String>, Option<String>, Option<String>), String> {
    let args: Vec<String> = env::args().collect();
    let program_name = args[0].clone();

    if args.len() < 2 {
        return Err("No words provided. Use --random <N> or provide words as arguments.".to_string());
    }

    let mut words = Vec::new();
    let mut random_count: Option<usize> = None;
    let mut ascii_input: Option<String> = None;
    let mut verbose = false;
    let mut seed: Option<u64> = None;
    let mut variations = 1;
    let mut best_of: usize = 1;
    let mut highlight_mode = HighlightMode::None;
    let mut generation_mode = GenerationMode::Body;
    let mut dialect_name = "body".to_string();
    let mut language = "english".to_string();
    let mut language_was_explicit = false;
    let mut show_grammar = false;
    let mut k_min = 3;
    let mut k_max = 20;
    let mut k_min_explicit = false;
    let mut k_max_explicit = false;
    let mut length_mode = SentenceLengthMode::Compact;
    let mut length_mode_explicit = false;
    let mut width = 80;
    let mut delimiter = " ".to_string();
    let mut merkle_mode = false;
    let mut demerkle_mode = false;
    let mut decode_mode = false;
    let mut merkle_highlight_mode: Option<HighlightMode> = None;
    let mut wordlist = "default".to_string();
    let mut meta_instruction: Option<String> = None;
    let mut pipeline_into: Option<String> = None;
    let mut pipeline_from: Option<String> = None;
    // Verb flags (encode/transcode/decode) drive a pipeline together with the
    // --from/--to endpoint helpers. `--decode` doubles as the legacy word→bytes
    // flag; it only drives a pipeline when an endpoint is also given.
    let mut encode_flag = false;
    let mut transcode_flag = false;
    let mut render_image: Option<String> = None;
    let mut render_output: Option<String> = None;
    let mut i = 1;
    
    while i < args.len() {
        match args[i].as_str() {
            "--help" | "-h" => {
                print_usage(&program_name);
                std::process::exit(0);
            }
            "--merkle" => {
                merkle_mode = true;
                i += 1;
            }
            "--demerkle" => {
                demerkle_mode = true;
                i += 1;
            }
            "--decode" => {
                decode_mode = true;
                i += 1;
            }
            "--verbose" | "-v" => {
                verbose = true;
                i += 1;
            }
            "--random" => {
                if i + 1 >= args.len() {
                    return Err("--random requires a value".to_string());
                }
                random_count = Some(args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --random: {}", args[i + 1]))?);
                i += 2;
            }
            "--dialect" => {
                if i + 1 >= args.len() {
                    return Err("--dialect requires a value".to_string());
                }
                let raw = &args[i + 1];
                // Support combined "language-dialect" syntax (e.g., "latin-body", "cs-nip04")
                if let Some(pos) = raw.find('-') {
                    language = raw[..pos].to_string();
                    language_was_explicit = true;
                    dialect_name = raw[pos + 1..].to_string();
                } else {
                    dialect_name = raw.clone();
                }
                generation_mode = if dialect_name.starts_with("subject") {
                    GenerationMode::Subject
                } else if dialect_name == "payload_only" {
                    GenerationMode::PayloadOnly
                } else {
                    GenerationMode::Body
                };
                i += 2;
            }
            "--highlight-payload" => {
                if i + 1 >= args.len() {
                    return Err("--highlight-payload requires a value".to_string());
                }
                highlight_mode = match args[i + 1].as_str() {
                    "none" => HighlightMode::None,
                    "bars" => HighlightMode::Bars,
                    color_str => {
                        match parse_color(color_str) {
                            Ok(color_code) => HighlightMode::Color(color_code),
                            Err(e) => return Err(e),
                        }
                    }
                };
                i += 2;
            }
            "--madlib" => {
                highlight_mode = HighlightMode::Madlib;
                i += 1;
            }
            "--highlight-merkle" => {
                if i + 1 >= args.len() {
                    return Err("--highlight-merkle requires a value".to_string());
                }
                merkle_highlight_mode = Some(match args[i + 1].as_str() {
                    "none" => HighlightMode::None,
                    "bars" => HighlightMode::Bars,
                    color_str => {
                        match parse_color(color_str) {
                            Ok(color_code) => HighlightMode::Color(color_code),
                            Err(e) => return Err(e),
                        }
                    }
                });
                i += 2;
            }
            "--seed" => {
                if i + 1 >= args.len() {
                    return Err("--seed requires a value".to_string());
                }
                seed = Some(args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --seed: {}", args[i + 1]))?);
                i += 2;
            }
            "--best-of" => {
                if i + 1 >= args.len() {
                    return Err("--best-of requires a value".to_string());
                }
                best_of = args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --best-of: {}", args[i + 1]))?;
                i += 2;
            }
            "--variations" => {
                if i + 1 >= args.len() {
                    return Err("--variations requires a value".to_string());
                }
                variations = args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --variations: {}", args[i + 1]))?;
                if variations == 0 {
                    return Err("--variations must be at least 1".to_string());
                }
                i += 2;
            }
            "--language" | "-l" => {
                if i + 1 >= args.len() {
                    return Err("--language requires a value".to_string());
                }
                language = args[i + 1].clone();
                language_was_explicit = true;
                i += 2;
            }
            "--wordlist" => {
                if i + 1 >= args.len() {
                    return Err("--wordlist requires a value".to_string());
                }
                wordlist = args[i + 1].clone();
                i += 2;
            }
            "--show-grammar" => {
                show_grammar = true;
                i += 1;
            }
            "--k-min" => {
                if i + 1 >= args.len() {
                    return Err("--k-min requires a value".to_string());
                }
                k_min = args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --k-min: {}", args[i + 1]))?;
                if k_min < 3 {
                    return Err("--k-min must be at least 3".to_string());
                }
                k_min_explicit = true;
                i += 2;
            }
            "--k-max" => {
                if i + 1 >= args.len() {
                    return Err("--k-max requires a value".to_string());
                }
                k_max = args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --k-max: {}", args[i + 1]))?;
                if k_max < k_min {
                    return Err(format!("--k-max ({}) must be >= --k-min ({})", k_max, k_min));
                }
                k_max_explicit = true;
                i += 2;
            }
            "--length-mode" => {
                if i + 1 >= args.len() {
                    return Err("--length-mode requires a value".to_string());
                }
                length_mode = match args[i + 1].as_str() {
                    "compact" => SentenceLengthMode::Compact,
                    "natural" => SentenceLengthMode::Natural,
                    _ => return Err(format!("Invalid length mode: {}. Use 'compact' or 'natural'", args[i + 1])),
                };
                length_mode_explicit = true;
                i += 2;
            }
            "--width" => {
                if i + 1 >= args.len() {
                    return Err("--width requires a value".to_string());
                }
                width = args[i + 1].parse()
                    .map_err(|_| format!("Invalid number for --width: {}", args[i + 1]))?;
                if width < 1 {
                    return Err("--width must be at least 1".to_string());
                }
                i += 2;
            }
            "--delimiter" => {
                if i + 1 >= args.len() {
                    return Err("--delimiter requires a value".to_string());
                }
                delimiter = args[i + 1].clone();
                i += 2;
            }
            "--meta" => {
                if i + 1 >= args.len() {
                    return Err("--meta requires a value (pipeline instruction)".to_string());
                }
                meta_instruction = Some(args[i + 1].clone());
                i += 2;
            }
            "--encode" => {
                encode_flag = true;
                i += 1;
            }
            "--transcode" => {
                transcode_flag = true;
                i += 1;
            }
            "--into" | "--to" => {
                if i + 1 >= args.len() {
                    return Err(format!("{} requires a value (target language or format)", args[i]));
                }
                pipeline_into = Some(args[i + 1].clone());
                i += 2;
            }
            "--from" => {
                if i + 1 >= args.len() {
                    return Err("--from requires a value (source language or format)".to_string());
                }
                pipeline_from = Some(args[i + 1].clone());
                i += 2;
            }
            "--from-ascii" => {
                if i + 1 >= args.len() {
                    return Err("--from-ascii requires a value (text string or '-' for stdin)".to_string());
                }
                let input = args[i + 1].clone();
                if input == "-" {
                    // Read from stdin
                    use std::io::{self, Read};
                    let mut buffer = String::new();
                    io::stdin().read_to_string(&mut buffer)
                        .map_err(|e| format!("Failed to read from stdin: {}", e))?;
                    ascii_input = Some(buffer);
                } else {
                    ascii_input = Some(input);
                }
                i += 2;
            }
            "--render-image" => {
                if i + 1 >= args.len() {
                    return Err("--render-image requires a value (text notation file or '-' for stdin)".to_string());
                }
                render_image = Some(args[i + 1].clone());
                i += 2;
            }
            "-o" | "--output" => {
                if i + 1 >= args.len() {
                    return Err("-o/--output requires a value (output file path)".to_string());
                }
                render_output = Some(args[i + 1].clone());
                i += 2;
            }
            arg if arg.starts_with("--") => {
                return Err(format!("Unknown option: {}", arg));
            }
            word => {
                // Split arguments that contain spaces or commas (e.g., "2 3 5" or "2,3,5")
                // This handles cases where users quote multiple words as a single argument
                if word.contains(' ') || word.contains(',') {
                    // Split by whitespace or comma, filter empty strings
                    let split_words: Vec<&str> = word
                        .split(|c: char| c.is_whitespace() || c == ',')
                        .filter(|s| !s.is_empty())
                        .collect();
                    words.extend(split_words.iter().map(|s| s.to_string()));
                } else {
                    words.push(word.to_string());
                }
                i += 1;
            }
        }
    }
    
    // Verb-flag pipeline: synthesize a meta instruction from --encode/--transcode/
    // --decode plus --from/--to and run it through the meta parser (which makes the
    // verb authoritative). `--decode` only drives a pipeline when an endpoint is
    // given; on its own it stays the legacy word→bytes flag.
    if encode_flag && transcode_flag {
        return Err("Use only one of --encode or --transcode.".to_string());
    }
    let verb_word = if encode_flag {
        Some("encode")
    } else if transcode_flag {
        Some("transcode")
    } else if decode_mode && (pipeline_from.is_some() || pipeline_into.is_some()) {
        Some("decode")
    } else {
        None
    };
    if let Some(vw) = verb_word {
        if meta_instruction.is_some() {
            return Err("Use either --meta or the --encode/--transcode/--decode flags, not both.".to_string());
        }
        if (encode_flag || transcode_flag) && pipeline_into.is_none() {
            return Err(format!("--{} requires a target: add --to <language|format>.", vw));
        }
        let mut s = String::from(vw);
        if let Some(ref f) = pipeline_from {
            s.push_str(" from ");
            s.push_str(f);
        }
        if let Some(ref t) = pipeline_into {
            s.push_str(" into ");
            s.push_str(t);
        }
        meta_instruction = Some(s);
    }

    if random_count.is_some() && !words.is_empty() {
        return Err("Cannot use --random with explicit words. Use one or the other.".to_string());
    }
    
    if random_count.is_some() && ascii_input.is_some() {
        return Err("Cannot use --random with --from-ascii. Use one or the other.".to_string());
    }
    
    if ascii_input.is_some() && !words.is_empty() {
        return Err("Cannot use --from-ascii with explicit words. Use one or the other.".to_string());
    }
    
    if random_count.is_none() && words.is_empty() && ascii_input.is_none() && !show_grammar && !decode_mode
        && meta_instruction.is_none() && pipeline_into.is_none() && pipeline_from.is_none()
        && render_image.is_none() {
        return Err("No words provided. Use --random <N>, --from-ascii <text>, --decode <words>, --meta, --into, --from, --render-image, or provide words as arguments.".to_string());
    }

    // Default length mode depends on grammar mode unless explicitly overridden:
    // - subject: compact (shortest-first)
    // - body: natural (sample from grammar length distribution)
    if !length_mode_explicit {
        length_mode = match generation_mode {
            GenerationMode::Subject => SentenceLengthMode::Compact,
            GenerationMode::Body => SentenceLengthMode::Natural,
            GenerationMode::PayloadOnly => SentenceLengthMode::Compact,
        };
    }
    
    Ok((words, random_count, ascii_input, verbose, seed, variations, best_of, highlight_mode, generation_mode, dialect_name, language, language_was_explicit, show_grammar, k_min, k_max, k_min_explicit, k_max_explicit, length_mode, width, delimiter, merkle_mode, demerkle_mode, decode_mode, merkle_highlight_mode, wordlist, meta_instruction, pipeline_into, pipeline_from, render_image, render_output))
}

/// Get the wordlist file path for a given language.
/// Returns the path to the payload.yaml file.

/// Find the languages directory by checking multiple locations
fn find_languages_dir() -> Option<String> {
    let probe = "languages/english/payload_bip39.yaml";
    // Try current directory first (for development)
    if std::path::Path::new(probe).exists() {
        return Some("languages".to_string());
    }

    // Try relative to the binary location (for installed binaries)
    if let Ok(exe_path) = std::env::current_exe() {
        if let Some(exe_dir) = exe_path.parent() {
            // Check if languages dir is sibling to bin dir (cargo install layout)
            let languages_path = exe_dir.join("../share/glossia/languages");
            if languages_path.join("english/payload_bip39.yaml").exists() {
                return Some(languages_path.to_string_lossy().to_string());
            }

            // Check if languages dir is in the same directory as binary
            let languages_path = exe_dir.join("languages");
            if languages_path.join("english/payload_bip39.yaml").exists() {
                return Some(languages_path.to_string_lossy().to_string());
            }

            // Check if we're in a cargo install location, go up to find languages
            // ~/.cargo/bin/glossia -> ~/.cargo/share/glossia/languages
            if exe_dir.ends_with("bin") {
                let share_path = exe_dir.join("../share/glossia/languages");
                if share_path.join("english/payload_bip39.yaml").exists() {
                    return Some(share_path.to_string_lossy().to_string());
                }
            }
        }
    }

    // Try home directory cargo share location
    if let Some(home) = std::env::var_os("HOME") {
        let cargo_share = std::path::Path::new(&home).join(".cargo/share/glossia/languages");
        if cargo_share.join("english/payload_bip39.yaml").exists() {
            return Some(cargo_share.to_string_lossy().to_string());
        }
    }

    // Try CARGO_HOME if set
    if let Ok(cargo_home) = std::env::var("CARGO_HOME") {
        let cargo_share = std::path::Path::new(&cargo_home).join("share/glossia/languages");
        if cargo_share.join("english/payload_bip39.yaml").exists() {
            return Some(cargo_share.to_string_lossy().to_string());
        }
    }

    None
}

/// Exits with error if payload file doesn't exist.
/// Only used for languages without embedded files (those use embedded files instead).
fn get_wordlist_path(language: &str, wordlist: &str) -> Result<String, String> {
    // Languages with embedded files should not call this function
    if has_embedded_files(language) {
        return Err(format!("Language '{}' uses embedded files and should not call get_wordlist_path", language));
    }

    let (payload_filename, _) = glossia::generator::wordlist_filenames(language, wordlist);

    // Try to find languages directory
    let languages_dir = find_languages_dir()
        .ok_or_else(|| format!("Could not find languages directory. Please ensure the 'languages' folder is accessible.\nTried: current directory, binary location, ~/.cargo/share/glossia/languages"))?;

    // Try to find payload file (supports subdirectories and recursive search)
    if let Some(payload_yaml) = find_language_file(language, &payload_filename) {
        return Ok(payload_yaml);
    }

    // Fallback: construct expected path for error message
    let expected_path = format!("{}/{}/{}", languages_dir, language, payload_filename);
    Err(format!("Wordlist file not found for language '{}'. Expected: {}\nOnly languages with a payload file are supported.",
                language, expected_path))
}

/// Find a language file (payload.yaml, cover.yaml, etc.) recursively
fn find_language_file(language: &str, filename: &str) -> Option<String> {
    let languages_dir = find_languages_dir()?;
    
    // First try exact path (supports subdirectories like "math/primes")
    let exact_path = format!("{}/{}/{}", languages_dir, language, filename);
    if std::path::Path::new(&exact_path).exists() {
        return Some(exact_path);
    }
    
    // If exact path doesn't exist, search recursively
    let languages_path = std::path::Path::new(&languages_dir);
    if let Ok(entries) = std::fs::read_dir(languages_path) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                if let Some(found) = find_language_file_recursive(&path, language, filename) {
                    return Some(found);
                }
            }
        }
    }
    
    None
}

/// Recursively search for a language file matching the language name
fn find_language_file_recursive(dir: &std::path::Path, language: &str, filename: &str) -> Option<String> {
    // Check if this directory matches the language name (last component)
    if let Some(dir_name) = dir.file_name().and_then(|n| n.to_str()) {
        // If language is just the directory name (e.g., "primes" matches "primes/")
        if dir_name == language {
            let file_path = dir.join(filename);
            if file_path.exists() {
                return Some(file_path.to_string_lossy().to_string());
            }
        }
        
        // If language contains slashes (e.g., "math/primes"), check if path ends with it
        if language.contains('/') {
            let lang_path = std::path::Path::new(language);
            if dir.ends_with(lang_path) {
                let file_path = dir.join(filename);
                if file_path.exists() {
                    return Some(file_path.to_string_lossy().to_string());
                }
            }
        }
    }
    
    // Recursively search subdirectories
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                if let Some(found) = find_language_file_recursive(&path, language, filename) {
                    return Some(found);
                }
            }
        }
    }
    
    None
}


/// Render text notation to SVG and write to a file.
///
/// Extracts color tokens from the text (CIELAB or hex format),
/// maps the dialect name to an SVG layout, and writes the SVG.
fn render_text_to_svg(text: &str, output_path: &str, dialect_name: &str, seed: Option<u64>) {
    use glossia::image_codec::render;
    use glossia::image_codec::svg::{self, SvgConfig, Layout};

    let hex_colors = render::extract_hex_colors(text);
    if hex_colors.is_empty() {
        eprintln!("No color tokens found in text notation");
        std::process::exit(1);
    }

    let layout = match dialect_name {
        "grid" | "patches" => Layout::Grid,
        "constellation" => Layout::Constellation,
        _ => Layout::Voronoi, // voronoi, mosaic, etc.
    };

    let cols = match dialect_name {
        "patches" => 4,
        "grid" => 8,
        _ => 8,
    };

    // "disk" is a Modal cover word in the text notation that triggers circular clipping.
    let disk = text.split_whitespace().any(|w| w == "disk");

    let config = SvgConfig {
        layout,
        seed: seed.unwrap_or(42),
        cols,
        circular: disk,
        color_scatter: true,
        ..Default::default()
    };

    let color_refs: Vec<&str> = hex_colors.iter().map(|s| s.as_str()).collect();
    let svg_content = svg::render_svg(&color_refs, &config);

    std::fs::write(output_path, &svg_content).unwrap_or_else(|e| {
        eprintln!("Error writing {}: {}", output_path, e);
        std::process::exit(1);
    });

    eprintln!("Rendered {} color cells as {:?} SVG to {}",
        hex_colors.len(), layout, output_path);
}

// ── Image banner subcommand ─────────────────────────────────────────
// glossia image encode --payload-hex <hex> --output banner.png
// glossia image decode --input banner.png [--expected-hex <hex>]
fn run_image_subcommand(args: &[String]) -> Result<(), String> {
    use glossia::image_codec::render::viridis_approx_curve;
    use glossia::image_codec::frame::BishopFrame;
    use glossia::image_codec::capacity::select_encoding_params;
    use glossia::image_codec::banner;
    use glossia::image_codec::color::Srgb;

    if args.is_empty() {
        return Err("Usage: glossia image encode|decode [options]".to_string());
    }

    match args[0].as_str() {
        "encode" => {
            let mut payload_hex: Option<String> = None;
            let mut output: Option<String> = None;
            let mut width = 1500usize;
            let mut height = 500usize;
            let mut nsym = 16usize;
            let mut seed = 42u64;

            let mut i = 1;
            while i < args.len() {
                match args[i].as_str() {
                    "--payload-hex" => {
                        i += 1;
                        payload_hex = Some(args.get(i).cloned()
                            .ok_or("--payload-hex requires a value")?);
                    }
                    "--output" | "-o" => {
                        i += 1;
                        output = Some(args.get(i).cloned()
                            .ok_or("--output requires a value")?);
                    }
                    "--width" => {
                        i += 1;
                        width = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--width requires a number")?;
                    }
                    "--height" => {
                        i += 1;
                        height = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--height requires a number")?;
                    }
                    "--nsym" => {
                        i += 1;
                        nsym = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--nsym requires a number")?;
                    }
                    "--seed" => {
                        i += 1;
                        seed = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--seed requires a number")?;
                    }
                    other => return Err(format!("Unknown option: {}", other)),
                }
                i += 1;
            }

            let hex = payload_hex.ok_or("--payload-hex is required")?;
            let output_path = output.ok_or("--output is required")?;

            // Parse hex
            let payload = parse_hex_payload(&hex)?;

            // Build curve
            let curve = viridis_approx_curve();
            let frame = BishopFrame::new(&curve, 500);

            // Select encoding params
            let params = select_encoding_params(&curve, &frame, 50)
                .ok_or("No valid encoding configuration found")?;

            eprintln!("Encoding {} bytes (N={}, epsilon={:.4}, nsym={})",
                payload.len(), params.n, params.epsilon, nsym);

            // Encode banner
            let encoded = banner::encode_banner(
                &payload, &curve, &frame,
                params.n, params.epsilon, nsym,
                width, height, seed, 10,
            )?;

            eprintln!("  Cells: {} payload + 1 header = {}",
                encoded.meta.n_payload_cells, encoded.meta.n_total_cells);
            eprintln!("  Bits/cell: {:.2}", encoded.meta.bits_per_cell);
            eprintln!("  RS: {} parity bytes (corrects up to {} byte errors)",
                encoded.meta.rs_parity_bytes, encoded.meta.max_correctable_bytes);

            // Render to PNG
            banner::render_banner_png(
                &encoded, &output_path, 2.5, Srgb::new(10, 10, 25),
            )?;

            eprintln!("  Saved: {}", output_path);
            Ok(())
        }
        "decode" => {
            let mut input: Option<String> = None;
            let mut expected_hex: Option<String> = None;
            let mut nsym = 16usize;
            let mut seed = 42u64;

            let mut i = 1;
            while i < args.len() {
                match args[i].as_str() {
                    "--input" | "-i" => {
                        i += 1;
                        input = Some(args.get(i).cloned()
                            .ok_or("--input requires a value")?);
                    }
                    "--expected-hex" => {
                        i += 1;
                        expected_hex = Some(args.get(i).cloned()
                            .ok_or("--expected-hex requires a value")?);
                    }
                    "--nsym" => {
                        i += 1;
                        nsym = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--nsym requires a number")?;
                    }
                    "--seed" => {
                        i += 1;
                        seed = args.get(i).and_then(|s| s.parse().ok())
                            .ok_or("--seed requires a number")?;
                    }
                    other => return Err(format!("Unknown option: {}", other)),
                }
                i += 1;
            }

            let input_path = input.ok_or("--input is required")?;

            let curve = viridis_approx_curve();
            let frame = BishopFrame::new(&curve, 500);

            eprintln!("Decoding {}...", input_path);
            let (recovered, meta) = banner::decode_banner_png(
                &input_path, &curve, &frame, nsym, seed, 10,
            )?;

            eprintln!("  N={}, epsilon={:.4}", meta.n_palette, meta.epsilon);
            eprintln!("  Cells decoded: {}", meta.n_cells);
            eprintln!("  Success: {}", meta.success);

            if meta.success {
                eprintln!("  Errors corrected: {}", meta.errors_corrected);
                let hex_str: String = recovered.iter().map(|b| format!("{:02x}", b)).collect();
                println!("{}", hex_str);

                if let Some(expected) = expected_hex {
                    let expected_bytes = parse_hex_payload(&expected)?;
                    if recovered == expected_bytes {
                        eprintln!("  Verification: PASS");
                    } else {
                        eprintln!("  Verification: FAIL");
                        eprintln!("  Expected: {}", expected);
                        std::process::exit(1);
                    }
                }
            } else {
                eprintln!("  Error: {}", meta.error_message.unwrap_or_default());
                std::process::exit(1);
            }
            Ok(())
        }
        other => Err(format!("Unknown image subcommand: {}. Use 'encode' or 'decode'.", other)),
    }
}

fn parse_hex_payload(hex: &str) -> Result<Vec<u8>, String> {
    let hex = hex.trim().trim_start_matches("0x").trim_start_matches("0X");
    if hex.len() % 2 != 0 {
        return Err("Hex string must have even length".to_string());
    }
    (0..hex.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16)
            .map_err(|_| format!("Invalid hex at position {}", i)))
        .collect()
}

// --- CLI usage ---
fn main() {
    // Check for "image" subcommand early, before complex arg parsing
    let args: Vec<String> = env::args().collect();
    if args.len() >= 2 && args[1] == "image" {
        if let Err(e) = run_image_subcommand(&args[2..]) {
            eprintln!("Error: {}", e);
            std::process::exit(1);
        }
        return;
    }

    let (mut words, random_count, ascii_input, verbose, seed, variations, best_of, highlight_mode, generation_mode, dialect_name, mut language, language_was_explicit, show_grammar, mut k_min, mut k_max, k_min_explicit, k_max_explicit, length_mode, width, delimiter, merkle_mode, demerkle_mode, decode_mode, merkle_highlight_mode, wordlist, meta_instruction, pipeline_into, pipeline_from, render_image, render_output) = match parse_args() {
        Ok(args) => args,
        Err(e) => {
            eprintln!("Error: {}", e);
            eprintln!();
            print_usage(&env::args().next().unwrap_or_else(|| "glossia".to_string()));
            std::process::exit(1);
        }
    };

    // ── Pipeline mode ────────────────────────────────────────────────────
    // --meta, --into, or --from flags route through Pipeline, bypassing
    // the traditional generation path.
    if meta_instruction.is_some() || pipeline_into.is_some() || pipeline_from.is_some() {
        use glossia::Pipeline;
        use std::io::Read;

        let pipeline = if let Some(ref meta) = meta_instruction {
            Pipeline::from_meta(meta).unwrap_or_else(|e| {
                eprintln!("Error parsing meta instruction: {}", e);
                std::process::exit(1);
            })
        } else {
            let source = match pipeline_from {
                Some(ref f) => parse_endpoint(f, &wordlist),
                None => glossia::Endpoint::Auto,
            };
            let target = match pipeline_into {
                Some(ref t) => parse_endpoint(t, &wordlist),
                None => glossia::Endpoint::Auto,
            };
            Pipeline::from_params(source, target)
        };

        let seed_value = seed.unwrap_or_else(|| rand::thread_rng().gen::<u64>());
        let pipeline = pipeline.with_seed(seed_value).with_verbose(verbose);

        if verbose {
            eprintln!("Pipeline: {} -> {}", pipeline.source, pipeline.target);
            if seed.is_some() {
                eprintln!("Using seed: {}", seed_value);
            }
        }

        // Get input: from --from-ascii, positional words, or stdin.
        let input = if let Some(ref ascii) = ascii_input {
            ascii.clone()
        } else if !words.is_empty() {
            words.join(" ")
        } else {
            let mut buffer = String::new();
            std::io::stdin().read_to_string(&mut buffer).unwrap_or_else(|e| {
                eprintln!("Error reading from stdin: {}", e);
                std::process::exit(1);
            });
            buffer.trim_end().to_string()
        };

        match pipeline.execute(&input) {
            Ok(output) => {
                println!("{}", glossia::generator::word_wrap(&output, width));
            }
            Err(e) => {
                eprintln!("Pipeline error: {}", e);
                std::process::exit(1);
            }
        }
        return;
    }

    // ── Render image mode ─────────────────────────────────────────────
    // --render-image takes a file of text notation and renders it to an image.
    // Output format is determined by file extension:
    //   .svg  → SVG (Voronoi, Grid, or Constellation layout)
    //   .png  → PNG grid (requires "native" feature)
    if let Some(ref input_path) = render_image {
        use glossia::image_codec::render;

        let text = if input_path == "-" {
            use std::io::Read;
            let mut buffer = String::new();
            std::io::stdin().read_to_string(&mut buffer).unwrap_or_else(|e| {
                eprintln!("Error reading from stdin: {}", e);
                std::process::exit(1);
            });
            buffer
        } else {
            std::fs::read_to_string(input_path).unwrap_or_else(|e| {
                eprintln!("Error reading {}: {}", input_path, e);
                std::process::exit(1);
            })
        };

        let output_path = render_output.as_deref().unwrap_or("output.svg");

        if output_path.ends_with(".svg") {
            render_text_to_svg(&text, output_path, &dialect_name, seed);
        } else if output_path.ends_with(".png") {
            let n_palette = 64;
            let cell_size = 32;
            let cols = 8;
            match render::render_text_to_png(&text, output_path, n_palette, cell_size, cols) {
                Ok(()) => {
                    let n = render::extract_hex_colors(&text).len();
                    eprintln!("Rendered {} color cells to {}", n, output_path);
                }
                Err(e) => {
                    eprintln!("Render error: {}", e);
                    std::process::exit(1);
                }
            }
        } else {
            eprintln!("Unknown output format: {}. Use .svg or .png extension.", output_path);
            std::process::exit(1);
        }
        return;
    }

    // Build DialectConfig to get dialect-specified wordlists.
    // The dialect knows its preferred payload/cover wordlists (from grammar.yaml).
    let dialect_config = {
        glossia::DialectConfig::from_language_dialect(&language, &dialect_name)
            .unwrap_or_else(|e| {
                eprintln!("Error loading dialect config for '{}' / '{}': {}", language, dialect_name, e);
                std::process::exit(1);
            })
    };

    // Resolve wordlist: CLI --wordlist overrides dialect default.
    // "default" means "use whatever the dialect specifies".
    let wordlist_was_explicit = wordlist != "default";
    let mut wordlist = if wordlist_was_explicit {
        wordlist
    } else {
        // Use the dialect's declared payload wordlist
        let dialect_wl = dialect_config.payload_wordlist().to_string();
        if dialect_wl == "default" {
            // "default" in the dialect means payload.yaml — resolve to language's primary profile
            glossia::generator::default_wordlist(&language).to_string()
        } else {
            dialect_wl
        }
    };

    // Apply the final wordlist choice to the dialect config (for cover resolution)
    let dialect_config = if wordlist_was_explicit {
        dialect_config.with_payload_wordlist(&wordlist)
    } else {
        dialect_config
    };

    // Handle --demerkle mode: parse merkleized sequence and output payload words
    if demerkle_mode {
        use glossia::merkle::parse_merkleized;
        use glossia::generator::data::load_payload_tree;
        
        let payload_tree = match load_payload_tree(&language) {
            Ok(tree) => tree,
            Err(e) => {
                eprintln!("Error loading payload tree: {}", e);
                std::process::exit(1);
            }
        };
        
        // If words provided, parse them; otherwise read from stdin
        let sequence = if !words.is_empty() {
            words
        } else {
            use std::io::{self, Read};
            let mut buffer = String::new();
            io::stdin().read_to_string(&mut buffer)
                .map_err(|e| {
                    eprintln!("Error reading from stdin: {}", e);
                    std::process::exit(1);
                })
                .unwrap();
            buffer.split_whitespace().map(|s| s.to_string()).collect()
        };
        
        match parse_merkleized(&sequence, &payload_tree) {
            Ok(payload_words) => {
                println!("{}", payload_words.join(" "));
            }
            Err(e) => {
                eprintln!("Error parsing merkleized sequence: {}", e);
                std::process::exit(1);
            }
        }
        return;
    }

    // Handle --decode mode: decode payload words back to bytes
    if decode_mode {
        // Collect raw input tokens first (before any dialect-specific processing)
        let raw_tokens: Vec<String> = if !words.is_empty() {
            words.clone()
        } else {
            use std::io::{self, Read};
            let mut buffer = String::new();
            io::stdin().read_to_string(&mut buffer)
                .map_err(|e| {
                    eprintln!("Error reading from stdin: {}", e);
                    std::process::exit(1);
                })
                .unwrap();
            buffer.split_whitespace().map(|s| s.to_string()).collect()
        };

        // Auto-detect dialect when neither --language nor --wordlist was explicitly set
        let wordlist_was_explicit = wordlist != glossia::generator::default_wordlist(&language);
        if !language_was_explicit && !wordlist_was_explicit {
            let dialect_matches = glossia::generator::detect_dialect(&raw_tokens);
            if let Some(best) = dialect_matches.first() {
                language = best.language.clone();
                wordlist = best.wordlist.clone();

                eprintln!("Detected dialect: {}/{} ({}/{} words matched, {:.1}%)",
                    best.language, best.wordlist,
                    best.hits, best.total, best.hit_rate * 100.0);
                if best.dialects.len() > 1 {
                    eprintln!("  Available grammars: {}", best.dialects.join(", "));
                }

                // Show runner-up candidates in verbose mode
                if verbose && dialect_matches.len() > 1 {
                    eprintln!("  Other candidates:");
                    for candidate in dialect_matches.iter().skip(1).take(3) {
                        eprintln!("    {}", candidate);
                    }
                }
            }
        } else if verbose {
            eprintln!("Using explicit dialect: {}/{}", language, wordlist);
        }

        let all_words = match glossia::generator::load_payload_words_for_wordlist(&language, &wordlist) {
            Ok(w) => w,
            Err(e) => {
                eprintln!("Error loading wordlist: {}", e);
                std::process::exit(1);
            }
        };
        let tree = glossia::WordlistTree::new(all_words.clone());

        // Load grammar to check payload separator
        let grammar = Grammar::from_language_dialect(&language, "body");
        let payload_separator = grammar.as_ref()
            .map(|g| g.payload_separator().to_string())
            .unwrap_or_else(|_| " ".to_string());

        // Process raw tokens according to payload separator rules
        let input_words: Vec<String> = if payload_separator.is_empty() {
            // For concatenated payloads, extract chars only from tokens where
            // ALL characters are in the payload set (i.e. pure payload blocks).
            // Tokens with mixed characters (like "BEGIN") are cover words — skip them.
            let payload_set: std::collections::HashSet<String> = all_words.iter()
                .map(|w| w.to_lowercase())
                .collect();
            raw_tokens.iter()
                .flat_map(|w| {
                    let trimmed = w.trim_matches(|c: char| {
                        !c.is_alphanumeric() && !payload_set.contains(&c.to_lowercase().to_string())
                    });
                    let all_in_payload = !trimmed.is_empty() && trimmed.chars()
                        .all(|c| payload_set.contains(&c.to_lowercase().to_string()));
                    if all_in_payload {
                        trimmed.chars()
                            .map(|c| c.to_lowercase().to_string())
                            .collect::<Vec<_>>()
                    } else {
                        vec![]  // Skip cover words
                    }
                })
                .collect()
        } else {
            // Standard whitespace-delimited: strip punctuation and filter to payload words only.
            // Cover words must be removed before the codec sees the sequence.
            raw_tokens.iter()
                .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()).to_lowercase())
                .filter(|w| !w.is_empty() && tree.contains(w))
                .collect()
        };

        // Calculate input bits from payload word count
        let wordlist_size = all_words.len();
        let codec_str = grammar.as_ref()
            .map(|g| g.codec().to_string())
            .unwrap_or_else(|_| "bitpack".to_string());
        if verbose {
            eprintln!("Decoding {} words", input_words.len());
            if let Some(input_bits) = calculate_output_bits(input_words.len(), wordlist_size, &codec_str) {
                let bits_per_word = wordlist_size.trailing_zeros();
                if codec_str == "bitpack" && wordlist_size > 1 && wordlist_size.is_power_of_two() && input_words.len() > 0 {
                    eprintln!("Input bits: {} ({} data words × {} bits/word, +1 padding word)",
                             input_bits, input_words.len() - 1, bits_per_word);
                } else {
                    eprintln!("Input bits: {} ({} words × {} bits/word)",
                             input_bits, input_words.len(), bits_per_word);
                }
            }
        }

        // Codec-controlled decoding: decode to bytes, then detect mode by inspection.
        let decode_result = glossia::codec::decode_base_n(&input_words, &tree, &codec_str)
            .map(|bytes| {
                let mode = if std::str::from_utf8(&bytes).is_ok() {
                    glossia::DataMode::Ascii7  // valid UTF-8: print as text
                } else {
                    glossia::DataMode::Hex     // binary: re-encode as hex string
                };
                (mode, bytes)
            });

        match decode_result {
            Ok((mode, bytes)) => {
                use std::io::Write;
                // For text modes (hex, base64, ascii), reconstruct the original string.
                // For raw bytes mode, write bytes directly.
                match mode {
                    glossia::DataMode::Hex => {
                        let hex = glossia::hex_encode(&bytes);
                        print!("{}", hex);
                        println!();
                    }
                    glossia::DataMode::Base64 => {
                        let b64 = glossia::base64_encode(&bytes);
                        print!("{}", b64);
                        println!();
                    }
                    glossia::DataMode::Ascii7 => {
                        std::io::stdout().write_all(&bytes).unwrap_or_else(|e| {
                            eprintln!("Error writing output: {}", e);
                            std::process::exit(1);
                        });
                        println!();
                    }
                    glossia::DataMode::Bytes8 => {
                        std::io::stdout().write_all(&bytes).unwrap_or_else(|e| {
                            eprintln!("Error writing output: {}", e);
                            std::process::exit(1);
                        });
                        // Add trailing newline if output looks like text
                        if bytes.iter().all(|&b| b.is_ascii()) {
                            println!();
                        }
                    }
                }
                if verbose {
                    eprintln!("Decoded mode: {}", mode);
                    eprintln!("Output bytes: {}", bytes.len());
                    eprintln!("Output bits: {} bytes × 8 = {} bits", bytes.len(), bytes.len() * 8);
                }
            }
            Err(e) => {
                eprintln!("Error decoding: {}", e);
                std::process::exit(1);
            }
        }
        return;
    }

    // If --show-grammar flag is set, display grammar and continue
    if show_grammar {
        let dialect = match generation_mode {
            GenerationMode::Subject => "subject",
            GenerationMode::Body => "body",
            GenerationMode::PayloadOnly => "payload_only",
        };
        let grammar = Grammar::from_language_dialect(&language, dialect);
        
        match grammar {
            Ok(g) => {
                let mode_str = match generation_mode {
                    GenerationMode::Subject => "subject",
                    GenerationMode::Body => "body",
                    GenerationMode::PayloadOnly => "payload_only",
                };
                println!("Grammar: {}", mode_str);
                let grammar_output = g.format_concise();
                print!("{}", grammar_output);
                let start_symbols = if generation_mode == GenerationMode::Body {
                    vec!["S"]
                } else {
                    vec!["S", "S_N", "S_V", "S_Adj", "S_Adv", "S_Prep", "S_Det"]
                };
                let mut by_start_symbol = HashMap::new();
                for start_symbol in start_symbols {
                    let sequences_by_k = g.precompute_sequences_with_probability(start_symbol, k_max);
                    if !sequences_by_k.is_empty() {
                        by_start_symbol.insert(start_symbol.to_string(), sequences_by_k);
                    }
                }
                print_sentence_kinds_once(mode_str, k_max, &by_start_symbol, true);
                println!(); // Add blank line after grammar
            }
            Err(e) => {
                eprintln!("Error loading grammar: {}", e);
                std::process::exit(1);
            }
        }
    }

    // Use seeded RNG if seed provided, otherwise generate random seed from thread_rng
    let seed_value = if let Some(s) = seed {
        s
    } else {
        // Generate a random seed from thread_rng for non-deterministic behavior
        rand::thread_rng().gen::<u64>()
    };
    let mut rng = StdRng::seed_from_u64(seed_value);
    
    if verbose {
        if seed.is_some() {
            eprintln!("Using seed: {}", seed_value);
        }
        let mode_str = match generation_mode {
            GenerationMode::Subject => "subject",
            GenerationMode::Body => "body",
            GenerationMode::PayloadOnly => "payload_only",
        };
        eprintln!("Mode: {}", mode_str);
        eprintln!("Language: {}", language);
        eprintln!("Dialect: {} (payload_wordlist={}, cover_wordlist={})",
            dialect_config.dialect(), dialect_config.payload_wordlist(), dialect_config.cover_wordlist());
        eprintln!("Wordlist: {}{}", wordlist,
            if wordlist_was_explicit { " (--wordlist override)" } else { "" });
    }

    // Validate language early (languages with embedded files don't need filesystem check)
    if has_embedded_files(&language) {
        if verbose {
            eprintln!("Using embedded wordlist for language: {}", language);
        }
    } else {
        match get_wordlist_path(&language, &wordlist) {
            Ok(path) => {
                if verbose {
                    eprintln!("Using wordlist: {}", path);
                }
            }
            Err(e) => {
                eprintln!("Error: {}", e);
                std::process::exit(1);
            }
        };
    }

    // Track data mode and input text for bit calculations
    let mut data_mode: Option<glossia::DataMode> = None;
    let mut input_text: Option<String> = None;

    // If ASCII input provided, encode it to words
    if let Some(ascii_text) = ascii_input {
        input_text = Some(ascii_text.clone());
        words = match encode_ascii_to_words(&ascii_text, &language, &wordlist, dialect_config.dialect()) {
            Ok((encoded_words, mode)) => {
                data_mode = Some(mode);
                if verbose {
                    eprintln!("Encoded {} bytes to {} words (mode: {})", ascii_text.len(), encoded_words.len(), mode);
                }
                encoded_words
            }
            Err(e) => {
                eprintln!("Error encoding ASCII to words: {}", e);
                std::process::exit(1);
            }
        };
    }
    
    // Load wordlist once (needed for random selection and validation)
    let wordlist_words = match glossia::generator::load_payload_words_for_wordlist(&language, &wordlist) {
        Ok(words) => words,
        Err(e) => {
            eprintln!("Error loading wordlist: {}", e);
            std::process::exit(1);
        }
    };
    let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();
    
    // If random words requested, select them now (using already-loaded wordlist)
    if let Some(count) = random_count {
        if wordlist_words.is_empty() || count == 0 {
            words = Vec::new();
        } else {
            // Sample WITH replacement so duplicates are possible (and therefore decodable).
            let mut selected = Vec::with_capacity(count);
            for _ in 0..count {
                selected.push(wordlist_words.choose(&mut rng).unwrap().clone());
            }
            words = selected;
            if verbose {
                eprintln!("Selected {} random payload words: {}", count, words.join(" "));
            }
        }
    }

    // Tag each word with POS tags using language-specific mapping
    let pos_mapping = glossia::generator::build_pos_mapping_for_wordlist(&language, &wordlist)
        .unwrap_or_else(|e| {
            eprintln!("Warning: Failed to load POS mapping for language '{}': {}", language, e);
            eprintln!("Falling back to English POS mapping");
            glossia::generator::build_pos_mapping("english").unwrap_or_else(|_| HashMap::new())
        });
    
    // Handle merkle mode: expand payload with Merkle proof
    let (expanded_payload, merkle_words_set, original_payload_set) = if merkle_mode {
        use glossia::merkle::merkleize;
        use glossia::generator::data::{load_cover_tree, load_cover_word_pos_tags};
        
        if words.len() < 2 {
            eprintln!("Error: --merkle requires at least 2 payload words (got {})", words.len());
            std::process::exit(1);
        }
        
        let cover_tree = load_cover_tree(&language);
        let cover_pos_tags = load_cover_word_pos_tags(&language);
        
        // Track original payload words for highlighting (only these should be highlighted)
        let original_payload: HashSet<String> = words.iter().map(|w| w.to_lowercase()).collect();
        
        // Merkleize the payload words
        let merkle_result = match merkleize(&words, &cover_tree) {
            Ok(result) => result,
            Err(e) => {
                eprintln!("Error merkleizing payload: {}", e);
                std::process::exit(1);
            }
        };
        
        if verbose {
            eprintln!("Merkleized {} payload words to {} total words ({} Merkle nodes)", 
                     merkle_result.n_leaves, 
                     merkle_result.sequence.len(),
                     merkle_result.n_merkle);
        }
        
        // Convert merkleized sequence to PayloadTok with POS tags
        let mut expanded: Vec<PayloadTok> = Vec::new();
        let mut merkle_set: HashSet<String> = HashSet::new();
        
        // Track which cover words are Merkle words (first N-1)
        for i in 0..merkle_result.n_merkle {
            if let Some(merkle_word) = cover_tree.get(i) {
                merkle_set.insert(merkle_word.to_lowercase());
            }
        }
        
        for word in &merkle_result.sequence {
            let word_lower = word.to_lowercase();
            // Check if this is a Merkle word (cover word) or payload word
            let tags = if merkle_set.contains(&word_lower) {
                // Merkle word: use POS tags from cover.yaml
                cover_pos_tags.get(&word_lower).cloned().unwrap_or_default()
            } else {
                // Payload word: use POS tags from payload.yaml
                pos_mapping.get(&word_lower).cloned().unwrap_or_default()
            };
            expanded.push(PayloadTok::new(word.clone(), &tags));
        }
        
        (expanded, merkle_set, original_payload)
    } else {
        // Normal mode: just tag the payload words
        let payload: Vec<PayloadTok> = words
            .iter()
            .map(|word| {
                let word_lower = word.to_lowercase();
                let tags = pos_mapping.get(&word_lower).cloned().unwrap_or_default();
                PayloadTok::new(word.clone(), &tags)
            })
            .collect();
        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        (payload, HashSet::new(), payload_set)
    };
    
    // For highlighting: only highlight original payload words, not Merkle words
    let payload_set: HashSet<String> = original_payload_set.clone();
    let payload_set_clone = payload_set.clone(); // Keep a copy for later statistics
    
    // Load cover words with explicit POS tags from cover.yaml
    // Uses the dialect config's cover_wordlist, which may differ from the payload wordlist.
    // No fallback defaults — cover.yaml must provide all needed POS categories.
    // Disjointness with payload is validated at build time.
    let cover_wl = dialect_config.cover_wordlist();
    let (cover_by_pos, refined_cover) = glossia::generator::load_cover_words_by_pos_for_wordlist(&wordlist_set, &language, cover_wl);

    if verbose {
        eprintln!("Loaded cover words from cover.yaml:");
        for (pos, words) in &cover_by_pos {
            eprintln!("  {:?}: {}", pos, words.len());
        }
    }

    // Build refined_payload map from dialect's scale definition (if any).
    // Maps refinement tags to valid payload word sets for runtime validation.
    let refined_payload = {
        let mut rp = std::collections::HashMap::new();
        if let Some(ref_tag) = dialect_config.n_refinement_tag() {
            // The payload_set already contains the scale-filtered words
            rp.insert(ref_tag, payload_set.clone());
        }
        rp
    };

    // Build Lexicon from all POS categories in cover_by_pos.
    // This handles any POS the grammar uses (including Dot for CS grammar).
    let lex = if merkle_mode {
        // Filter out Merkle words from each POS category
        let mut lex = Lexicon::new(payload_set.clone(), wordlist_set.clone());
        for (pos, words) in &cover_by_pos {
            let filtered: Vec<String> = words.iter()
                .filter(|w| !merkle_words_set.contains(&w.to_lowercase()))
                .cloned()
                .collect();
            lex = lex.with_words(*pos, &filtered.iter().map(|s| s.as_str()).collect::<Vec<_>>());
        }
        lex.with_refined_cover(refined_cover.clone())
            .with_refined_payload(refined_payload)
    } else {
        let mut lex = Lexicon::new(payload_set.clone(), wordlist_set.clone());
        for (pos, words) in &cover_by_pos {
            lex = lex.with_words(*pos, &words.iter().map(|s| s.as_str()).collect::<Vec<_>>());
        }
        lex.with_refined_cover(refined_cover.clone())
            .with_refined_payload(refined_payload)
    };

    // Attach the optional semantic model (softly biases sentence planning toward
    // coherent verb-argument pairings; absent for languages without semantics.yaml).
    let lex = match glossia::generator::data::load_semantics(&language) {
        Some(model) => lex.with_semantics(std::sync::Arc::new(model)),
        None => lex,
    };

    let input_word_count = if merkle_mode {
        words.len()  // Original payload word count
    } else {
        expanded_payload.len()
    };
    // For validation: only check original payload words, not Merkle words
    let expected_words: Vec<String> = words.iter().map(|w| w.to_lowercase()).collect();

    // Calculate input statistics by POS
    let mut input_pos_counts: HashMap<Pos, usize> = HashMap::new();
    for tok in &expanded_payload {
        for pos in &tok.allowed {
            *input_pos_counts.entry(*pos).or_insert(0) += 1;
        }
    }

    // Load grammar for payload format (needed for validation of concatenated payloads)
    let encode_grammar = Grammar::from_language_dialect(&language, &dialect_name).ok();
    let encode_payload_separator = encode_grammar.as_ref()
        .map(|g| g.payload_separator().to_string())
        .unwrap_or_else(|| " ".to_string());

    // Compute dynamic k_min/k_max from grammar when not explicitly set by user.
    // Concatenated-payload grammars (payload_separator="") encode all payload
    // in a single sentence — force k_min = k_max = overhead + payload_count.
    let concat_payload = encode_grammar.as_ref()
        .map(|g| g.payload_separator().is_empty())
        .unwrap_or(false);
    if let Some(ref grammar) = encode_grammar {
        let min_k = grammar.min_sentence_length().unwrap_or(5);
        if !k_min_explicit && !k_max_explicit && concat_payload {
            let k = min_k + expanded_payload.len().saturating_sub(1);
            k_min = k;
            k_max = k;
        }
        // Respect grammar's max_k declaration (caps enumeration for recursive grammars)
        if !k_max_explicit {
            if let Some(grammar_max_k) = grammar.max_k() {
                if k_max > grammar_max_k {
                    k_max = grammar_max_k;
                }
            }
        }
        if verbose {
            eprintln!("Grammar min sentence length: {}", min_k);
            eprintln!("Using k_min={}, k_max={}", k_min, k_max);
        }
    }

    // Generate multiple variations and select the most compact one
    let mut best_text: Option<String> = None;
    let mut best_compactness = 0.0;
    let mut best_output_count = 0;
    let mut variation_stats: Vec<f64> = Vec::new();
    let mut valid_variation_texts: Vec<String> = Vec::new();

    if verbose && variations > 1 {
        eprintln!("Generating {} variations to maximize compactness...", variations);
    }

    for variation in 0..variations {
        // Space each variation's seed range by best_of so best-of-N candidate
        // seeds don't overlap between variations.
        let variation_seed = seed_value.wrapping_add((variation as u64).wrapping_mul(best_of.max(1) as u64));

        // Call library's generate_text (without highlighting). With --best-of N,
        // sample N candidates and keep the densest / most coherent (see
        // generate_text_best_of); otherwise a single generation.
        // Pass original_payload_set so it can validate sentences contain actual payload words
        let (text_unhighlighted, _payload_set_from_gen) = if best_of > 1 {
            generate_text_best_of(
                variation_seed,
                best_of,
                &lex,
                &expanded_payload,
                Some(&original_payload_set),
                variations == 1 && verbose,
                generation_mode,
                &language,
                Some(&dialect_name),
                k_min,
                k_max,
                length_mode,
                &delimiter,
            )
        } else {
            let mut variation_rng = StdRng::seed_from_u64(variation_seed);
            generate_text_with_original_payload(
                &mut variation_rng,
                &lex,
                &expanded_payload,
                Some(&original_payload_set),
                variations == 1 && verbose,
                generation_mode,
                &language,
                Some(&dialect_name),
                k_min,
                k_max,
                length_mode,
                &delimiter,
            )
        };
        
        // Apply highlighting based on highlight_mode
        // Use payload_set (original payload words only, not Merkle words) for highlighting
        // If --highlight-merkle is set, also highlight Merkle words with the specified mode
        let merkle_set_for_highlighting = if merkle_highlight_mode.is_some() && merkle_mode {
            Some(&merkle_words_set)
        } else {
            None
        };
        let text = apply_highlighting(&text_unhighlighted, &payload_set, &expanded_payload, highlight_mode, merkle_set_for_highlighting, merkle_highlight_mode);
        
        // Validate that the generated text contains exactly the input BIP39 words in order
        // Skip validation in madlib mode since words are replaced with [POS] placeholders
        // Skip validation in payload-only mode since it's a direct pass-through
        if highlight_mode != HighlightMode::Madlib && generation_mode != GenerationMode::PayloadOnly {
            let extracted_wordlist_words: Vec<String> = if encode_payload_separator.is_empty() {
                // For concatenated payloads (CS grammar), extract chars only from tokens
                // where ALL characters are in the payload set (pure payload blocks).
                text.split_whitespace()
                    .flat_map(|token| {
                        let trimmed = token.trim_matches(|c: char| {
                            !c.is_alphanumeric() && !payload_set.contains(&c.to_lowercase().to_string())
                        });
                        let all_in_payload = !trimmed.is_empty() && trimmed.chars()
                            .all(|c| payload_set.contains(&c.to_lowercase().to_string()));
                        if all_in_payload {
                            trimmed.chars()
                                .map(|c| c.to_lowercase().to_string())
                                .collect::<Vec<_>>()
                        } else {
                            vec![]  // Skip cover words
                        }
                    })
                    .collect()
            } else {
                text
                    .split_whitespace()
                    .map(normalize_token_for_bip39)
                    .filter(|w| !w.is_empty() && payload_set.contains(w))
                    .collect()
            };
            
            if extracted_wordlist_words != expected_words {
                if verbose || variations == 1 {
                    eprintln!("Variation {}: ERROR - Generated BIP39 words do not match input words!", variation + 1);
                    eprintln!("  Expected: {:?}", expected_words);
                    eprintln!("  Got:      {:?}", extracted_wordlist_words);
                    eprintln!("  Text:     {}", text.chars().take(500).collect::<String>());
                    eprintln!("  Payload set size: {}", payload_set.len());
                    eprintln!("  Sample payload words: {:?}", payload_set.iter().take(5).collect::<Vec<_>>());
                    // Debug: show what words are being extracted
                    let all_normalized: Vec<String> = text
                        .split_whitespace()
                        .map(normalize_token_for_bip39)
                        .filter(|w| !w.is_empty())
                        .collect();
                    eprintln!("  All normalized words (first 20): {:?}", all_normalized.iter().take(20).collect::<Vec<_>>());
                    let matching_words: Vec<String> = all_normalized
                        .iter()
                        .filter(|w| payload_set.contains(*w))
                        .cloned()
                        .collect();
                    eprintln!("  Matching words: {:?}", matching_words);
                }
                continue;
            }
        }

        // Save this valid variation for optional printing
        valid_variation_texts.push(text.clone());
        
        // Calculate compactness score for this variation (based on character counts)
        // In payload-only mode, split by delimiter instead of whitespace
        let words: Vec<&str> = if matches!(generation_mode, GenerationMode::PayloadOnly) {
            text.split(&delimiter).collect()
        } else {
            text.split_whitespace().collect()
        };
        let output_word_count = words.len();
        let mut bip39_chars = 0;
        let mut non_bip39_chars = 0;
        
        // In payload-only mode, all words are payload words
        if matches!(generation_mode, GenerationMode::PayloadOnly) {
            // In payload-only mode, compactness is always 1.0 (all payload words)
            let compactness = 1.0;
            variation_stats.push(compactness);
            
            if verbose && variations > 1 {
                eprintln!("Variation {}: compactness {:.3} (all payload words)", 
                          variation + 1, compactness);
            }
            
            // Always use first variation in payload-only mode (they're all the same)
            if best_text.is_none() {
                best_compactness = compactness;
                best_text = Some(text);
                best_output_count = output_word_count;
            }
        } else {
            // In madlib mode, count [POS] placeholders as BIP39 chars
            if highlight_mode == HighlightMode::Madlib {
                for word in &words {
                    if word.starts_with('[') && word.contains(']') {
                        // Count [POS] placeholder as BIP39 (approximate length)
                        bip39_chars += 5; // Approximate: "[N]" = 3 chars, but use 5 for weighting
                    } else {
                        let normalized = normalize_token_for_bip39(word);
                        non_bip39_chars += normalized.chars().count();
                    }
                }
            } else if encode_payload_separator.is_empty() {
                // For concatenated payloads (CS grammar), count characters individually
                for word in &words {
                    let trimmed = word.trim_matches(|c: char| {
                        !c.is_alphanumeric() && !payload_set.contains(&c.to_lowercase().to_string())
                    });
                    for ch in trimmed.chars() {
                        let ch_str = ch.to_lowercase().to_string();
                        if payload_set.contains(&ch_str) {
                            bip39_chars += 1;
                        } else {
                            non_bip39_chars += 1;
                        }
                    }
                }
            } else {
                for word in &words {
                    let normalized = normalize_token_for_bip39(word);
                    if !normalized.is_empty() && payload_set.contains(&normalized) {
                        // Count characters in BIP39 word (use normalized length)
                        bip39_chars += normalized.chars().count();
                    } else {
                        // Count characters in non-BIP39 word (use normalized length, excluding punctuation)
                        non_bip39_chars += normalized.chars().count();
                    }
                }
            }
            
            let total_chars = bip39_chars + non_bip39_chars;
            let compactness = if total_chars > 0 {
                bip39_chars as f64 / total_chars as f64
            } else {
                0.0
            };

            variation_stats.push(compactness);

            if verbose && variations > 1 {
                eprintln!("Variation {}: compactness {:.3} ({} BIP39 chars / {} total chars)", 
                          variation + 1, compactness, bip39_chars, total_chars);
            }

            // Keep track of the best (most compact) variation
            if compactness >= best_compactness || best_text.is_none() {
                best_compactness = compactness;
                best_text = Some(text);
                best_output_count = output_word_count;
            }
        }
    }

    // Output variations (if requested) or the best variation (default)
    let text = match best_text {
        Some(t) => t,
        None => {
            eprintln!("Error: Failed to generate any valid variations after {} attempts.", variations);
            eprintln!("This may happen if:");
            eprintln!("  - The grammar cannot accommodate all input words");
            eprintln!("  - There are POS tagging issues with some words");
            eprintln!("  - The word extraction logic is failing");
            if highlight_mode == HighlightMode::Madlib {
                eprintln!("  Note: Validation is skipped in madlib mode, so this error should not occur.");
            }
            eprintln!("\nTry running with --verbose to see detailed error messages.");
            std::process::exit(1);
        }
    };
    
    // ── Image rendering: if -o <file.svg|file.png> with image language ──
    // When the image language is used with -o, render directly to an image
    // file instead of printing text notation. The text notation is an
    // intermediate representation — the final output is the image.
    if let Some(ref output_path) = render_output {
        if language == "image" {
            if output_path.ends_with(".svg") {
                render_text_to_svg(&text, output_path, &dialect_name, seed);
                // Also print the text notation to stdout for piping/debugging
                if verbose {
                    eprintln!("Text notation:\n{}", &text);
                }
                return;
            } else if output_path.ends_with(".png") {
                use glossia::image_codec::render;
                let n_palette = 128;
                let cell_size = 32;
                let cols = 8;
                match render::render_text_to_png(&text, output_path, n_palette, cell_size, cols) {
                    Ok(()) => {
                        let n = render::extract_hex_colors(&text).len();
                        eprintln!("Rendered {} color cells to {}", n, output_path);
                    }
                    Err(e) => {
                        eprintln!("Render error: {}", e);
                        std::process::exit(1);
                    }
                }
                if verbose {
                    eprintln!("Text notation:\n{}", &text);
                }
                return;
            }
            // Non-image extension: fall through to normal text output
        }
    }

    // Word wrap the output to specified width
    if variations > 1 {
        for (i, vtext) in valid_variation_texts.iter().enumerate() {
            if i > 0 {
                println!();
                println!();
            }
            println!("{}", word_wrap(vtext, width));
        }
    } else {
        println!("{}", word_wrap(&text, width));
    }

    // Calculate detailed statistics from the best text
    let sentences: Vec<&str> = text.split('.').filter(|s| !s.trim().is_empty()).collect();
    let sentence_count = sentences.len();
    let avg_words_per_sentence = if sentence_count > 0 {
        best_output_count as f64 / sentence_count as f64
    } else {
        0.0
    };

    // Count payload vs cover words and characters
    let payload_words_in_output: HashSet<String> = payload_set_clone.iter().cloned().collect();
    let words: Vec<&str> = text.split_whitespace().collect();
    let mut payload_word_count = 0;
    let mut cover_word_count = 0;
    let mut bip39_char_count = 0;
    let mut non_bip39_char_count = 0;
    
    for word in &words {
        let normalized = normalize_token_for_bip39(word);
        if !normalized.is_empty() && payload_words_in_output.contains(&normalized) {
            payload_word_count += 1;
            bip39_char_count += normalized.chars().count();
        } else {
            cover_word_count += 1;
            non_bip39_char_count += normalized.chars().count();
        }
    }
    
    let total_output_chars = bip39_char_count + non_bip39_char_count;

    // Print comprehensive statistics (only when generating multiple variations)
    if variations > 1 {
        eprintln!("\n=== Statistics ===");
        eprintln!("Input:");
        eprintln!("  Total words: {}", input_word_count);

        // Report input bits if we have the information
        if let (Some(txt), Some(mode)) = (&input_text, data_mode) {
            let input_bits = calculate_input_bits(txt, mode);
            eprintln!("  Input bits: {} (mode: {})", input_bits, mode);
        }

        if !input_pos_counts.is_empty() {
            eprintln!("  POS breakdown:");
            let mut pos_vec: Vec<_> = input_pos_counts.iter().collect();
            pos_vec.sort_by_key(|(pos, _)| {
                match pos {
                    Pos::N => 1,
                    Pos::V => 2,
                    Pos::Adj => 3,
                    Pos::Adv => 4,
                    Pos::Prep => 5,
                    Pos::Det => 6,
                    _ => 7,
                }
            });
            for (pos, count) in pos_vec {
                let pos_name = pos.description();
                eprintln!("    {}: {}", pos_name, count);
            }
        }

        eprintln!("Output:");
        eprintln!("  Total words: {}", best_output_count);
        eprintln!("  Payload words: {} ({:.1}%)", payload_word_count,
                  (payload_word_count as f64 / best_output_count as f64) * 100.0);
        eprintln!("  Cover words: {} ({:.1}%)", cover_word_count,
                  (cover_word_count as f64 / best_output_count as f64) * 100.0);

        // Report output bits based on payload word count and wordlist size
        let wordlist_size = wordlist_words.len();
        let encode_codec = Grammar::from_language_dialect(&language, "body")
            .map(|g| g.codec().to_string())
            .unwrap_or_else(|_| "bitpack".to_string());
        if let Some(output_bits) = calculate_output_bits(payload_word_count, wordlist_size, &encode_codec) {
            let bits_per_word = wordlist_size.trailing_zeros();
            if encode_codec == "bitpack" && wordlist_size > 1 && wordlist_size.is_power_of_two() && payload_word_count > 0 {
                eprintln!("  Output bits: {} ({} data words × {} bits/word, +1 padding word)",
                         output_bits, payload_word_count - 1, bits_per_word);
            } else {
                eprintln!("  Output bits: {} ({} words × {} bits/word)",
                         output_bits, payload_word_count, bits_per_word);
            }
        }

        eprintln!("  Sentences: {}", sentence_count);
        eprintln!("  Avg words per sentence: {:.1}", avg_words_per_sentence);
        
        eprintln!("Compactness:");
        eprintln!("  Score: {:.3} ({} BIP39 chars / {} total chars)", 
                  best_compactness, bip39_char_count, total_output_chars);
        eprintln!("  Efficiency: {:.1}% BIP39 characters", (best_compactness * 100.0));
        
        if !variation_stats.is_empty() {
        let min_compactness = variation_stats.iter().fold(f64::INFINITY, |a, &b| a.min(b));
        let max_compactness = variation_stats.iter().fold(0.0_f64, |a, &b| a.max(b));
        let avg_compactness = variation_stats.iter().sum::<f64>() / variation_stats.len() as f64;
        eprintln!("Variations:");
        eprintln!("  Tested: {}", variations);
        eprintln!("  Min compactness: {:.3}", min_compactness);
        eprintln!("  Max compactness: {:.3}", max_compactness);
        eprintln!("  Avg compactness: {:.3}", avg_compactness);
        eprintln!("  Improvement: {:.1}% better than average", 
                  ((best_compactness - avg_compactness) / avg_compactness * 100.0).max(0.0));
        }
    }

    // Decoding: split on whitespace/punct, keep only tokens that are in the BIP39 set.
    // The payload_set contains all the BIP39 words, so filtering is straightforward.
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::RngCore;
    use rand::SeedableRng;
    use glossia::GrammarChecker;
    use glossia::generator::{
        tag_word, select_random_words, max_subsequence_embedding,
        plan_sentence, fill_slots, load_payload_words,
        generate_text, load_cover_words_by_pos, build_pos_mapping_for_wordlist,
    };
    use glossia::generator::utils::{payload_fits, starts_with_vowel_sound};

    #[test]
    fn test_max_subsequence_embedding() {
        // Test basic embedding: slots [Det, N, V, Dot], payload [N, V]
        let slots = vec![Pos::Det, Pos::N, Pos::V, Pos::Dot];
        let payload = vec![
            PayloadTok::new("word1", &[Pos::N]),
            PayloadTok::new("word2", &[Pos::V]),
        ];
        
        let placement = max_subsequence_embedding(&slots, &payload, 0, 2);
        assert!(placement.is_some());
        let placement = placement.unwrap();
        assert_eq!(placement.len(), 2);
        assert_eq!(placement.get(&1), Some(&0)); // N slot -> payload[0]
        assert_eq!(placement.get(&2), Some(&1)); // V slot -> payload[1]
        
        // Test with gaps: slots [Det, N, Adj, V, Dot], payload [N, V]
        let slots = vec![Pos::Det, Pos::N, Pos::Adj, Pos::V, Pos::Dot];
        let placement = max_subsequence_embedding(&slots, &payload, 0, 2);
        assert!(placement.is_some());
        let placement = placement.unwrap();
        assert_eq!(placement.len(), 2);
        assert_eq!(placement.get(&1), Some(&0)); // N slot -> payload[0]
        assert_eq!(placement.get(&3), Some(&1)); // V slot -> payload[1]
        
        // Test impossible: slots [Det, N, Dot], payload [N, V] - can't fit V
        let impossible_slots = vec![Pos::Det, Pos::N, Pos::Dot];
        let placement = max_subsequence_embedding(&impossible_slots, &payload, 0, 2);
        assert!(placement.is_none());
        
        // Test j=0
        let placement = max_subsequence_embedding(&slots, &payload, 0, 0);
        assert!(placement.is_some());
        assert!(placement.unwrap().is_empty());
    }

    #[test]
    fn test_plan_sentence_max_j() {
        // Load cache
        let cache = SequenceCache::load(GenerationMode::Body, "english", 20, false).expect("Failed to load cache");
        let mut rng = StdRng::seed_from_u64(42);
        
        // Create payload with N, V, N
        let payload = vec![
            PayloadTok::new("noun1", &[Pos::N]),
            PayloadTok::new("verb1", &[Pos::V]),
            PayloadTok::new("noun2", &[Pos::N]),
        ];
        
        // Plan for k=5 (should fit at least 2 words, maybe 3)
        let result = plan_sentence(&mut rng, &cache, "S", 5, &payload, 0, false, None);
        assert!(result.is_some());
        let (slots, _refs, forced_placements, j) = result.unwrap();
        assert!(j >= 1, "Should embed at least 1 word");
        assert_eq!(forced_placements.len(), j);
        
        // Verify placements are valid
        for (&slot_idx, &payload_idx) in &forced_placements {
            assert!(slot_idx < slots.len());
            assert!(payload_idx < payload.len());
            assert!(payload_fits(&payload[payload_idx], slots[slot_idx]));
        }
        
        // Verify order preservation: if slot_i < slot_j, then payload_i < payload_j
        let mut sorted_slots: Vec<(usize, usize)> = forced_placements.iter().map(|(&s, &p)| (s, p)).collect();
        sorted_slots.sort_by_key(|(s, _)| *s);
        for i in 1..sorted_slots.len() {
            assert!(sorted_slots[i-1].1 < sorted_slots[i].1, "Placements must preserve payload order");
        }
    }

    #[test]
    fn test_ordered_payload_extraction() {
        // Smoke test: generate text and verify payload words appear in order
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec!["abandon".to_string(), "ability".to_string(), "able".to_string()];
        
        let payload: Vec<PayloadTok> = words
            .iter()
            .map(|word| {
                let tags = tag_word(word);
                PayloadTok::new(word.clone(), &tags)
            })
            .collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();
        
        let lex = setup_test_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(&mut rng, &lex, &payload, false, GenerationMode::Body, "english", 3, 20, SentenceLengthMode::Compact, " ");

        // Extract BIP39 words in order
        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();
        
        assert_eq!(extracted, words.iter().map(|w| w.to_lowercase()).collect::<Vec<_>>(), 
                   "Extracted words should match input words in order");
    }

    #[derive(Default)]
    struct ZeroRng;

    impl RngCore for ZeroRng {
        fn next_u32(&mut self) -> u32 {
            0
        }
        fn next_u64(&mut self) -> u64 {
            0
        }
        fn fill_bytes(&mut self, dest: &mut [u8]) {
            dest.fill(0);
        }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
            self.fill_bytes(dest);
            Ok(())
        }
    }

    #[test]
    fn test_article_selection_uses_forced_next_word_and_normalizes() {
        // Ensure a/an selection uses the *actual* next token in forced-placement mode
        // (body/subject planner) and that highlight/bars don't break vowel detection.
        assert!(starts_with_vowel_sound("ivory"));
        assert!(starts_with_vowel_sound("|ivory|"));
        assert!(starts_with_vowel_sound("\x1b[32mivory\x1b[0m"));

        // Slots: N Det Adj N Dot
        let slots = vec![Pos::N, Pos::Det, Pos::Adj, Pos::N, Pos::Dot];

        // Payload: force "ivory" into the Adj slot (index 2). It starts with a vowel, so determiner should be "an".
        let payload = vec![PayloadTok::new("ivory", &[Pos::Adj])];
        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_set: HashSet<String> = HashSet::new();
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &["the", "a", "an", "each", "some"])
            .with_words(Pos::Adj, &["clear", "simple"])
            .with_words(Pos::N, &["user", "note"]);

        let mut forced: HashMap<usize, usize> = HashMap::new();
        forced.insert(2, 0);

        let mut rng = ZeroRng::default();
        let mut payload_i = 0usize;
        // Det slot at index 1 needs "indef" refinement to trigger a/an phonological rule
        let refs = vec![None, Some("indef".to_string()), None, None, None];
        let out = fill_slots(
            &mut rng,
            &lex,
            &slots,
            &refs,
            &payload,
            &mut payload_i,
            &[],
            None,
            Some(&forced),
            false,
            false,
            true,
            None,
        );

        assert_eq!(out[1], "an", "Expected 'an' before forced 'ivory'");
        assert_eq!(out[2], "ivory");
    }

    #[test]
    fn test_modal_requires_bare_verb_cover_choice() {
        // Ensure our lightweight constraint "Modal -> bare V only" is applied when choosing cover verbs.
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();

        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Modal, &["can"])
            .with_words(Pos::V, &["walked", "walking", "walk"]);

        let slots = vec![Pos::Modal, Pos::V, Pos::Dot];
        let mut rng = ZeroRng::default();
        let mut payload_i = 0usize;
        let refs = vec![None; slots.len()];
        let out = fill_slots(&mut rng, &lex, &slots, &refs, &payload, &mut payload_i, &[], None, None, false, false, true, None);

        assert_eq!(out[0], "can");
        assert_eq!(out[1].trim_end_matches('.'), "walk", "Expected bare verb after modal");
    }

    #[test]
    fn test_v_np_prefers_transitive_cover_verb() {
        // Ensure "V -> NP only if transitive" is applied as a bias for cover verbs.
        // Slots: Det N V Det N Dot  (so V is followed by an NP start)
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();

        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &["the"])
            .with_words(Pos::N, &["user", "note"])
            .with_words(Pos::V, &["sleep", "send"]);

        let slots = vec![Pos::Det, Pos::N, Pos::V, Pos::Det, Pos::N, Pos::Dot];
        let mut rng = ZeroRng::default();
        let mut payload_i = 0usize;
        let refs = vec![None; slots.len()];
        let out = fill_slots(&mut rng, &lex, &slots, &refs, &payload, &mut payload_i, &[], None, None, false, false, true, None);

        assert_eq!(out[2], "send", "Expected transitive verb before NP object");
    }

    /// Fixed seed for reproducible tests
    const TEST_SEED: u64 = 42;

    /// Helper function to set up a test lexicon with minimal cover words
    fn setup_test_lexicon(payload_set: HashSet<String>, wordlist_set: HashSet<String>) -> Lexicon {
        let det_words = ["the", "a", "an", "each", "some"];
        let modal_words = ["should", "could", "would", "might", "may"];
        let aux_words = ["do", "does"];
        let cop_words = ["is", "are"];
        let to_words = ["to"];
        let conj_words = ["and", "but", "or"];
        let prefix_words = ["re", "fwd", "fw", "attn"];
        let adj_words = ["clear", "plain", "dim", "quiet", "steady", "blunt"];
        let n_words = ["user", "node", "server", "batch"];
        let v_words = ["send", "relay", "log", "sift", "forge"];
        let prep_words = ["amid", "beside", "along", "beneath", "via"];
        let adv_words = ["hence", "well", "quite", "thus"];

        Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &det_words)
            .with_words(Pos::Modal, &modal_words)
            .with_words(Pos::Aux, &aux_words)
            .with_words(Pos::Cop, &cop_words)
            .with_words(Pos::To, &to_words)
            .with_words(Pos::Conj, &conj_words)
            .with_words(Pos::Prefix, &prefix_words)
            .with_words(Pos::Adj, &adj_words)
            .with_words(Pos::N, &n_words)
            .with_words(Pos::V, &v_words)
            .with_words(Pos::Prep, &prep_words)
            .with_words(Pos::Adv, &adv_words)
    }

    /// Extract individual sentences from generated text
    fn extract_sentences(text: &str) -> Vec<String> {
        text.split('.')
            .map(|s| s.trim())
            .filter(|s| !s.is_empty())
            .map(|s| {
                // Remove BIP39 markers (|word|)
                s.replace('|', "")
            })
            .collect()
    }

    #[test]
    fn test_generated_sentences_grammar() {
        // Skip if grammar checker files are not available
        let grammar_checker = match GrammarChecker::from_language(glossia::Language::English) {
            Ok(checker) => checker,
            Err(_) => {
                eprintln!("Skipping grammar test: nlprule binary files not found");
                return;
            }
        };

        // Generate sentences with random BIP39 words (using fixed seed for reproducibility)
        let mut rng = StdRng::seed_from_u64(TEST_SEED);
        let words = select_random_words(&mut rng, 10, "english").unwrap();
        
        let payload: Vec<PayloadTok> = words
            .iter()
            .map(|word| {
                let tags = tag_word(word);
                PayloadTok::new(word.clone(), &tags)
            })
            .collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();
        
        let lex = setup_test_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(&mut rng, &lex, &payload, false, GenerationMode::Body, "english", 3, 20, SentenceLengthMode::Natural, " ");

        // Extract individual sentences (splits on '.' — requires Body mode which produces Dot tokens)
        let sentences = extract_sentences(&text);

        assert!(!sentences.is_empty(), "Should generate at least one sentence");

        // Check grammar for each sentence
        let mut total_suggestions = 0;
        let mut sentences_with_errors = 0;

        for sentence in &sentences {
            let suggestions = grammar_checker.check(sentence);
            total_suggestions += suggestions.len();

            if !suggestions.is_empty() {
                sentences_with_errors += 1;
                eprintln!("Sentence with grammar issues: \"{}\"", sentence);
                for suggestion in &suggestions {
                    eprintln!("  Suggestion: {:?}", suggestion);
                }
            }
        }

        // Report results
        eprintln!("\nGrammar check results:");
        eprintln!("  Total sentences: {}", sentences.len());
        eprintln!("  Sentences with errors: {}", sentences_with_errors);
        eprintln!("  Total suggestions: {}", total_suggestions);
        eprintln!("  Average suggestions per sentence: {:.2}",
                  total_suggestions as f64 / sentences.len() as f64);

        // Allow some grammar errors (generated sentences may not be perfect)
        // But we want most sentences to be reasonably correct
        let error_rate = sentences_with_errors as f64 / sentences.len() as f64;
        assert!(
            error_rate < 0.5,
            "More than 50% of sentences have grammar errors (error rate: {:.2}%)",
            error_rate * 100.0
        );
    }

    #[test]
    fn test_grammar_with_different_payload_sizes() {
        // Skip if grammar checker files are not available
        let grammar_checker = match GrammarChecker::from_language(glossia::Language::English) {
            Ok(checker) => checker,
            Err(_) => {
                eprintln!("Skipping grammar test: nlprule binary files not found");
                return;
            }
        };

        // Test with different payload sizes (using same seed for reproducibility)
        for word_count in [5, 8, 12] {
            let mut rng = StdRng::seed_from_u64(TEST_SEED);
            let words = select_random_words(&mut rng, word_count, "english").unwrap();
            
            let payload: Vec<PayloadTok> = words
                .iter()
                .map(|word| {
                    let tags = tag_word(word);
                    PayloadTok::new(word.clone(), &tags)
                })
                .collect();

            let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
            let wordlist_words = load_payload_words("english").unwrap();
            let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

            let lex = setup_test_lexicon(payload_set, wordlist_set);
            let (text, _) = generate_text(&mut rng, &lex, &payload, false, GenerationMode::Body, "english", 3, 20, SentenceLengthMode::Natural, " ");

            let sentences = extract_sentences(&text);
            
            // Check that at least some sentences are grammatically reasonable
            let mut correct_sentences = 0;
            for sentence in &sentences {
                let suggestions = grammar_checker.check(sentence);
                if suggestions.is_empty() {
                    correct_sentences += 1;
                }
            }

            eprintln!("Payload size {}: {}/{} sentences grammatically correct", 
                     word_count, correct_sentences, sentences.len());
            
            // At least one sentence should be correct
            assert!(
                correct_sentences > 0 || sentences.is_empty(),
                "No grammatically correct sentences generated for payload size {}",
                word_count
            );
        }
    }

    #[test]
    fn test_sentence_structure() {
        // Test that generated sentences have reasonable structure (using fixed seed for reproducibility)
        let mut rng = StdRng::seed_from_u64(TEST_SEED);
        let words = select_random_words(&mut rng, 5, "english").unwrap();
        
        let payload: Vec<PayloadTok> = words
            .iter()
            .map(|word| {
                let tags = tag_word(word);
                PayloadTok::new(word.clone(), &tags)
            })
            .collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set, wordlist_set);
        let (text, _) = generate_text(&mut rng, &lex, &payload, false, GenerationMode::Subject, "english", 3, 20, SentenceLengthMode::Compact, " ");

        let sentences = extract_sentences(&text);
        
        // Each sentence should have at least a few words
        for sentence in &sentences {
            let word_count = sentence.split_whitespace().count();
            assert!(
                word_count >= 3,
                "Sentence too short ({} words): \"{}\"",
                word_count,
                sentence
            );
        }

        // Should have generated at least one sentence
        assert!(!sentences.is_empty(), "Should generate at least one sentence");
    }

    #[test]
    fn test_compute_k_candidates_compact_mode() {
        let cache = SequenceCache::load(GenerationMode::Body, "english", 20, false).expect("Failed to load cache");
        let mut rng = StdRng::seed_from_u64(42);
        
        // Compact mode should return k_min..=k_max in order
        let candidates = compute_k_candidates(
            &mut rng,
            &cache,
            "S",
            3,
            10,
            SentenceLengthMode::Compact,
            false,
        );
        
        assert_eq!(candidates, vec![3, 4, 5, 6, 7, 8, 9, 10]);
        
        // With require_prefix, should start at k_min+1
        let candidates_prefix = compute_k_candidates(
            &mut rng,
            &cache,
            "S",
            3,
            10,
            SentenceLengthMode::Compact,
            true,
        );
        
        assert_eq!(candidates_prefix, vec![4, 5, 6, 7, 8, 9, 10]);
    }

    #[test]
    fn test_compute_k_candidates_natural_mode() {
        let cache = SequenceCache::load(GenerationMode::Body, "english", 20, false).expect("Failed to load cache");
        let mut rng = StdRng::seed_from_u64(42);
        
        // Natural mode should ignore k_min and sample from all available k values
        let candidates = compute_k_candidates(
            &mut rng,
            &cache,
            "S",
            3,  // k_min is ignored in natural mode
            10,
            SentenceLengthMode::Natural,
            false,
        );
        
        // Verify all candidates are in valid range (k_min is ignored, so can be < k_min)
        for &k in &candidates {
            assert!(k >= 1 && k <= 10, "k={} should be in range [1, 10] (k_min ignored)", k);
        }
        
        // Verify no duplicates
        let mut seen = std::collections::HashSet::new();
        for &k in &candidates {
            assert!(!seen.contains(&k), "Duplicate k={} in candidates", k);
            seen.insert(k);
        }
        
        // Verify sampled k is first
        assert!(!candidates.is_empty(), "Should have at least one candidate");
        let sampled_k = candidates[0];
        assert!(sampled_k >= 1 && sampled_k <= 10, "Sampled k should be in valid range");
    }

    #[test]
    fn test_compute_k_candidates_natural_mode_prefix() {
        let cache = SequenceCache::load(GenerationMode::Subject, "english", 20, false).expect("Failed to load cache");
        let mut rng = StdRng::seed_from_u64(42);
        
        // Natural mode with require_prefix should start at k=2 (k_min ignored)
        let candidates = compute_k_candidates(
            &mut rng,
            &cache,
            "S",
            3,  // k_min is ignored in natural mode
            10,
            SentenceLengthMode::Natural,
            true,
        );
        
        // Verify all candidates are >= 2 (need at least Prefix + one more terminal)
        for &k in &candidates {
            assert!(k >= 2 && k <= 10, "k={} should be in range [2, 10] when require_prefix=true (k_min ignored)", k);
        }
        
        // Verify no duplicates
        let mut seen = std::collections::HashSet::new();
        for &k in &candidates {
            assert!(!seen.contains(&k), "Duplicate k={} in candidates", k);
            seen.insert(k);
        }
        
        // Verify sampled k is first
        assert!(!candidates.is_empty(), "Should have at least one candidate");
        let sampled_k = candidates[0];
        assert!(sampled_k >= 2 && sampled_k <= 10, "Sampled k should be in valid range");
    }

    #[test]
    fn test_compute_k_candidates_natural_mode_weight_ordering() {
        let cache = SequenceCache::load(GenerationMode::Body, "english", 20, false).expect("Failed to load cache");
        let mut rng = StdRng::seed_from_u64(42);
        
        // Test that after the sampled k, remaining k's are in descending weight order
        let candidates = compute_k_candidates(
            &mut rng,
            &cache,
            "S",
            3,
            8,
            SentenceLengthMode::Natural,
            false,
        );
        
        // Compute weights for verification
        let mut k_weights: Vec<(usize, f64)> = Vec::new();
        for k in 3..=8 {
            if let Some(sequences) = cache.get("S", k) {
                let weight: f64 = sequences.iter()
                    .map(|seq_prob| seq_prob.probability)
                    .sum();
                if weight > 0.0 {
                    k_weights.push((k, weight));
                }
            }
        }
        
        // Find the sampled k (first in candidates)
        let sampled_k = candidates[0];
        
        // Extract remaining k's (excluding sampled_k)
        let remaining: Vec<usize> = candidates.iter()
            .skip(1)
            .filter(|&&k| k != sampled_k)
            .cloned()
            .collect();
        
        // Verify remaining k's are in descending weight order
        let mut prev_weight = f64::INFINITY;
        for &k in &remaining {
            if let Some((_, weight)) = k_weights.iter().find(|(k_val, _)| *k_val == k) {
                assert!(*weight <= prev_weight,
                    "k={} has weight {} which should be <= previous weight {}",
                    k, weight, prev_weight);
                prev_weight = *weight;
            }
        }
    }

    // ========== fill_slots refinement tests ==========

    #[test]
    fn test_fill_slots_det_indef_before_vowel_produces_an() {
        // Det[indef] before a vowel-initial word should produce "an"
        let payload = vec![PayloadTok::new("apple", &[Pos::N])];
        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_set: HashSet<String> = HashSet::new();
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &["the", "a", "an", "each", "some"])
            .with_words(Pos::N, &["user", "note"]);

        let slots = vec![Pos::Det, Pos::N, Pos::Dot];
        let refinements = vec![Some("indef".to_string()), None, None];
        let mut forced: HashMap<usize, usize> = HashMap::new();
        forced.insert(1, 0); // Force "apple" into N slot

        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, Some(&forced), false, false, true, None,
        );

        assert_eq!(out[0], "an", "Det[indef] before 'apple' (vowel) should produce 'an'");
        // Dot is appended to the last word, so "apple" becomes "apple."
        assert_eq!(out[1], "apple.", "payload word should have Dot appended");
    }

    #[test]
    fn test_fill_slots_det_indef_before_consonant_produces_a() {
        // Det[indef] before a consonant-initial word should produce "a"
        let payload = vec![PayloadTok::new("basket", &[Pos::N])];
        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_set: HashSet<String> = HashSet::new();
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &["the", "a", "an", "each", "some"])
            .with_words(Pos::N, &["user", "note"]);

        let slots = vec![Pos::Det, Pos::N, Pos::Dot];
        let refinements = vec![Some("indef".to_string()), None, None];
        let mut forced: HashMap<usize, usize> = HashMap::new();
        forced.insert(1, 0); // Force "basket" into N slot

        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, Some(&forced), false, false, true, None,
        );

        assert_eq!(out[0], "a", "Det[indef] before 'basket' (consonant) should produce 'a'");
        // Dot is appended to the last word, so "basket" becomes "basket."
        assert_eq!(out[1], "basket.", "payload word should have Dot appended");
    }

    #[test]
    fn test_fill_slots_det_def_produces_definite() {
        // Det[def] should produce a definite determiner (the, its, our, etc.)
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();
        let mut refined_cover: HashMap<(Pos, String), Vec<String>> = HashMap::new();
        refined_cover.insert(
            (Pos::Det, "def".to_string()),
            vec!["the".to_string(), "its".to_string(), "our".to_string()],
        );
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Det, &["the", "a", "an", "its", "our", "some"])
            .with_words(Pos::N, &["user", "note"])
            .with_refined_cover(refined_cover);

        let slots = vec![Pos::Det, Pos::N, Pos::Dot];
        let refinements = vec![Some("def".to_string()), None, None];

        let def_words: HashSet<&str> = ["the", "its", "our"].iter().copied().collect();
        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, None, false, false, true, None,
        );

        assert!(
            def_words.contains(out[0].as_str()),
            "Det[def] should produce definite determiner, got: '{}'", out[0]
        );
    }

    #[test]
    fn test_fill_slots_cop_sg_produces_is() {
        // Cop[sg] should produce "is"
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();
        let mut refined_cover: HashMap<(Pos, String), Vec<String>> = HashMap::new();
        refined_cover.insert(
            (Pos::Cop, "sg".to_string()),
            vec!["is".to_string()],
        );
        refined_cover.insert(
            (Pos::Cop, "pl".to_string()),
            vec!["are".to_string()],
        );
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Cop, &["is", "are"])
            .with_words(Pos::Adj, &["clear", "plain"])
            .with_words(Pos::N, &["user", "note"])
            .with_refined_cover(refined_cover);

        let slots = vec![Pos::N, Pos::Cop, Pos::Adj, Pos::Dot];
        let refinements = vec![None, Some("sg".to_string()), None, None];

        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, None, false, false, true, None,
        );

        assert_eq!(out[1], "is", "Cop[sg] should produce 'is'");
    }

    #[test]
    fn test_fill_slots_cop_pl_produces_are() {
        // Cop[pl] should produce "are"
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();
        let mut refined_cover: HashMap<(Pos, String), Vec<String>> = HashMap::new();
        refined_cover.insert(
            (Pos::Cop, "sg".to_string()),
            vec!["is".to_string()],
        );
        refined_cover.insert(
            (Pos::Cop, "pl".to_string()),
            vec!["are".to_string()],
        );
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::Cop, &["is", "are"])
            .with_words(Pos::Adj, &["clear", "plain"])
            .with_words(Pos::N, &["user", "note"])
            .with_refined_cover(refined_cover);

        let slots = vec![Pos::N, Pos::Cop, Pos::Adj, Pos::Dot];
        let refinements = vec![None, Some("pl".to_string()), None, None];

        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, None, false, false, true, None,
        );

        assert_eq!(out[1], "are", "Cop[pl] should produce 'are'");
    }

    #[test]
    fn test_fill_slots_unrefined_fallback() {
        // Slots without refinement should still produce valid cover words
        let payload: Vec<PayloadTok> = vec![];
        let payload_set: HashSet<String> = HashSet::new();
        let wordlist_set: HashSet<String> = HashSet::new();
        let lex = Lexicon::new(payload_set, wordlist_set)
            .with_words(Pos::N, &["user", "node"])
            .with_words(Pos::V, &["send", "relay"])
            .with_words(Pos::Adj, &["clear", "plain"]);

        let slots = vec![Pos::Adj, Pos::N, Pos::V, Pos::N, Pos::Dot];
        let refinements = vec![None, None, None, None, None];

        let mut rng = StdRng::seed_from_u64(42);
        let mut payload_i = 0usize;
        let out = fill_slots(
            &mut rng, &lex, &slots, &refinements, &payload, &mut payload_i,
            &[], None, None, false, false, true, None,
        );

        // 4 word tokens + Dot merged into last word
        assert_eq!(out.len(), 4, "Should produce 4 tokens (Dot is appended)");
        for word in &out {
            let cleaned = word.trim_end_matches('.');
            assert!(!cleaned.is_empty(), "Every slot should produce a non-empty word");
        }
    }

    // ========== Cross-language validation tests ==========

    #[test]
    fn test_english_no_leakage_into_latin() {
        // Generate Latin text and verify no English function words appear
        // Note: "is" is excluded because it's a valid Latin pronoun (cover word).
        // "a" is excluded because it's a valid Latin preposition.
        let english_function_words: HashSet<&str> = [
            "the", "an", "are", "does", "do",
        ].iter().copied().collect();

        for seed in [42u64, 123, 999, 7, 2024] {
            let mut rng = StdRng::seed_from_u64(seed);
            let words = select_random_words(&mut rng, 6, "latin").unwrap();

            let payload: Vec<PayloadTok> = words.iter().map(|word| {
                let tags = glossia::generator::build_pos_mapping("latin")
                    .unwrap()
                    .get(&word.to_lowercase())
                    .cloned()
                    .unwrap_or_default();
                PayloadTok::new(word.clone(), &tags)
            }).collect();

            if payload.iter().any(|t| t.allowed.is_empty()) {
                continue; // Skip seeds with untagged words
            }

            let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
            let wordlist_words = load_payload_words("latin").unwrap();
            let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

            let (by_pos, refined_cover) = glossia::generator::load_cover_words_by_pos(&wordlist_set, "latin");
            let lex = Lexicon::new(payload_set, wordlist_set)
                .with_words(Pos::N, &by_pos.get(&Pos::N).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::V, &by_pos.get(&Pos::V).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::Adj, &by_pos.get(&Pos::Adj).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::Adv, &by_pos.get(&Pos::Adv).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::Prep, &by_pos.get(&Pos::Prep).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::Cop, &by_pos.get(&Pos::Cop).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::Modal, &by_pos.get(&Pos::Modal).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_words(Pos::To, &by_pos.get(&Pos::To).map(|v| v.iter().map(|s| s.as_str()).collect::<Vec<_>>()).unwrap_or_default())
                .with_refined_cover(refined_cover);

            let (text, _) = generate_text(
                &mut rng, &lex, &payload, false, GenerationMode::Body, "latin",
                3, 20, SentenceLengthMode::Compact, " ",
            );

            for token in text.split_whitespace() {
                let cleaned = normalize_token_for_bip39(token);
                assert!(
                    !english_function_words.contains(cleaned.as_str()),
                    "Latin output (seed={}) contains English function word '{}' in text: {}",
                    seed, cleaned, text,
                );
            }
        }
    }

    #[test]
    fn test_grammar_introspection_english_has_dot() {
        let grammar = Grammar::from_language_dialect("english", "body").expect("English body grammar");
        assert!(grammar.grammar_uses_pos(Pos::Dot), "English body grammar should use Dot");
    }

    #[test]
    fn test_grammar_introspection_latin_has_dot() {
        let grammar = Grammar::from_language_dialect("latin", "body").expect("Latin body grammar");
        assert!(grammar.grammar_uses_pos(Pos::Dot), "Latin body grammar should use Dot");
    }

    // ========== Round-trip verification tests ==========

    #[test]
    fn test_english_round_trip() {
        // Encode 4+ BIP39 words, extract from output, verify they match in order
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec![
            "abandon".to_string(),
            "ability".to_string(),
            "able".to_string(),
            "about".to_string(),
            "above".to_string(),
        ];

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = tag_word(word);
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        // Extract payload words from the output text (decode step)
        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();

        assert_eq!(
            extracted,
            words.iter().map(|w| w.to_lowercase()).collect::<Vec<_>>(),
            "Extracted payload words should match input in order.\nInput: {:?}\nOutput text: {}\nExtracted: {:?}",
            words, text, extracted,
        );
    }

    #[test]
    fn test_deterministic_output_same_seed() {
        // Same payload + same seed produces identical output across runs
        let words = vec!["abandon".to_string(), "ability".to_string(), "able".to_string()];
        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = tag_word(word);
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set, wordlist_set);

        let mut rng1 = StdRng::seed_from_u64(42);
        let (text1, _) = generate_text(
            &mut rng1, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        let mut rng2 = StdRng::seed_from_u64(42);
        let (text2, _) = generate_text(
            &mut rng2, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        assert_eq!(text1, text2, "Same seed should produce identical output");
    }

    // ========== Edge case tests ==========

    #[test]
    fn test_single_payload_word_appears_in_output() {
        let mut rng = StdRng::seed_from_u64(42);
        let payload = vec![PayloadTok::new("abandon", &[Pos::N, Pos::V])];
        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();

        assert_eq!(extracted, vec!["abandon"], "Single payload word should appear in output");
    }

    #[test]
    fn test_many_payload_words_round_trip() {
        // Test with 12 random BIP39 words
        let mut rng = StdRng::seed_from_u64(99);
        let words = select_random_words(&mut rng, 12, "english").unwrap();

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = tag_word(word);
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        // Extract payload words from output
        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();

        assert_eq!(
            extracted,
            words.iter().map(|w| w.to_lowercase()).collect::<Vec<_>>(),
            "All 12 payload words should appear in order in the output"
        );
    }

    #[test]
    fn test_output_ends_with_period_body_mode() {
        // Body mode text should end with a period for English
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec!["abandon".to_string(), "ability".to_string()];
        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = tag_word(word);
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set, wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        assert!(
            text.ends_with('.'),
            "Body mode English text should end with period. Got: '{}'",
            text
        );
    }

    // ========== Meta-language integration tests ==========

    /// Build a Lexicon for meta language from its actual cover.yaml data.
    fn setup_meta_lexicon(payload_set: HashSet<String>, wordlist_set: HashSet<String>) -> Lexicon {
        let (by_pos, refined_cover) = load_cover_words_by_pos(&wordlist_set, "meta");
        let mut lex = Lexicon::new(payload_set, wordlist_set)
            .with_refined_cover(refined_cover);
        for (pos, words) in &by_pos {
            let word_refs: Vec<&str> = words.iter().map(|s: &String| s.as_str()).collect();
            lex = lex.with_words(*pos, &word_refs);
        }
        lex
    }

    #[test]
    fn test_meta_grammar_loads_with_max_k() {
        let grammar = Grammar::from_language_dialect("meta", "body")
            .expect("Should load meta body grammar");
        assert_eq!(grammar.max_k(), Some(12), "Meta grammar should declare max_k=12");
        assert!(grammar.grammar_uses_pos(Pos::Dot), "Meta grammar should use Dot");
    }

    #[test]
    fn test_meta_payload_words_load() {
        let words = load_payload_words("meta").unwrap();
        assert_eq!(words.len(), 19, "Meta should have exactly 19 payload words");
        // Spot-check some dialect identifiers
        assert!(words.contains(&"latin".to_string()), "Should contain 'latin'");
        assert!(words.contains(&"english".to_string()), "Should contain 'english'");
        assert!(words.contains(&"czech".to_string()), "Should contain 'czech'");
        assert!(words.contains(&"hex".to_string()), "Should contain 'hex'");
        assert!(words.contains(&"nostr".to_string()), "Should contain 'nostr'");
        assert!(words.contains(&"html".to_string()), "Should contain 'html'");
    }

    #[test]
    fn test_meta_pos_mapping() {
        let mapping = build_pos_mapping_for_wordlist("meta", "default")
            .expect("Should load meta POS mapping");
        // All meta payload words should have POS tags
        let words = load_payload_words("meta").unwrap();
        for word in &words {
            let word_lower = word.to_lowercase();
            assert!(
                mapping.contains_key(&word_lower),
                "Payload word '{}' should have POS mapping", word
            );
            assert!(
                !mapping[&word_lower].is_empty(),
                "Payload word '{}' should have at least one POS tag", word
            );
        }
    }

    #[test]
    fn test_meta_round_trip_4_words() {
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec![
            "base64".to_string(),
            "latin".to_string(),
            "nostr".to_string(),
            "english".to_string(),
        ];

        // Use meta POS mapping (not English tag_word)
        let pos_mapping = build_pos_mapping_for_wordlist("meta", "default")
            .expect("Meta POS mapping");

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = pos_mapping.get(&word.to_lowercase())
                .cloned()
                .unwrap_or_default();
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("meta").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_meta_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "meta",
            3, 12, SentenceLengthMode::Compact, " ",
        );

        eprintln!("Meta 4-word output: {}", text);

        // Round-trip: extract payload words from output
        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();

        assert_eq!(
            extracted,
            words.iter().map(|w| w.to_lowercase()).collect::<Vec<_>>(),
            "Extracted payload words should match input in order.\nInput: {:?}\nOutput text: {}\nExtracted: {:?}",
            words, text, extracted,
        );
    }

    #[test]
    fn test_meta_round_trip_8_words() {
        let mut rng = StdRng::seed_from_u64(100);
        let words = select_random_words(&mut rng, 8, "meta").unwrap();

        let pos_mapping = build_pos_mapping_for_wordlist("meta", "default")
            .expect("Meta POS mapping");

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = pos_mapping.get(&word.to_lowercase())
                .cloned()
                .unwrap_or_default();
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("meta").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_meta_lexicon(payload_set.clone(), wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "meta",
            3, 12, SentenceLengthMode::Compact, " ",
        );

        eprintln!("Meta 8-word output: {}", text);

        let extracted: Vec<String> = text
            .split_whitespace()
            .map(normalize_token_for_bip39)
            .filter(|w| !w.is_empty() && payload_set.contains(w))
            .collect();

        assert_eq!(
            extracted,
            words.iter().map(|w| w.to_lowercase()).collect::<Vec<_>>(),
            "All 8 meta payload words should appear in order.\nInput: {:?}\nOutput: {}\nExtracted: {:?}",
            words, text, extracted,
        );
    }

    #[test]
    fn test_meta_deterministic_output() {
        let words = vec!["hex".to_string(), "latin".to_string(), "bits".to_string()];
        let pos_mapping = build_pos_mapping_for_wordlist("meta", "default")
            .expect("Meta POS mapping");

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = pos_mapping.get(&word.to_lowercase())
                .cloned()
                .unwrap_or_default();
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("meta").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_meta_lexicon(payload_set, wordlist_set);

        let mut rng1 = StdRng::seed_from_u64(42);
        let (text1, _) = generate_text(
            &mut rng1, &lex, &payload, false, GenerationMode::Body, "meta",
            3, 12, SentenceLengthMode::Compact, " ",
        );

        let mut rng2 = StdRng::seed_from_u64(42);
        let (text2, _) = generate_text(
            &mut rng2, &lex, &payload, false, GenerationMode::Body, "meta",
            3, 12, SentenceLengthMode::Compact, " ",
        );

        assert_eq!(text1, text2, "Same seed should produce identical meta output");
    }

    #[test]
    fn test_meta_output_ends_with_period() {
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec!["pgp".to_string(), "nostr".to_string()];
        let pos_mapping = build_pos_mapping_for_wordlist("meta", "default")
            .expect("Meta POS mapping");

        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = pos_mapping.get(&word.to_lowercase())
                .cloned()
                .unwrap_or_default();
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("meta").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_meta_lexicon(payload_set, wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "meta",
            3, 12, SentenceLengthMode::Compact, " ",
        );

        assert!(
            text.ends_with('.'),
            "Meta body mode should end with period. Got: '{}'",
            text
        );
    }

    #[test]
    fn test_first_word_is_capitalized() {
        // First word of output should be capitalized
        let mut rng = StdRng::seed_from_u64(42);
        let words = vec!["abandon".to_string(), "ability".to_string()];
        let payload: Vec<PayloadTok> = words.iter().map(|word| {
            let tags = tag_word(word);
            PayloadTok::new(word.clone(), &tags)
        }).collect();

        let payload_set: HashSet<String> = payload.iter().map(|t| t.word.to_lowercase()).collect();
        let wordlist_words = load_payload_words("english").unwrap();
        let wordlist_set: HashSet<String> = wordlist_words.iter().map(|w| w.to_lowercase()).collect();

        let lex = setup_test_lexicon(payload_set, wordlist_set);
        let (text, _) = generate_text(
            &mut rng, &lex, &payload, false, GenerationMode::Body, "english",
            3, 20, SentenceLengthMode::Compact, " ",
        );

        let first_char = text.chars().next().unwrap();
        assert!(
            first_char.is_uppercase(),
            "First character should be uppercase. Got: '{}'",
            first_char
        );
    }
}