use super::gguf::Gguf;
use super::{
    ChatTurn, ModelFamily, VISION_END_TOKEN, VISION_IMAGE_TOKEN, VISION_START_TOKEN,
    VISION_TOKEN_END, VISION_TOKEN_START,
};
use crate::settings::ReasoningMode;
use serde::Deserialize;
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use unicode_normalization_alignments::UnicodeNormalization;

const HIGH_REASONING_PREFIX: &str = "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n\
You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n\
Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n";
const MAX_REASONING_PREFIX: &str = "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n\
You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n\
Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n";

fn reasoning_prefix(family: ModelFamily, reasoning: ReasoningMode) -> Option<&'static str> {
    match (family, reasoning) {
        (ModelFamily::DeepSeek, ReasoningMode::High) => Some(HIGH_REASONING_PREFIX),
        (ModelFamily::DeepSeek, ReasoningMode::Max) => Some(MAX_REASONING_PREFIX),
        (ModelFamily::Glm, ReasoningMode::Low) => Some("Reasoning Effort: Low"),
        (ModelFamily::Glm, ReasoningMode::High) => Some("Reasoning Effort: High"),
        (ModelFamily::Glm, ReasoningMode::Max) => Some("Reasoning Effort: Max"),
        (ModelFamily::Qwen, _) => None,
        (
            _,
            ReasoningMode::Direct
            | ReasoningMode::Low
            | ReasoningMode::Medium
            | ReasoningMode::XHigh,
        ) => None,
    }
}

pub(super) struct Tokenizer {
    family: ModelFamily,
    tokens: Vec<Vec<u8>>,
    token_to_id: HashMap<Vec<u8>, i32>,
    merge_rank: HashMap<Vec<u8>, usize>,
    bos: i32,
    eos: i32,
    system: i32,
    user: i32,
    assistant: i32,
    observation: i32,
    sop: i32,
    think_start: i32,
    think_end: i32,
    alternate_eos: i32,
    rendered_specials: Vec<(Vec<u8>, i32)>,
}

#[derive(Deserialize)]
struct QwenTokenizerFile {
    normalizer: QwenNormalizer,
    added_tokens: Vec<QwenAddedToken>,
    model: QwenBpe,
}

#[derive(Deserialize)]
#[serde(tag = "type")]
enum QwenNormalizer {
    #[serde(rename = "NFC")]
    Nfc,
}

#[derive(Deserialize)]
struct QwenAddedToken {
    id: usize,
    content: String,
}

#[derive(Deserialize)]
struct QwenBpe {
    #[serde(rename = "type")]
    kind: String,
    vocab: HashMap<String, usize>,
    merges: Vec<String>,
    byte_fallback: bool,
}

impl Tokenizer {
    pub(super) fn load(model: &Gguf, family: ModelFamily) -> Result<Self, String> {
        let tokens = model.strings("tokenizer.ggml.tokens")?.to_vec();
        let merges = model.strings("tokenizer.ggml.merges")?;
        let token_to_id = tokens
            .iter()
            .enumerate()
            .map(|(id, token)| (token.clone(), id as i32))
            .collect::<HashMap<_, _>>();
        let merge_rank = merges
            .iter()
            .enumerate()
            .map(|(rank, merge)| (merge.clone(), rank))
            .collect::<HashMap<_, _>>();
        let lookup = |token: &[u8]| token_to_id.get(token).copied().unwrap_or(-1);
        let required = |token: &[u8]| {
            token_to_id.get(token).copied().ok_or_else(|| {
                format!(
                    "required tokenizer token is missing: {}",
                    String::from_utf8_lossy(token)
                )
            })
        };

        let (bos, eos, system, user, assistant, observation, sop, think_start, think_end) =
            match family {
                ModelFamily::DeepSeek => (
                    required("<｜begin▁of▁sentence｜>".as_bytes())?,
                    required("<｜end▁of▁sentence｜>".as_bytes())?,
                    -1,
                    required("<｜User｜>".as_bytes())?,
                    required("<｜Assistant｜>".as_bytes())?,
                    -1,
                    -1,
                    required(b"<think>")?,
                    required(b"</think>")?,
                ),
                ModelFamily::Glm => (
                    model
                        .token_id("tokenizer.ggml.bos_token_id")
                        .ok()
                        .unwrap_or_else(|| lookup(b"<sop>")),
                    model
                        .token_id("tokenizer.ggml.eos_token_id")
                        .ok()
                        .unwrap_or_else(|| lookup(b"<|endoftext|>")),
                    lookup(b"<|system|>"),
                    lookup(b"<|user|>"),
                    lookup(b"<|assistant|>"),
                    lookup(b"<|observation|>"),
                    lookup(b"<sop>"),
                    lookup(b"<think>"),
                    lookup(b"</think>"),
                ),
                ModelFamily::Qwen => {
                    return Err("Qwen tokenizers must be loaded from tokenizer.json".into());
                }
            };
        if [bos, user, assistant, think_start, think_end]
            .into_iter()
            .any(|token| token < 0)
            || (family == ModelFamily::Glm && system < 0)
        {
            return Err("tokenizer does not provide the required DS4 chat markers".into());
        }
        let rendered_specials = [
            ("<｜begin▁of▁sentence｜>".as_bytes(), bos),
            ("<｜end▁of▁sentence｜>".as_bytes(), eos),
            (b"[gMASK]".as_slice(), bos),
            (b"<sop>".as_slice(), sop),
            (b"<|system|>".as_slice(), system),
            ("<｜User｜>".as_bytes(), user),
            ("<｜Assistant｜>".as_bytes(), assistant),
            (b"<|user|>".as_slice(), user),
            (b"<|assistant|>".as_slice(), assistant),
            (b"<|observation|>".as_slice(), observation),
            (b"<think>".as_slice(), think_start),
            (b"</think>".as_slice(), think_end),
            (b"<tool_call>".as_slice(), lookup(b"<tool_call>")),
            (b"</tool_call>".as_slice(), lookup(b"</tool_call>")),
            (b"<tool_response>".as_slice(), lookup(b"<tool_response>")),
            (b"</tool_response>".as_slice(), lookup(b"</tool_response>")),
            (b"<arg_key>".as_slice(), lookup(b"<arg_key>")),
            (b"</arg_key>".as_slice(), lookup(b"</arg_key>")),
            (b"<arg_value>".as_slice(), lookup(b"<arg_value>")),
            (b"</arg_value>".as_slice(), lookup(b"</arg_value>")),
            ("｜DSML｜".as_bytes(), lookup("｜DSML｜".as_bytes())),
        ]
        .into_iter()
        .filter(|(_, token)| *token >= 0)
        .map(|(marker, token)| (marker.to_vec(), token))
        .collect();

        Ok(Self {
            family,
            tokens,
            token_to_id,
            merge_rank,
            bos,
            eos,
            system,
            user,
            assistant,
            observation,
            sop,
            think_start,
            think_end,
            alternate_eos: -1,
            rendered_specials,
        })
    }

