import gradio as gr
from transformers import WhisperTokenizer

# Supported languages with their Whisper language token names
LANGUAGES = {
    "Hindi": "hi",
    "English": "en",
    "Urdu": "ur",
    "Bengali": "bn",
    "Tamil": "ta",
    "Telugu": "te",
    "Marathi": "mr",
    "Gujarati": "gu",
    "Punjabi": "pa",
    "Kannada": "kn",
    "Malayalam": "ml",
    "Odia": "or",
    "Arabic": "ar",
    "French": "fr",
    "Spanish": "es",
    "German": "de",
    "Chinese": "zh",
    "Japanese": "ja",
    "Korean": "ko",
    "Russian": "ru",
    "Portuguese": "pt",
    "Italian": "it",
    "Dutch": "nl",
    "Turkish": "tr",
    "Polish": "pl",
    "Indonesian": "id",
}

# Cache loaded tokenizers so we don't reload on every call
_tokenizer_cache = {}

def get_tokenizer(multilingual: bool):
    key = "multilingual" if multilingual else "english"
    if key not in _tokenizer_cache:
        model_name = "openai/whisper-large-v3" if multilingual else "openai/whisper-small.en"
        _tokenizer_cache[key] = WhisperTokenizer.from_pretrained(model_name)
    return _tokenizer_cache[key]

def parse_tokens(raw_input: str):
    """
    Accepts tokens in any format:
      50258, 50276, 50359, 50363, 13, 2958
      50258\n50276\n50359
      50258 50276 50359
    """
    cleaned = raw_input.replace("\n", ",").replace(" ", ",")
    parts = [p.strip() for p in cleaned.split(",") if p.strip()]
    if not parts:
        raise ValueError("No tokens found in input.")
    tokens = []
    for p in parts:
        if not p.isdigit():
            raise ValueError(f"'{p}' is not a valid integer token ID.")
        tokens.append(int(p))
    return tokens

def decode_tokens(raw_tokens: str, language: str, skip_special: bool):
    if not raw_tokens or not raw_tokens.strip():
        return "Please enter token IDs.", "", ""

    try:
        token_ids = parse_tokens(raw_tokens)
    except ValueError as e:
        return f"Parse error: {e}", "", ""

    lang_code = LANGUAGES.get(language, "hi")
    is_multilingual = lang_code != "en"

    try:
        tokenizer = get_tokenizer(is_multilingual)
    except Exception as e:
        return f"Failed to load tokenizer: {e}", "", ""

    try:
        decoded_full = tokenizer.decode(token_ids, skip_special_tokens=False)
        decoded_clean = tokenizer.decode(token_ids, skip_special_tokens=True)
    except Exception as e:
        return f"Decode error: {e}", "", ""

    # Per-token breakdown
    breakdown_lines = []
    for tid in token_ids:
        try:
            word = tokenizer.decode([tid], skip_special_tokens=False)
            word_display = repr(word) if word.strip() == "" else word
            breakdown_lines.append(f"  {tid:>6}  ->  {word_display}")
        except Exception:
            breakdown_lines.append(f"  {tid:>6}  ->  [decode error]")
    breakdown = "\n".join(breakdown_lines)

    result = decoded_clean if skip_special else decoded_full
    return result, decoded_full, breakdown


with gr.Blocks(title="Whisper Token Decoder", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Whisper Token Decoder
    Paste raw Whisper token IDs from your Android app and decode them to text.
    Tokens can be comma-separated, space-separated, or one per line.
    """)

    with gr.Row():
        with gr.Column(scale=2):
            token_input = gr.Textbox(
                label="Token IDs",
                placeholder="e.g.\n50258,\n50276,\n50359,\n50363,\n13,\n2958",
                lines=10,
                max_lines=30,
            )
            with gr.Row():
                language_dropdown = gr.Dropdown(
                    choices=list(LANGUAGES.keys()),
                    value="Hindi",
                    label="Language",
                )
                skip_special_cb = gr.Checkbox(
                    value=True,
                    label="Skip special tokens",
                )
            decode_btn = gr.Button("Decode Tokens", variant="primary", size="lg")

        with gr.Column(scale=2):
            output_text = gr.Textbox(
                label="Decoded Text",
                lines=4,
                interactive=False,
            )
            output_full = gr.Textbox(
                label="Full decode (with special tokens)",
                lines=3,
                interactive=False,
            )
            output_breakdown = gr.Textbox(
                label="Per-token breakdown",
                lines=12,
                interactive=False,
            )

    decode_btn.click(
        fn=decode_tokens,
        inputs=[token_input, language_dropdown, skip_special_cb],
        outputs=[output_text, output_full, output_breakdown],
    )

    gr.Examples(
        examples=[
            ["50258,\n50276,\n50359,\n50363,\n13,\n2958", "Hindi", True],
            ["50258, 50359, 50363, 2264, 526, 345", "English", True],
        ],
        inputs=[token_input, language_dropdown, skip_special_cb],
        label="Try these examples",
    )

    gr.Markdown("""
    ---
    **Special token reference:**
    `50258` = start-of-transcript | `50359` = Hindi language tag | `50363` = no-timestamps | `50256` = end-of-text
    """)

demo.launch()