// 🌌 [BioPhys 5.0 네이티브 BPE 바이트페어 토크나이저 엔진] (src/native_bpe_tokenizer.rs)
// 외부 파이썬 의존성 0%! 임의의 한국어/영어/기호 텍스트를 BPE 바이트 단위로 분절 및 정수 ID 인코딩/디코딩
use std::collections::HashMap;

pub struct NativeBpeTokenizer {
    pub token_to_id: HashMap<String, usize>,
    pub id_to_token: Vec<String>,
    pub byte_tokens: Vec<String>,
}

impl NativeBpeTokenizer {
    pub fn new() -> Self {
        let mut token_to_id = HashMap::new();
        let mut id_to_token = Vec::new();

        // 1. 특수 토큰 등록
        let special_tokens = vec!["<PAD>", "<BOS>", "<EOS>", "<UNK>", "<MASK>"];
        for &tok in &special_tokens {
            let id = id_to_token.len();
            id_to_token.push(tok.to_string());
            token_to_id.insert(tok.to_string(), id);
        }

        // 2. 256개 기본 바이트 토큰 등록 (0x00 ~ 0xFF)
        let mut byte_tokens = Vec::new();
        for b in 0..=255u8 {
            let tok_str = format!("<0x{:02X}>", b);
            let id = id_to_token.len();
            id_to_token.push(tok_str.clone());
            token_to_id.insert(tok_str.clone(), id);
            byte_tokens.push(tok_str);
        }

        // 3. 주요 한국어 음절 및 기술 키워드 서브워드 등록
        let common_korean_subwords = vec![
            "인공지능", "블랙홀", "사건의", "지평선", "스파이킹", "양자화", "뉴런",
            "가중치", "압축", "초저지연", "엔진", "Rust", "Svelte", "Tauri",
            "한국어", "문법", "항상성", "특이점", "물리", "동역학", "입니다", "합니다",
            "으로", "에서", "의", "을", "를", "이", "가", "은", "는", "과", "와"
        ];

        for &word in &common_korean_subwords {
            let id = id_to_token.len();
            id_to_token.push(word.to_string());
            token_to_id.insert(word.to_string(), id);
        }

        NativeBpeTokenizer {
            token_to_id,
            id_to_token,
            byte_tokens,
        }
    }

    /// [텍스트 ➔ 정규 BPE 토큰 ID 벡터 인코딩]
    pub fn encode(&self, text: &str) -> Vec<usize> {
        let mut token_ids = Vec::new();
        token_ids.push(1); // <BOS>

        let mut chars = text.chars().peekable();
        let mut current_buf = String::new();

        while let Some(c) = chars.next() {
            current_buf.push(c);
            
            // 등록된 서브워드인지 탐색
            if let Some(&id) = self.token_to_id.get(&current_buf) {
                token_ids.push(id);
                current_buf.clear();
            } else if c.is_whitespace() || chars.peek().is_none() {
                // 공백 또는 단어 끝에서 바이트 단위 폴백
                for b in current_buf.as_bytes() {
                    let byte_tok = format!("<0x{:02X}>", b);
                    let id = self.token_to_id.get(&byte_tok).copied().unwrap_or(3); // <UNK>
                    token_ids.push(id);
                }
                current_buf.clear();
            }
        }

        token_ids.push(2); // <EOS>
        token_ids
    }

    /// [토큰 ID 벡터 ➔ 사람이 읽을 수 있는 UTF-8 텍스트 디코딩]
    pub fn decode(&self, token_ids: &[usize]) -> String {
        let mut byte_buffer = Vec::new();
        let mut decoded_text = String::new();

        for &id in token_ids {
            if id >= self.id_to_token.len() { continue; }
            let tok_str = &self.id_to_token[id];

            if tok_str == "<BOS>" || tok_str == "<EOS>" || tok_str == "<PAD>" {
                continue;
            }

            if tok_str.starts_with("<0x") && tok_str.ends_with(">") {
                if let Ok(b) = u8::from_str_radix(&tok_str[3..5], 16) {
                    byte_buffer.push(b);
                }
            } else {
                if !byte_buffer.is_empty() {
                    decoded_text.push_str(&String::from_utf8_lossy(&byte_buffer));
                    byte_buffer.clear();
                }
                decoded_text.push_str(tok_str);
            }
        }

        if !byte_buffer.is_empty() {
            decoded_text.push_str(&String::from_utf8_lossy(&byte_buffer));
        }

        decoded_text
    }
}