    pub(super) fn load_qwen(path: &Path) -> Result<Self, String> {
        let file: QwenTokenizerFile = serde_json::from_slice(
            &fs::read(path).map_err(|error| format!("{}: {error}", path.display()))?,
        )
        .map_err(|error| format!("{}: {error}", path.display()))?;
        let QwenNormalizer::Nfc = file.normalizer;
        if file.model.kind != "BPE" || file.model.byte_fallback {
            return Err("Qwen tokenizer must use BPE without byte fallback".into());
        }
        let maximum = file
            .model
            .vocab
            .values()
            .copied()
            .chain(file.added_tokens.iter().map(|token| token.id))
            .max()
            .ok_or_else(|| "Qwen tokenizer vocabulary is empty".to_owned())?;
        if maximum != 248_076 {
            return Err(format!(
                "Qwen tokenizer ends at token id {maximum}, expected 248076"
            ));
        }
        // The model head has 248320 rows; the final 243 ids are deliberately
        // reserved and have no tokenizer spelling in the pinned source.
        let mut tokens = vec![Vec::new(); 248_320];
        for (token, id) in file.model.vocab {
            if !tokens[id].is_empty() {
                return Err(format!("Qwen tokenizer has duplicate token id {id}"));
            }
            tokens[id] = token.into_bytes();
        }
        let mut rendered_specials = Vec::new();
        for token in file.added_tokens {
            if !tokens[token.id].is_empty() {
                return Err(format!(
                    "Qwen tokenizer has duplicate token id {}",
                    token.id
                ));
            }
            tokens[token.id] = token.content.as_bytes().to_vec();
            rendered_specials.push((token.content.into_bytes(), token.id as i32));
        }
        if tokens[..=maximum].iter().any(Vec::is_empty) {
            return Err("Qwen tokenizer token ids are not contiguous".into());
        }
        let token_to_id = tokens
            .iter()
            .enumerate()
            .filter(|(_, token)| !token.is_empty())
            .map(|(id, token)| (token.clone(), id as i32))
            .collect::<HashMap<_, _>>();
        let required = |token: &[u8]| {
            token_to_id.get(token).copied().ok_or_else(|| {
                format!(
                    "required Qwen tokenizer token is missing: {}",
                    String::from_utf8_lossy(token)
                )
            })
        };
        let merge_rank = file
            .model
            .merges
            .into_iter()
            .enumerate()
            .map(|(rank, merge)| (merge.into_bytes(), rank))
            .collect();
        let alternate_eos = required(b"<|endoftext|>")?;
        let im_start = required(b"<|im_start|>")?;
        let im_end = required(b"<|im_end|>")?;
        let think_start = required(b"<think>")?;
        let think_end = required(b"</think>")?;
        Ok(Self {
            family: ModelFamily::Qwen,
            tokens,
            token_to_id,
            merge_rank,
            bos: alternate_eos,
            eos: im_end,
            system: im_start,
            user: im_start,
            assistant: im_start,
            observation: im_start,
            sop: im_end,
            think_start,
            think_end,
            alternate_eos,
            rendered_specials,
        })
    }

    pub(super) fn vocab_size(&self) -> usize {
        self.tokens.len()
    }

    pub(super) fn validate_qwen_contract(&self) -> Result<(), String> {
        if self.family != ModelFamily::Qwen {
            return Err("tokenizer is not Qwen".into());
        }
        for (token, expected) in [
            (b"<|endoftext|>".as_slice(), 248_044),
            (b"<|im_start|>".as_slice(), 248_045),
            (b"<|im_end|>".as_slice(), 248_046),
            (b"<tool_call>".as_slice(), 248_058),
            (b"</tool_call>".as_slice(), 248_059),
            (b"<tool_response>".as_slice(), 248_066),
            (b"</tool_response>".as_slice(), 248_067),
            (b"<think>".as_slice(), 248_068),
            (b"</think>".as_slice(), 248_069),
        ] {
            if self.token_to_id.get(token).copied() != Some(expected) {
                return Err(format!(
                    "Qwen special token {} does not have id {expected}",
                    String::from_utf8_lossy(token)
                ));
            }
        }
        if ![248_044, 248_045, 248_046]
            .into_iter()
            .all(|token| self.is_ple_reset(token))
            || ![248_044, 248_046]
                .into_iter()
                .all(|token| self.is_stop(token))
        {
            return Err("Qwen EOS or PLE reset markers are invalid".into());
        }
        Ok(())
    }

