use regex::Regex;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::Path;

pub struct Tokenizer {
    encoder: HashMap<Vec<u8>, u32>,
    decoder: Vec<Vec<u8>>,
    merge_ranks: HashMap<(u32, u32), usize>,
    special_tokens: HashMap<String, u32>,
    #[allow(dead_code)]
    special_token_ids: HashMap<u32, String>,
    pre_tokenize_re: Regex,
    eos_token_id: Option<u32>,
    eos_token_ids: Vec<u32>,
    byte_fallback: bool,
}

#[derive(Deserialize)]
struct TokenizerJson {
    model: ModelSection,
    #[serde(default)]
    added_tokens: Vec<AddedToken>,
    #[serde(default)]
    pre_tokenizer: Option<PreTokenizerSection>,
}

#[derive(Deserialize)]
struct PreTokenizerSection {
    #[serde(default, rename = "type")]
    kind: Option<String>,
    #[serde(default)]
    pattern: Option<PatternSpec>,
    #[serde(default)]
    pretokenizers: Option<Vec<PreTokenizerSection>>,
}

#[derive(Deserialize)]
struct PatternSpec {
    #[serde(rename = "Regex")]
    regex: Option<String>,
}

#[derive(Deserialize)]
struct ModelSection {
    vocab: HashMap<String, u32>,
    merges: Vec<MergeEntry>,
    #[serde(default)]
    byte_fallback: bool,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum MergeEntry {
    Str(String),
    Pair(Vec<String>),
}

#[derive(Deserialize)]
struct AddedToken {
    id: u32,
    content: String,
    #[allow(dead_code)]
    special: bool,
}

impl Tokenizer {
    pub fn from_file(path: &Path) -> Self {
        let data = std::fs::read_to_string(path)
            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
        let tj: TokenizerJson = serde_json::from_str(&data)
            .unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));

        // Build encoder: token bytes → ID
        // All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
        let mut encoder = HashMap::new();
        for (token_str, &id) in &tj.model.vocab {
            let bytes = token_str_to_bytes(token_str);
            encoder.insert(bytes, id);
        }

        // Build decoder: ID → token bytes
        let max_id = tj.model.vocab.values().copied().max().unwrap_or(0);
        let added_max = tj.added_tokens.iter().map(|t| t.id).max().unwrap_or(0);
        let vocab_size = (max_id.max(added_max) + 1) as usize;
        let mut decoder = vec![vec![]; vocab_size];
        for (token_str, &id) in &tj.model.vocab {
            decoder[id as usize] = token_str_to_bytes(token_str);
        }

        // Parse merges (supports both "a b" string format and ["a", "b"] array format)
        let byte_fallback = tj.model.byte_fallback;
        let mut merge_ranks = HashMap::new();
        for (rank, entry) in tj.model.merges.iter().enumerate() {
            let (a_str, b_str) = match entry {
                MergeEntry::Str(s) => {
                    let parts: Vec<&str> = s.splitn(2, ' ').collect();
                    if parts.len() != 2 {
                        continue;
                    }
                    (parts[0].to_string(), parts[1].to_string())
                }
                MergeEntry::Pair(v) => {
                    if v.len() != 2 {
                        continue;
                    }
                    (v[0].clone(), v[1].clone())
                }
            };
            let a_bytes = token_str_to_bytes(&a_str);
            let b_bytes = token_str_to_bytes(&b_str);
            if let (Some(&a_id), Some(&b_id)) = (encoder.get(&a_bytes), encoder.get(&b_bytes)) {
                merge_ranks.insert((a_id, b_id), rank);
            }
        }

        // Added tokens are matched as indivisible tokens by HF tokenizers,
        // even when their `special` flag is false (for example Qwen3's
        // <think> and </think> tokens).
        let mut special_tokens = HashMap::new();
        let mut special_token_ids = HashMap::new();
        for at in &tj.added_tokens {
            special_tokens.insert(at.content.clone(), at.id);
            special_token_ids.insert(at.id, at.content.clone());
            decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
            decoder[at.id as usize] = at.content.as_bytes().to_vec();
        }
        // End-of-generation tokens, in priority order. Families differ:
        // Qwen uses <|im_end|>, Llama <|end_of_text|>, GPT-2 <|endoftext|>.
        // gpt-oss (harmony) ends the assistant turn with <|return|> and also
        // treats <|call|> (tool call) and <|endoftext|> as terminators
        // (see generation_config.json eos_token_id = [200002, 199999, 200012]).
        let eos_names = [
            "<|im_end|>",
            "<|end_of_text|>",
            "<|return|>",
            "<|call|>",
            "<|endoftext|>",
        ];
        let mut eos_token_ids: Vec<u32> = Vec::new();
        for name in eos_names {
            if let Some(&id) = special_tokens.get(name) {
                if !eos_token_ids.contains(&id) {
                    eos_token_ids.push(id);
                }
            }
        }
        let eos_token_id = eos_token_ids.first().copied();

