from google import genai
import gradio as gr
import os

models = [
    'gemini-2.5-flash',
    'gemini-2.5-pro',
    'gemini-2.0-flash',
    'gemini-2.0-flash-lite',
    'gemini-pro-latest',
    'gemini-2.5-flash-lite',
    'gemini-2.5-flash-image',
    'gemini-3-flash-preview',
    'gemini-3.1-pro-image-preview',
    'gemini-3.1-pro-preview',
]

color_map = {
    '3': '#ccbfee',
    '2': '#beedc6',
    '1': '#f4d7ab',
    '0': '#f4aeb1',
    # '4': '#a4dcf3',
}

client_vertex = genai.Client(
    vertexai=True,
    api_key=os.getenv("VERTEX_API_KEY")
)

def tokens_len(prompt, model):
    result = client_vertex.models.compute_tokens(
        model=model,
        contents=prompt
    )
    tokens = result.tokens_info[0].tokens

    return (
        f'{len(tokens)} tokens, {len(prompt)} chracters',
        [
            (
                token.decode('utf-8').replace("\n", "\\n"), 
                str(i % len(color_map))
            ) for i, token in enumerate(tokens)
        ]
    )

web_tokenizer = gr.Interface(
    fn = tokens_len,
    inputs = [
        gr.Textbox(
            label='輸入提示文字',
            lines=6),
        gr.Dropdown(
            label='模型名稱',
            choices = models,
            value=models[-1])
    ],
    outputs = [
        gr.Textbox(
            label='統計量'),
        gr.Highlightedtext(
            label='切割結果：',
            color_map = color_map,
            show_inline_category=False)
    ]
)

web_tokenizer.launch(share=True, debug=True)