    pub(super) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
        let token = self.tokens.get(usize::try_from(token).ok()?)?;
        if token.windows(3).any(|window| window == [0xef, 0xbd, 0x9c]) {
            return Some(token.clone());
        }
        let text = std::str::from_utf8(token).ok()?;
        Some(text.chars().filter_map(gpt2_codepoint_to_byte).collect())
    }

    pub(super) fn tokenize(&self, text: &str) -> Vec<i32> {
        let mut output = Vec::new();
        match self.family {
            ModelFamily::DeepSeek => self.tokenize_joyai(text, &mut output),
            ModelFamily::Glm => self.tokenize_gpt4(text, &mut output, 3),
            ModelFamily::Qwen => self.tokenize_qwen(text, &mut output),
        }
        output
    }

    pub(super) fn tokenize_rendered(&self, text: &str) -> Vec<i32> {
        let bytes = text.as_bytes();
        let mut output = Vec::new();
        let mut span = 0;
        let mut position = 0;
        while position < bytes.len() {
            if bytes[position..].starts_with(VISION_TOKEN_START.as_bytes())
                && let Some(relative_end) =
                    text[position + VISION_TOKEN_START.len()..].find(VISION_TOKEN_END)
            {
                let count_start = position + VISION_TOKEN_START.len();
                let count_end = count_start + relative_end;
                if let Ok(count) = text[count_start..count_end].parse::<usize>() {
                    self.tokenize_plain(&text[span..position], &mut output);
                    output.push(VISION_START_TOKEN);
                    output.extend(std::iter::repeat_n(VISION_IMAGE_TOKEN, count));
                    output.push(VISION_END_TOKEN);
                    position = count_end + VISION_TOKEN_END.len();
                    span = position;
                    continue;
                }
            }
            let special = self
                .rendered_specials
                .iter()
                .find(|(marker, _)| bytes[position..].starts_with(marker));
            if let Some((marker, token)) = special {
                self.tokenize_plain(&text[span..position], &mut output);
                output.push(*token);
                position += marker.len();
                span = position;
            } else {
                position += 1;
            }
        }
        self.tokenize_plain(&text[span..], &mut output);
        output
    }

    fn tokenize_plain(&self, text: &str, output: &mut Vec<i32>) {
        match self.family {
            ModelFamily::DeepSeek => self.tokenize_joyai(text, output),
            ModelFamily::Glm => self.tokenize_gpt4(text, output, 3),
            ModelFamily::Qwen => self.tokenize_qwen(text, output),
        }
    }

    pub(super) fn encode_chat(
        &self,
        system_prompt: &str,
        prompt: &str,
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.encode_conversation(
            system_prompt,
            &[ChatTurn {
                user: true,
                tool: false,
                system: false,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: prompt.to_owned(),
            }],
            reasoning,
        )
    }

    pub(super) fn encode_conversation(
        &self,
        system_prompt: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.encode_messages(system_prompt, messages, reasoning, true)
    }

    pub(super) fn encode_history(
        &self,
        system_prompt: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.encode_messages(system_prompt, messages, reasoning, false)
    }

    fn encode_messages(
        &self,
        system_prompt: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
        continue_assistant: bool,
    ) -> Vec<i32> {
        if self.family == ModelFamily::Qwen {
            return self.encode_qwen_messages(
                Some(system_prompt),
                messages,
                reasoning,
                continue_assistant,
            );
        }
        let mut output = vec![self.bos];
        if self.family == ModelFamily::Glm && self.sop >= 0 {
            output.push(self.sop);
        }
        if let Some(prefix) = reasoning_prefix(self.family, reasoning) {
            if self.family == ModelFamily::Glm {
                output.push(self.system);
            }
            output.extend(self.tokenize(prefix));
        }
        if !system_prompt.is_empty() {
            if self.family == ModelFamily::Glm {
                output.push(self.system);
            }
            output.extend(self.tokenize_rendered(system_prompt));
        }
        for (index, message) in messages.iter().enumerate() {
            if message.system {
                if self.family == ModelFamily::Glm {
                    output.push(self.system);
                }
                output.extend(self.tokenize_rendered(&message.content));
                continue;
            }
            if message.tool {
                if self.family == ModelFamily::Glm {
                    output.push(self.observation);
                    output.extend(self.tokenize_rendered("<tool_response>"));
                    output.extend(
                        self.tokenize_rendered(
                            &message
                                .content
                                .replace("</tool_response>", "&lt;/tool_response>"),
                        ),
                    );
                    output.extend(self.tokenize_rendered("</tool_response>"));
                } else {
                    output.push(self.user);
                    output.extend(self.tokenize("<tool_result>"));
                    output.extend(
                        self.tokenize(
                            &message
                                .content
                                .replace("</tool_result>", "&lt;/tool_result>"),
                        ),
                    );
                    output.extend(self.tokenize("</tool_result>"));
                }
                continue;
            }
            output.push(if message.user {
                self.user
            } else {
                self.assistant
            });
            if !message.user {
                if let Some(reasoning) = &message.reasoning {
                    output.push(self.think_start);
                    output.extend(self.tokenize_rendered(reasoning));
                    if message.reasoning_complete {
                        output.push(self.think_end);
                    }
                } else if self.family == ModelFamily::Glm {
                    output.extend([self.think_start, self.think_end]);
                } else {
                    output.push(self.think_end);
                }
            }
            output.extend(self.tokenize_rendered(&message.content));
            if !message.user
                && self.family == ModelFamily::DeepSeek
                && (continue_assistant || index + 1 < messages.len())
            {
                output.push(self.eos);
            }
        }
        if continue_assistant {
            output.push(self.assistant);
            if reasoning != ReasoningMode::Direct {
                output.push(self.think_start);
            } else if self.family == ModelFamily::Glm {
                output.extend([self.think_start, self.think_end]);
            } else {
                output.push(self.think_end);
            }
        }
        output
    }

    fn encode_qwen_messages(
        &self,
        system_prompt: Option<&str>,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
        continue_assistant: bool,
    ) -> Vec<i32> {
        let messages = qwen_history(messages, reasoning);
        let rendered =
            Self::render_qwen_messages(system_prompt, &messages, reasoning, continue_assistant);
        // MTPLX preserves the generation-time BPE boundary in thinking
        // histories, including the two separate newlines of an empty thought.
        // Direct mode tokenizes the closed scaffold together.
        if reasoning != ReasoningMode::Direct
            && messages
                .iter()
                .any(|message| !message.user && !message.tool && !message.system)
        {
            let mut output = Vec::new();
            let mut start = 0;
            for (at, marker) in rendered.match_indices(QWEN_ASSISTANT_THINK_PROMPT) {
                let end = at + marker.len();
                if end < rendered.len() {
                    output.extend(self.tokenize_rendered(&rendered[start..end]));
                    start = end;
                }
            }
            output.extend(self.tokenize_rendered(&rendered[start..]));
            output
        } else {
            self.tokenize_rendered(&rendered)
        }
    }

    fn render_qwen_messages(
        system_prompt: Option<&str>,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
        continue_assistant: bool,
    ) -> String {
        let instruction = match reasoning {
            ReasoningMode::Direct | ReasoningMode::Medium => "",
            ReasoningMode::Low => {
                "Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration."
            }
            ReasoningMode::XHigh => {
                "Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer."
            }
            ReasoningMode::High | ReasoningMode::Max => {
                unreachable!("legacy reasoning modes are not valid for Qwen")
            }
        };
        let mut rendered = String::new();
        // A checkpoint continuation already contains the initial system message.
        if let Some(system_prompt) = system_prompt.map(str::trim)
            && (!instruction.is_empty() || !system_prompt.is_empty())
        {
            rendered.push_str("<|im_start|>system\n");
            if !instruction.is_empty() {
                rendered.push_str(instruction);
                if !system_prompt.is_empty() {
                    rendered.push_str("\n\n");
                }
            }
            rendered.push_str(system_prompt);
            rendered.push_str("<|im_end|>\n");
        }
        for (index, message) in messages.iter().enumerate() {
            let content = message.content.trim();
            if message.system {
                rendered.push_str("<|im_start|>system\n");
                rendered.push_str(content);
                rendered.push_str("<|im_end|>\n");
            } else if message.tool {
                if index == 0 || !messages[index - 1].tool {
                    rendered.push_str("<|im_start|>user");
                }
                rendered.push_str("\n<tool_response>\n");
                rendered.push_str(content);
                rendered.push_str("\n</tool_response>");
                if index + 1 == messages.len() || !messages[index + 1].tool {
                    rendered.push_str("<|im_end|>\n");
                }
            } else if message.user {
                rendered.push_str("<|im_start|>user\n");
                rendered.push_str(content);
                rendered.push_str("<|im_end|>\n");
            } else {
                rendered.push_str("<|im_start|>assistant\n<think>\n");
                // MTPLX's preserve-echo policy excludes reasoning_content
                // when enable_thinking is false, retaining the empty scaffold.
                if reasoning != ReasoningMode::Direct
                    && let Some(thinking) = &message.reasoning
                {
                    rendered.push_str(thinking.trim());
                }
                rendered.push_str("\n</think>\n\n");
                rendered.push_str(content);
                rendered.push_str("<|im_end|>\n");
            }
        }
        if continue_assistant {
            rendered.push_str("<|im_start|>assistant\n<think>\n");
            if reasoning == ReasoningMode::Direct {
                rendered.push_str("\n</think>\n\n");
            }
        }
        rendered
    }

    pub(super) fn encode_continuation(
        &self,
        prompt: &str,
        reasoning: ReasoningMode,
        skip_previous_eos: bool,
    ) -> Vec<i32> {
        if self.family == ModelFamily::Qwen {
            let messages = [ChatTurn {
                user: true,
                tool: false,
                system: false,
                skip_previous_eos,
                reasoning: None,
                reasoning_complete: true,
                content: prompt.to_owned(),
            }];
            return self.encode_qwen_messages(None, &messages, reasoning, true);
        }
        let mut output = Vec::new();
        if self.family == ModelFamily::DeepSeek && !skip_previous_eos {
            output.push(self.eos);
        }
        output.push(self.user);
        output.extend(self.tokenize_rendered(prompt));
        output.push(self.assistant);
        if reasoning != ReasoningMode::Direct {
            output.push(self.think_start);
        } else if self.family == ModelFamily::Glm {
            output.extend([self.think_start, self.think_end]);
        } else {
            output.push(self.think_end);
        }
        output
    }

    pub(super) fn eos(&self) -> i32 {
        self.eos
    }

    pub(super) fn is_stop(&self, token: i32) -> bool {
        token == self.eos
            || token == self.alternate_eos
            || (self.family == ModelFamily::Glm
                && [self.system, self.user, self.assistant, self.observation].contains(&token))
    }

    pub(super) fn is_think_start(&self, token: i32) -> bool {
        token == self.think_start
    }

    pub(super) fn is_think_end(&self, token: i32) -> bool {
        token == self.think_end
    }

    pub(super) fn is_ple_reset(&self, token: i32) -> bool {
        self.family == ModelFamily::Qwen
            && [self.system, self.sop, self.alternate_eos].contains(&token)
    }

    fn emit_piece(&self, raw: &[u8], output: &mut Vec<i32>) {
        if raw.is_empty() {
            return;
        }
        let encoded = byte_encode(raw);
        let mut symbols = encoded
            .char_indices()
            .map(|(start, character)| {
                encoded.as_bytes()[start..start + character.len_utf8()].to_vec()
            })
            .collect::<Vec<_>>();
        loop {
            let best = symbols
                .windows(2)
                .enumerate()
                .filter_map(|(index, pair)| {
                    let mut key = Vec::with_capacity(pair[0].len() + pair[1].len() + 1);
                    key.extend(&pair[0]);
                    key.push(b' ');
                    key.extend(&pair[1]);
                    self.merge_rank.get(&key).map(|rank| (index, *rank))
                })
                .min_by_key(|(_, rank)| *rank);
            let Some((index, _)) = best else { break };
            let right = symbols.remove(index + 1);
            symbols[index].extend(right);
        }
        for symbol in symbols {
            if let Some(token) = self.token_to_id.get(symbol.as_slice()) {
                output.push(*token);
            } else {
                output.extend(
                    symbol
                        .iter()
                        .filter_map(|byte| self.token_to_id.get(std::slice::from_ref(byte)))
                        .copied(),
                );
            }
        }
    }

    fn tokenize_joyai(&self, text: &str, output: &mut Vec<i32>) {
        let bytes = text.as_bytes();
        let mut position = 0;
        while position < bytes.len() {
            let start = position;
            let byte = bytes[position];
            if byte.is_ascii_digit() {
                let mut digits = 0;
                while position < bytes.len() && bytes[position].is_ascii_digit() && digits < 3 {
                    position += 1;
                    digits += 1;
                }
            } else if cjk_at(text, position) {
                while position < bytes.len() && cjk_at(text, position) {
                    position = next_char(text, position);
                }
            } else if ascii_punct(byte)
                && position + 1 < bytes.len()
                && bytes[position + 1].is_ascii_alphabetic()
            {
                position += 1;
                while position < bytes.len() && bytes[position].is_ascii_alphabetic() {
                    position += 1;
                }
            } else if letter_like(text, position) {
                position = consume_letters(text, position);
            } else if !ascii_newline(byte)
                && !ascii_punct(byte)
                && position + 1 < bytes.len()
                && letter_like(text, position + 1)
            {
                position = consume_letters(text, position + 1);
            } else if byte == b' ' && position + 1 < bytes.len() && ascii_punct(bytes[position + 1])
            {
                position += 1;
                while position < bytes.len() && ascii_punct(bytes[position]) {
                    position += 1;
                }
                while position < bytes.len() && ascii_newline(bytes[position]) {
                    position += 1;
                }
            } else if ascii_punct(byte) {
                while position < bytes.len() && ascii_punct(bytes[position]) {
                    position += 1;
                }
                while position < bytes.len() && ascii_newline(bytes[position]) {
                    position += 1;
                }
            } else if byte.is_ascii_whitespace() {
                let mut scan = position;
                let mut last_newline = None;
                while scan < bytes.len() && bytes[scan].is_ascii_whitespace() {
                    scan += 1;
                    if ascii_newline(bytes[scan - 1]) {
                        last_newline = Some(scan);
                    }
                }
                position = last_newline.unwrap_or_else(|| {
                    if scan < bytes.len()
                        && scan > position + 1
                        && (letter_like(text, scan) || ascii_punct(bytes[scan]))
                    {
                        scan - 1
                    } else {
                        scan
                    }
                });
            } else {
                position = next_char(text, position);
            }
            if position == start {
                position = next_char(text, position);
            }
            self.emit_piece(&bytes[start..position], output);
        }
    }

    fn tokenize_qwen(&self, text: &str, output: &mut Vec<i32>) {
        // Added tokens are extracted first by tokenize_rendered. Normalize each
        // remaining span before the trained split and BPE, as in tokenizers.
        let normalized = qwen_normalize(text);
        for piece in qwen_pieces(&normalized) {
            self.emit_piece(piece.as_bytes(), output);
        }
    }

    fn tokenize_gpt4(&self, text: &str, output: &mut Vec<i32>, max_digits: usize) {
        let mut position = 0;
        while position < text.len() {
            let start = position;
            let current = char_info(text, position);
            if current.ch == '\'' {
                let next = char_info(text, current.next);
                let lower = next.ch.to_ascii_lowercase();
                if matches!(lower, 's' | 't' | 'm' | 'd') {
                    position = next.next;
                    self.emit_piece(&text.as_bytes()[start..position], output);
                    continue;
                }
                if next.next < text.len() {
                    let next2 = char_info(text, next.next);
                    if matches!(
                        (lower, next2.ch.to_ascii_lowercase()),
                        ('r', 'e') | ('v', 'e') | ('l', 'l')
                    ) {
                        position = next2.next;
                        self.emit_piece(&text.as_bytes()[start..position], output);
                        continue;
                    }
                }
            }
            let following = char_info(text, current.next);
            if !matches!(current.ch, '\r' | '\n')
                && !current.number
                && (current.letter || following.letter)
            {
                position = current.next;
                while position < text.len() && char_info(text, position).letter {
                    position = char_info(text, position).next;
                }
            } else if current.number {
                let mut digits = 0;
                while position < text.len()
                    && char_info(text, position).number
                    && digits < max_digits
                {
                    position = char_info(text, position).next;
                    digits += 1;
                }
            } else {
                let (punctuation, punct_start) = if current.ch == ' ' {
                    (char_info(text, current.next), current.next)
                } else {
                    (current, position)
                };
                if punctuation.punctuation {
                    position = punct_start;
                    while position < text.len() && char_info(text, position).punctuation {
                        position = char_info(text, position).next;
                    }
                    while position < text.len()
                        && matches!(char_info(text, position).ch, '\r' | '\n')
                    {
                        position = char_info(text, position).next;
                    }
                } else if current.whitespace {
                    let mut scan = position;
                    let mut last_newline = None;
                    let mut last_start = position;
                    let mut count = 0;
                    while scan < text.len() && char_info(text, scan).whitespace {
                        let item = char_info(text, scan);
                        last_start = scan;
                        if matches!(item.ch, '\r' | '\n') {
                            last_newline = Some(item.next);
                        }
                        scan = item.next;
                        count += 1;
                    }
                    position = last_newline.unwrap_or(if count > 1 && scan < text.len() {
                        last_start
                    } else {
                        scan
                    });
                } else {
                    position = current.next;
                }
            }
            if position == start {
                position = next_char(text, position);
            }
            self.emit_piece(&text.as_bytes()[start..position], output);
        }
    }
}