        // Pre-tokenization regex: prefer the model's own regex from tokenizer.json,
        // fall back to GPT-2/Qwen heuristic if not present or unsupported.
        let model_regex = tj.pre_tokenizer.as_ref().and_then(|pt| {
            // Direct Split with regex
            if pt.kind.as_deref() == Some("Split") {
                return pt.pattern.as_ref().and_then(|p| p.regex.clone());
            }
            // Sequence → find the Split entry
            if let Some(subs) = &pt.pretokenizers {
                for sub in subs {
                    if sub.kind.as_deref() == Some("Split") {
                        if let Some(r) = sub.pattern.as_ref().and_then(|p| p.regex.clone()) {
                            return Some(r);
                        }
                    }
                }
            }
            None
        });

        let pre_tokenize_re = if let Some(ref pat) = model_regex {
            // Strip unsupported lookahead (?!\S) — Rust regex doesn't support it.
            // The lookahead only affects trailing-whitespace edge cases.
            let cleaned = pat.replace(r"(?!\S)", "");
            match Regex::new(&cleaned) {
                Ok(re) => re,
                Err(e) => {
                    eprintln!("warning: model pre_tokenizer regex failed ({e}), using fallback");
                    if byte_fallback {
                        Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
                    } else {
                        Regex::new(
                            r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+",
                        )
                        .unwrap()
                    }
                }
            }
        } else if byte_fallback {
            Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
        } else {
            Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
        };

        Self {
            encoder,
            decoder,
            merge_ranks,
            special_tokens,
            special_token_ids,
            pre_tokenize_re,
            eos_token_id,
            eos_token_ids,
            byte_fallback,
        }
    }

    pub fn encode(&self, text: &str) -> Vec<u32> {
        let mut tokens = Vec::new();

        // Check for special tokens first (split around them)
        let mut remaining = text;
        while !remaining.is_empty() {
            // Find earliest special token
            let mut earliest: Option<(usize, &str, u32)> = None;
            for (st, &id) in &self.special_tokens {
                if let Some(pos) = remaining.find(st.as_str()) {
                    if earliest.is_none() || pos < earliest.unwrap().0 {
                        earliest = Some((pos, st, id));
                    }
                }
            }

            if let Some((pos, st, id)) = earliest {
                if pos > 0 {
                    self.encode_ordinary(&remaining[..pos], &mut tokens);
                }
                tokens.push(id);
                remaining = &remaining[pos + st.len()..];
            } else {
                self.encode_ordinary(remaining, &mut tokens);
                break;
            }
        }

        tokens
    }

    fn encode_ordinary(&self, text: &str, out: &mut Vec<u32>) {
        for mat in self.pre_tokenize_re.find_iter(text) {
            let word = mat.as_str();
            // Try to encode the whole word first
            if let Some(&id) = self.encoder.get(word.as_bytes()) {
                out.push(id);
                continue;
            }
            // Fall back to per-byte encoding
            let word_bytes: Vec<u8> = word.bytes().collect();
            let mut token_ids: Vec<u32> = word_bytes.iter().filter_map(|&b| {
                if let Some(&id) = self.encoder.get(&vec![b]) {
                    Some(id)
                } else if self.byte_fallback {
                    let hex_token = format!("<0x{:02X}>", b);
                    if let Some(&id) = self.special_tokens.get(&hex_token) {
                        Some(id)
                    } else if let Some(&id) = self.encoder.get(hex_token.as_bytes()) {
                        Some(id)
                    } else if let Some(&unk_id) = self.special_tokens.get("<unk>") {
                        eprintln!("warning: byte 0x{b:02X} not in vocab, using <unk> token");
                        Some(unk_id)
                    } else {
                        eprintln!("warning: byte 0x{b:02X} not in vocab and no fallback token, using token 0");
                        Some(0)
                    }
                } else {
                    eprintln!("warning: byte {b} (0x{b:02X}) not in vocab, skipping");
                    None
                }
            }).collect();

            // BPE merges
            loop {
                if token_ids.len() < 2 {
                    break;
                }
                let mut best_rank = usize::MAX;
                let mut best_idx = 0;
                for i in 0..token_ids.len() - 1 {
                    if let Some(&rank) = self.merge_ranks.get(&(token_ids[i], token_ids[i + 1])) {
                        if rank < best_rank {
                            best_rank = rank;
                            best_idx = i;
                        }
                    }
                }
                if best_rank == usize::MAX {
                    break;
                }

                let merged_bytes = [
                    self.decoder[token_ids[best_idx] as usize].as_slice(),
                    self.decoder[token_ids[best_idx + 1] as usize].as_slice(),
                ]
                .concat();
                let merged_id = *self.encoder.get(&merged_bytes).unwrap_or_else(|| {
                    panic!("merged token not in vocab");
                });
                token_ids[best_idx] = merged_id;
                token_ids.remove(best_idx + 1);
            }

            out.extend_from_slice(&token_ids);
        }
    }

    pub fn decode(&self, token_ids: &[u32]) -> String {
        let mut bytes = Vec::new();
        for &id in token_ids {
            if let Some(b) = self.decoder.get(id as usize) {
                bytes.extend_from_slice(b);
            }
        }
        String::from_utf8_lossy(&bytes).into_owned()
    }

    pub fn decode_token_stream(&self, token_id: u32, pending: &mut Vec<u8>) -> String {
        if let Some(bytes) = self.decoder.get(token_id as usize) {
            pending.extend_from_slice(bytes);
        }
        take_valid_utf8(pending)
    }

    pub fn flush_decode_stream(&self, pending: &mut Vec<u8>) -> String {
        let text = String::from_utf8_lossy(pending).into_owned();
        pending.clear();
        text
    }

    pub fn eos_token_id(&self) -> Option<u32> {
        self.eos_token_id
    }

    /// True if `id` is any end-of-generation token (a model may have several;
    /// gpt-oss/harmony ends on <|return|>, <|call|>, or <|endoftext|>).
    pub fn is_eos(&self, id: u32) -> bool {
        self.eos_token_ids.contains(&id)
    }

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

    pub fn special_token_id(&self, name: &str) -> Option<u32> {
        self.special_tokens.get(name).copied()
    }
}

