import sys
import io

import gradio
from transformers import AutoTokenizer

# List of supported models
MODELS = [
    "zai-org/GLM-5",
    "MiniMaxAI/MiniMax-M2.5",
    "deepseek-ai/DeepSeek-V3.2",
    "Qwen/Qwen3-14B",
    "Qwen/Qwen3-235B-A22B",
    "Qwen/Qwen3-30B-A3B",
    "Qwen/Qwen3-32B"
]

# Cache tokenizers to avoid repeated downloads
tokenizer_cache = {}

def highlight_text(tokens):
    return [
        (token, str(i%9))
        for i, token in enumerate(tokens)
    ]

def count_tokens(model_name, text_input):
    # textbytes = text_input.encode('utf-8')
    text = text_input

    if not text.strip():
        return 0, []

    # Load tokenizer (with caching)
    if model_name not in tokenizer_cache:
        tokenizer_cache[model_name] = AutoTokenizer.from_pretrained(
            model_name, trust_remote_code=True
        )
    tokenizer = tokenizer_cache[model_name]

    token_ids = tokenizer.encode(text, add_special_tokens=False)
    tokens = []
    for tid in token_ids:
        tokens.append(tokenizer.decode([tid]))

    return len(token_ids), len(text), highlight_text(tokens)


gradio.Interface(
    fn=count_tokens,
    inputs=[
        gradio.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
        gradio.Textbox(lines=5, label="Input Text")
    ],
    outputs=[
        gradio.Number(label="Tokens"),
        gradio.Number(label="Characters"),
        # gradio.Textbox(label="Tokens")
        gradio.HighlightedText(
            label="Tokens",
            show_inline_category=False,
            color_map={
                "1": "red", "2": "green", "3": "blue",
                "4": "yellow", "5": "orange", "6": "purple",
                "7": "pink", "8": "cyan", "9": "gray"
            }
        )
    ],
    title="Model Tokenizer",
    description="Select a model and input text to see token count and token list."
).launch()