fn qwen_normalize(text: &str) -> std::borrow::Cow<'_, str> {
    if text.is_ascii() {
        text.into()
    } else {
        text.nfc()
            .map(|(character, _)| character)
            .collect::<String>()
            .into()
    }
}

pub(super) fn qwen_normalization_changes(text: &str) -> bool {
    qwen_normalize(text) != text
}

pub(super) const QWEN_ASSISTANT_THINK_PROMPT: &str = "<|im_start|>assistant\n<think>\n";

pub(super) fn qwen_void_assistant(message: &ChatTurn) -> bool {
    !message.user && !message.tool && !message.system && message.content.trim().is_empty()
}

fn qwen_history(messages: &[ChatTurn], mode: ReasoningMode) -> Vec<ChatTurn> {
    // MTPLX drops void assistants, then merges adjacent user/assistant roles.
    // The server also drops reasoning-only assistants. Keep the
    // first assistant's reasoning when merging; trim user text after merging.
    let mut normalized: Vec<ChatTurn> = Vec::with_capacity(messages.len());
    for message in messages {
        if qwen_void_assistant(message) {
            continue;
        }
        let mut message = message.clone();
        if !message.user && !message.tool && !message.system {
            message.content = message.content.trim().to_owned();
            if mode == ReasoningMode::Direct {
                message.reasoning = None;
            }
        }
        if let Some(previous) = normalized.last_mut()
            && !previous.tool
            && !previous.system
            && !message.tool
            && !message.system
            && previous.user == message.user
        {
            if !message.content.is_empty() {
                if !previous.content.is_empty() {
                    previous.content.push_str("\n\n");
                }
                previous.content.push_str(&message.content);
            }
        } else {
            normalized.push(message);
        }
    }
    normalized
}