/// Convert a token string from HF vocab (which uses Unicode replacements for bytes)
/// back to raw bytes. GPT-2 uses a byte-to-unicode mapping where e.g. byte 0x20 (space)
/// is represented as 'Ġ' (U+0120).
fn token_str_to_bytes(s: &str) -> Vec<u8> {
    s.chars().map(|c| unicode_to_byte(c)).collect()
}

fn take_valid_utf8(pending: &mut Vec<u8>) -> String {
    match std::str::from_utf8(pending) {
        Ok(text) => {
            let text = text.to_string();
            pending.clear();
            text
        }
        Err(err) => {
            let valid_up_to = err.valid_up_to();
            if valid_up_to == 0 {
                if let Some(error_len) = err.error_len() {
                    let invalid_len = error_len.min(pending.len());
                    let text = String::from_utf8_lossy(&pending[..invalid_len]).into_owned();
                    pending.drain(..invalid_len);
                    return text;
                }
                return String::new();
            }
            let text = String::from_utf8_lossy(&pending[..valid_up_to]).into_owned();
            pending.drain(..valid_up_to);
            text
        }
    }
}

/// Convert a Unicode char back to the byte it represents in GPT-2 encoding.
fn unicode_to_byte(c: char) -> u8 {
    // Build the inverse map on first use
    use std::sync::OnceLock;
    static INV_MAP: OnceLock<HashMap<u32, u8>> = OnceLock::new();

    let map = INV_MAP.get_or_init(|| {
        let mut m = HashMap::new();
        // Build GPT-2's bytes_to_unicode forward map, then invert
        let mut n = 0u32;
        for b in 0..=255u16 {
            let byte = b as u8;
            let unicode = match byte {
                0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => byte as u32,
                _ => {
                    let u = 256 + n;
                    n += 1;
                    u
                }
            };
            m.insert(unicode, byte);
        }
        m
    });

    *map.get(&(c as u32))
        .unwrap_or_else(|| panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32))
}

#[cfg(test)]
mod tests {
    use super::{Tokenizer, take_valid_utf8};

    #[test]
    fn qwen_added_tokens_are_indivisible_and_im_end_is_eos() {
        let path =
            std::env::temp_dir().join(format!("xserv-tokenizer-test-{}.json", std::process::id()));
        std::fs::write(
            &path,
            r#"{
                "model": {
                    "vocab": {},
                    "merges": [],
                    "byte_fallback": false
                },
                "added_tokens": [
                    {"id":151643,"content":"<|endoftext|>","special":true},
                    {"id":151644,"content":"<|im_start|>","special":true},
                    {"id":151645,"content":"<|im_end|>","special":true},
                    {"id":151667,"content":"<think>","special":false},
                    {"id":151668,"content":"</think>","special":false}
                ]
            }"#,
        )
        .unwrap();

        let tokenizer = Tokenizer::from_file(&path);
        let _ = std::fs::remove_file(&path);

        assert_eq!(tokenizer.eos_token_id(), Some(151645));
        assert_eq!(tokenizer.encode("<think>"), vec![151667]);
        assert_eq!(tokenizer.encode("</think>"), vec![151668]);
        assert_eq!(tokenizer.decode(&[151645]), "<|im_end|>");
    }

    #[test]
    fn stream_decode_buffers_incomplete_utf8() {
        let mut pending = vec![0xF0, 0x9F];
        assert_eq!(take_valid_utf8(&mut pending), "");
        pending.extend_from_slice(&[0x98, 0x8A, b'!']);
        assert_eq!(take_valid_utf8(&mut pending), "😊!");
        assert!(pending.is_empty());
    }
}