"""
Gemma3 Tokenizer.

Wraps the HuggingFace tokenizers library with Gemma3-specific special
token handling. Uses a hybrid tokenization approach: special tokens
(<start_of_turn>, <end_of_turn>, <think>, etc.) are mapped directly
to their IDs, while regular text is encoded by the underlying BPE
tokenizer with add_special_tokens=False to prevent unwanted <bos>
insertion between fragments.

Chat template follows Gemma3's turn-based format:
  <start_of_turn>user\\n{content}<end_of_turn>\\n
  <start_of_turn>model\\n
"""

import re
import os
from typing import List, Dict, Union
from pathlib import Path
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
from src.utils.model_store import get_model_dir

class Gemma3Tokenizer:
    # 1. Define Gemma-specific specials
    _SPECIALS = [
        "<bos>", "<eos>", "<unk>", "<pad>",
        "<start_of_turn>", "<end_of_turn>",
        "<think>", "</think>" 
    ]

    def __init__(self, model_path_or_repo: str):
        # ... (Same init logic as before) ...
        # 1. Handle Auto-Download
        if os.path.exists(model_path_or_repo):
            if (Path(model_path_or_repo) / "tokenizer.json").exists():
                file_path = Path(model_path_or_repo) / "tokenizer.json"
            else:
                file_path = Path(model_path_or_repo) / "tokenizer.model"
        else:
            # Check centralized model store first
            repo_name = Path(model_path_or_repo).parts[-1]
            central_path = get_model_dir(repo_name) / "tokenizer.json"
            if central_path.exists():
                file_path = central_path
            else:
                print(f"⬇️ Downloading tokenizer from {model_path_or_repo}...")
                file_path = hf_hub_download(
                    repo_id=model_path_or_repo,
                    filename="tokenizer.json",
                    local_dir=str(get_model_dir(repo_name)),
                )

        # 2. Load the Tokenizer
        self._tok = Tokenizer.from_file(str(file_path))

        # 3. Build Special Token Map
        self._special_to_id = {}
        self._id_to_special = {}
        
        for t in self._SPECIALS:
            tid = self._tok.token_to_id(t)
            if tid is not None:
                self._special_to_id[t] = tid
                self._id_to_special[tid] = t
        
        # 4. Compile Regex
        pattern = "|".join(map(re.escape, self._SPECIALS))
        self._split_re = re.compile(f"({pattern})")

        # 5. Set Important IDs
        self.pad_token_id = self._special_to_id.get("<pad>", self._special_to_id.get("<eos>"))
        self.eos_token_id = self._special_to_id.get("<eos>", None)

    def apply_chat_template(
        self, 
        messages: List[Dict[str, str]], 
        add_generation_prompt: bool = False,
        add_thinking: bool = False
    ) -> str:
        formatted_text = ""
        
        # Optional: Add <bos> ONLY ONCE at the very start if needed
        # formatted_text += "<bos>" 

        for msg in messages:
            role = msg["role"]
            content = msg["content"]
            formatted_text += f"<start_of_turn>{role}\n{content}<end_of_turn>\n"

        if add_generation_prompt:
            formatted_text += "<start_of_turn>model\n"
            if add_thinking:
                formatted_text += "<think>\n"
        
        return formatted_text

    def encode(
        self, 
        text_or_messages: Union[str, List[Dict]], 
        add_generation_prompt: bool = False,
        add_thinking: bool = False
    ):
        # 1. Handle Chat Logic
        if isinstance(text_or_messages, list):
            text = self.apply_chat_template(
                text_or_messages, 
                add_generation_prompt=add_generation_prompt,
                add_thinking=add_thinking
            )
        else:
            text = text_or_messages

        # 2. Hybrid Tokenization
        ids = []
        parts = self._split_re.split(text)
        
        for part in parts:
            if not part: continue 

            if part in self._special_to_id:
                ids.append(self._special_to_id[part])
            else:
                # --- FIX IS HERE ---
                # We disable add_special_tokens so it doesn't add <bos> 
                # to every single text fragment (like "user", "model", etc.)
                ids.extend(self._tok.encode(part, add_special_tokens=False).ids)
        
        return ids

    def decode(self, ids, skip_special_tokens=False):
        return self._tok.decode(ids, skip_special_tokens=skip_special_tokens)

    @property
    def vocab_size(self):
        return self._tok.get_vocab_size()