/// The trained Qwen3 Split expression. Rust regex supports its Unicode
/// categories but not the penultimate `\s+(?!\S)` alternative: before a
/// non-whitespace character that branch backtracks by one whitespace scalar,
/// leaving it for the next word/punctuation match. At EOF it keeps the run.
fn qwen_pieces(text: &str) -> impl Iterator<Item = &str> {
    static SPLIT: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
        regex::Regex::new(
            r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|(?P<whitespace>\s+)",
        ).expect("valid Qwen3 pretokenizer expression")
    });
    let mut position = 0;
    std::iter::from_fn(move || {
        if position == text.len() {
            return None;
        }
        let captures = SPLIT
            .captures_at(text, position)
            .expect("Qwen3 covers every Unicode scalar");
        let matched = captures.get(0).unwrap();
        debug_assert_eq!(matched.start(), position);
        let mut end = matched.end();
        if captures.name("whitespace").is_some() && end < text.len() {
            let last = matched.as_str().char_indices().next_back().unwrap().0;
            if last != 0 {
                end = position + last;
            }
        }
        let piece = &text[position..end];
        position = end;
        Some(piece)
    })
}

fn byte_encode(bytes: &[u8]) -> String {
    bytes
        .iter()
        .map(|byte| char::from_u32(gpt2_byte_to_codepoint(*byte)).unwrap())
        .collect()
}

fn gpt2_byte_to_codepoint(byte: u8) -> u32 {
    if (33..=126).contains(&byte) || (161..=172).contains(&byte) || byte >= 174 {
        return u32::from(byte);
    }
    let mut index = 0;
    for candidate in 0..=255_u16 {
        let candidate = candidate as u8;
        if (33..=126).contains(&candidate) || (161..=172).contains(&candidate) || candidate >= 174 {
            continue;
        }
        if candidate == byte {
            return 256 + index;
        }
        index += 1;
    }
    u32::from(byte)
}

fn gpt2_codepoint_to_byte(character: char) -> Option<u8> {
    let codepoint = character as u32;
    if (33..=126).contains(&codepoint)
        || (161..=172).contains(&codepoint)
        || (174..=255).contains(&codepoint)
    {
        return Some(codepoint as u8);
    }
    let mut index = 0;
    for byte in 0..=255_u16 {
        let byte = byte as u8;
        if (33..=126).contains(&byte) || (161..=172).contains(&byte) || byte >= 174 {
            continue;
        }
        if codepoint == 256 + index {
            return Some(byte);
        }
        index += 1;
    }
    None
}

fn next_char(text: &str, position: usize) -> usize {
    match text.get(position..).and_then(|rest| rest.chars().next()) {
        Some(ch) => position + ch.len_utf8(),
        None => position.saturating_add(1).min(text.len()),
    }
}

fn cjk_at(text: &str, position: usize) -> bool {
    let Some(&byte) = text.as_bytes().get(position) else {
        return false;
    };
    if byte < 128 {
        return false;
    }
    let Some(character) = text
        .get(position..)
        .and_then(|rest| rest.chars().next())
        .map(|ch| ch as u32)
    else {
        return false;
    };
    (0x4e00..=0x9fa5).contains(&character)
        || (0x3040..=0x309f).contains(&character)
        || (0x30a0..=0x30ff).contains(&character)
}

fn letter_like(text: &str, position: usize) -> bool {
    // Non-ASCII bytes count as letters; CJK and kana are isolated first by cjk_at.
    let byte = text.as_bytes()[position];
    byte >= 128 || byte.is_ascii_alphabetic()
}

fn consume_letters(text: &str, mut position: usize) -> usize {
    while position < text.len() && letter_like(text, position) {
        position = next_char(text, position);
    }
    position
}

fn ascii_punct(byte: u8) -> bool {
    byte.is_ascii_punctuation()
}

fn ascii_newline(byte: u8) -> bool {
    matches!(byte, b'\n' | b'\r')
}

#[derive(Clone, Copy)]
struct CharInfo {
    ch: char,
    next: usize,
    letter: bool,
    number: bool,
    whitespace: bool,
    punctuation: bool,
}

fn char_info(text: &str, position: usize) -> CharInfo {
    let Some(ch) = text.get(position..).and_then(|rest| rest.chars().next()) else {
        return CharInfo {
            ch: '\0',
            next: position.saturating_add(1).min(text.len()),
            letter: false,
            number: false,
            whitespace: false,
            punctuation: false,
        };
    };
    let codepoint = ch as u32;
    let whitespace = unicode_whitespace(codepoint);
    let number = unicode_number(codepoint);
    let punctuation = unicode_punctuation(codepoint);
    CharInfo {
        ch,
        next: position + ch.len_utf8(),
        letter: if codepoint < 128 {
            ch.is_ascii_alphabetic()
        } else {
            !whitespace && !number && !punctuation
        },
        number,
        whitespace,
        punctuation,
    }
}

fn unicode_whitespace(cp: u32) -> bool {
    (cp < 128 && (cp as u8).is_ascii_whitespace())
        || matches!(
            cp,
            0x0085 | 0x00a0 | 0x1680 | 0x2028 | 0x2029 | 0x202f | 0x205f | 0x3000
        )
        || (0x2000..=0x200a).contains(&cp)
}

fn unicode_number(cp: u32) -> bool {
    (cp < 128 && (cp as u8).is_ascii_digit())
        || [
            0x0660..=0x0669,
            0x06f0..=0x06f9,
            0x07c0..=0x07c9,
            0x0966..=0x096f,
            0x09e6..=0x09ef,
            0x0a66..=0x0a6f,
            0x0ae6..=0x0aef,
            0x0b66..=0x0b6f,
            0x0be6..=0x0bef,
            0x0c66..=0x0c6f,
            0x0ce6..=0x0cef,
            0x0d66..=0x0d6f,
            0x0de6..=0x0def,
            0x0e50..=0x0e59,
            0x0ed0..=0x0ed9,
            0x0f20..=0x0f29,
            0x1040..=0x1049,
            0x1090..=0x1099,
            0x17e0..=0x17e9,
            0x1810..=0x1819,
            0xff10..=0xff19,
        ]
        .iter()
        .any(|range| range.contains(&cp))
}

fn unicode_punctuation(cp: u32) -> bool {
    if cp < 128 {
        return (cp as u8).is_ascii_punctuation();
    }
    matches!(
        cp,
        0x00b4
            | 0x00bb
            | 0x00bf
            | 0x00d7
            | 0x00f7
            | 0x0387
            | 0x05c3
            | 0x061b
            | 0x061e
            | 0x061f
            | 0x066a
            | 0x066d
            | 0x06d4
    ) || [
        0x00a1..=0x00a9,
        0x00ab..=0x00ac,
        0x00ae..=0x00b1,
        0x00b6..=0x00b8,
        0x02c2..=0x02df,
        0x02e5..=0x02eb,
        0x02ed..=0x02ff,
        0x0375..=0x037e,
        0x0384..=0x0385,
        0x055a..=0x055f,
        0x0589..=0x058a,
        0x05be..=0x05c0,
        0x05c6..=0x05c7,
        0x0609..=0x060a,
        0x060c..=0x060d,
        0x2000..=0x206f,
        0x20a0..=0x20cf,
        0x2100..=0x214f,
        0x2190..=0x23ff,
        0x2460..=0x24ff,
        0x2500..=0x2775,
        0x2794..=0x2bff,
        0x2e00..=0x2e7f,
        0x3000..=0x303f,
        0xfd3e..=0xfd3f,
        0xfe10..=0xfe6f,
        0xff01..=0xff0f,
        0xff1a..=0xff20,
        0xff3b..=0xff40,
        0xff5b..=0xff65,
        0x1f000..=0x1faff,
    ]
    .iter()
    .any(|range| range.contains(&cp))
}

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

    #[test]
    #[ignore = "CPU-only; requires DS4SERVER_CHAT_REFERENCE and its installed GGUF"]
    fn ds4_chat_matches_original_session_tokens() {
        use super::super::{Model, ModelRef, conversation_tag, render_text_prompt};
        use crate::model::ModelChoice;
        let reference = std::env::var_os("DS4SERVER_CHAT_REFERENCE").unwrap();
        let fixture: serde_json::Value =
            serde_json::from_slice(&fs::read(reference).unwrap()).unwrap();
        let choice = match fixture["family"].as_str().unwrap() {
            "deepseek" => ModelChoice::DeepSeekV4Flash0731,
            "glm" => ModelChoice::Glm53Flash,
            other => panic!("unsupported reference family: {other}"),
        };
        let model =
            Model::open_main(Path::new(fixture["model"].as_str().unwrap()), choice).unwrap();
        let tokenizer = &model.tokenizer;
        let cases = fixture["cases"].as_array().unwrap();
        assert_eq!(cases.len(), 3);
        let mut previous = Vec::new();
        for (index, case) in cases.iter().enumerate() {
            let prompt = case["prompt"].as_str().unwrap();
            let expected: Vec<i32> = serde_json::from_value(case["tokens"].clone()).unwrap();
            let actual = if index == 0 {
                tokenizer.encode_chat("", prompt, ReasoningMode::Low)
            } else {
                let mut tokens = previous;
                tokens.extend(tokenizer.encode_continuation(prompt, ReasoningMode::Low, false));
                tokens
            };
            if actual != expected {
                let first = actual
                    .iter()
                    .zip(&expected)
                    .position(|(a, b)| a != b)
                    .unwrap_or(actual.len().min(expected.len()));
                let context = |tokens: &[i32]| {
                    tokens[first.saturating_sub(4)..tokens.len().min(first + 8)]
                        .iter()
                        .map(|&id| {
                            (
                                id,
                                String::from_utf8_lossy(&tokenizer.token_bytes(id).unwrap())
                                    .into_owned(),
                            )
                        })
                        .collect::<Vec<_>>()
                };
                panic!(
                    "turn {} first mismatch {first}, lengths {}/{}; Rust {:?}; DS4 {:?}",
                    index + 1,
                    actual.len(),
                    expected.len(),
                    context(&actual),
                    context(&expected)
                );
            }
            if index == 0 {
                let settings = crate::settings::TurnSettings {
                    kv_cache: crate::settings::KvCachePreferences::default().settings(),
                    context_tokens: 32768,
                    max_generated_tokens: i32::MAX,
                    system_prompt: String::new(),
                    temperature: 0.6,
                    top_p: 0.95,
                    min_p: 0.0,
                    top_k: 0,
                    stops: vec![],
                    seed: Some(42),
                    reasoning_mode: ReasoningMode::Low,
                };
                let messages = [ChatTurn {
                    user: true,
                    tool: false,
                    system: false,
                    skip_previous_eos: false,
                    reasoning: None,
                    reasoning_complete: true,
                    content: prompt.to_owned(),
                }];
                let bootstrap = model.render_history("", &[], settings.reasoning_mode);
                let rendered = render_text_prompt(
                    ModelRef::Gguf(&model),
                    &bootstrap,
                    conversation_tag(model.shape.family, "", settings.reasoning_mode, &[]),
                    &messages,
                    &settings,
                );
                assert!(
                    rendered == expected,
                    "shared bootstrap renderer differs: {}/{} tokens; starts {:?}/{:?}",
                    rendered.len(),
                    expected.len(),
                    &rendered[..4],
                    &expected[..4]
                );
            }
            previous = expected;
            let reply: Vec<i32> = serde_json::from_value(case["reply_tokens"].clone()).unwrap();
            previous.extend(reply);
        }
    }

    #[test]
    #[ignore = "requires local tokenizer and tools/qwen-chat-reference.py output; no inference"]
    fn qwen_plain_chat_matches_mtplx_server_tokens() {
        let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE").unwrap();
        let reference = std::env::var_os("DS4SERVER_QWEN38_CHAT_REFERENCE").unwrap();
        let fixture: serde_json::Value =
            serde_json::from_slice(&fs::read(reference).unwrap()).unwrap();
        let tokenizer = Tokenizer::load_qwen(&Path::new(&root).join("tokenizer.json")).unwrap();
        let cases = fixture["cases"].as_array().unwrap();
        assert_eq!(cases.len(), 3);
        for (index, case) in cases.iter().enumerate() {
            let messages = case["messages"]
                .as_array()
                .unwrap()
                .iter()
                .map(|message| ChatTurn {
                    user: message["role"] == "user",
                    tool: false,
                    system: false,
                    skip_previous_eos: false,
                    reasoning: message["reasoning_content"].as_str().map(str::to_owned),
                    reasoning_complete: true,
                    content: message["content"].as_str().unwrap().to_owned(),
                })
                .collect::<Vec<_>>();
            let actual = tokenizer.encode_conversation("", &messages, ReasoningMode::Low);
            let rendered = actual
                .iter()
                .flat_map(|&id| tokenizer.token_bytes(id).unwrap())
                .collect::<Vec<_>>();
            assert_eq!(
                rendered,
                case["rendered"].as_str().unwrap().as_bytes(),
                "turn {} rendered bytes",
                index + 1
            );
            let expected: Vec<i32> = serde_json::from_value(case["token_ids"].clone()).unwrap();
            assert!(
                actual == expected,
                "turn {} token mismatch: actual={}, reference={}, first_difference={:?}",
                index + 1,
                actual.len(),
                expected.len(),
                actual.iter().zip(&expected).position(|(a, b)| a != b)
            );
            let mut resumed =
                tokenizer.encode_history("", &messages[..messages.len() - 1], ReasoningMode::Low);
            resumed.extend(tokenizer.encode_continuation(
                &messages.last().unwrap().content,
                ReasoningMode::Low,
                false,
            ));
            assert!(resumed == expected, "turn {} checkpoint tokens", index + 1);
        }
    }

    #[test]
    #[ignore = "requires the installed Qwen tokenizer, but no GPU or model weights"]
    fn qwen_checkpoint_continuation_matches_full_conversation() {
        let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE")
            .expect("set DS4SERVER_QWEN38_SOURCE to the installed artifact directory");
        let tokenizer = Tokenizer::load_qwen(&Path::new(&root).join("tokenizer.json")).unwrap();
        let user = ChatTurn {
            user: true,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: "ping".into(),
        };
        let assistant = ChatTurn {
            user: false,
            reasoning: Some("Brief thought.".into()),
            content: "pong".into(),
            ..user.clone()
        };
        for reasoning in [
            ReasoningMode::Direct,
            ReasoningMode::Low,
            ReasoningMode::Medium,
            ReasoningMode::XHigh,
        ] {
            for mut history in [vec![], vec![user.clone(), assistant.clone()]] {
                let mut resumed =
                    tokenizer.encode_history("System instruction", &history, reasoning);
                resumed.extend(tokenizer.encode_continuation("Next question", reasoning, false));
                history.push(ChatTurn {
                    content: "Next question".into(),
                    ..user.clone()
                });
                assert_eq!(
                    resumed,
                    tokenizer.encode_conversation("System instruction", &history, reasoning),
                    "checkpoint/full prompt mismatch: {reasoning:?}, {} messages",
                    history.len(),
                );
            }
        }
    }

    #[test]
    fn character_helpers_are_total_at_every_byte_offset() {
        let text = "Aé中😊";
        for position in 0..=text.len() + 1 {
            let next = next_char(text, position);
            assert!(next <= text.len());
            if position < text.len() {
                assert!(next > position);
            }
            assert!(char_info(text, position).next <= text.len());
            let _ = cjk_at(text, position);
        }

        assert_eq!(next_char(text, 1), 3);
        assert!(cjk_at(text, 3));
        assert!(!cjk_at(text, 6));
        assert_eq!(next_char("", 0), 0);
        assert!(!cjk_at("", 0));
    }

    #[test]
    fn reasoning_prefixes_match_each_model_family_contract() {
        assert_eq!(
            reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::Low),
            None
        );
        assert_eq!(
            reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::High),
            Some(HIGH_REASONING_PREFIX)
        );
        assert_eq!(
            reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::Max),
            Some(MAX_REASONING_PREFIX)
        );
        assert_eq!(
            reasoning_prefix(ModelFamily::Glm, ReasoningMode::Low),
            Some("Reasoning Effort: Low")
        );
        assert_eq!(
            reasoning_prefix(ModelFamily::Glm, ReasoningMode::High),
            Some("Reasoning Effort: High")
        );
        assert_eq!(
            reasoning_prefix(ModelFamily::Glm, ReasoningMode::Max),
            Some("Reasoning Effort: Max")
        );
    }
}

#[cfg(test)]
#[path = "tokenizer_issue111.rs"]
mod issue111;

#[cfg(test)]
#[path = "tokenizer_issue116.rs"]
mod issue116;

#[cfg(test)]
#[path = "tokenizer_issue117.rs"]
mod issue117;