"""
Tokenizer Playground — inspect how text is tokenized, view config and chat templates.
Runs locally and on Hugging Face Spaces.
"""
import html
import os
import random
import gradio as gr
from transformers import AutoTokenizer

# For gated repos: set TOKEN or HF_TOKEN (e.g. in .env or Hugging Face Space secrets)
HF_TOKEN = os.getenv("TOKEN") or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or None


def _token_color(index: int, total: int) -> str:
    """Distinct hue per token, pastel background (HSL)."""
    if total <= 0:
        return "#e0e0e0"
    hue = (index * 263) % 360  # spread hues
    return f"hsl({hue}, 75%, 88%)"


# Style for special tokens (BOS, EOS, PAD, etc.) so they stand out in the viz
SPECIAL_TOKEN_STYLE = (
    "background-color:#e8e0d0;border:1px dashed #9a7b4f;border-radius:2px;"
    "padding:0 1px;font-style:italic"
)


def _build_token_viz_html(
    text: str,
    tokenizer,
    ids: list,
    special_mask: list,
    offset_mapping: list | None,
) -> str:
    """Build HTML with input text and each token span colored by background.
    Special tokens get a distinct style (dashed border, muted bg).
    """
    tokens = tokenizer.convert_ids_to_tokens(ids)
    n = len(tokens)
    parts = []
    for i, (tok, (start, end)) in enumerate(zip(tokens, offset_mapping or [])):
        if start == 0 and end == 0:
            segment = tok.replace("▁", " ").replace("Ġ", " ")
            segment = segment if segment else repr(tok)
        else:
            segment = text[start:end]
        segment = html.escape(segment)
        is_special = i < len(special_mask) and special_mask[i] == 1
        if is_special:
            style = SPECIAL_TOKEN_STYLE
        else:
            color = _token_color(i, n)
            style = f"background-color:{color};border-radius:2px;padding:0 1px"
        parts.append(
            f'<span style="{style}" title="id={ids[i]}">{segment}</span>'
        )
    if not offset_mapping and tokens:
        parts = []
        for i, tok in enumerate(tokens):
            segment = tok.replace("▁", " ").replace("Ġ", " ")
            segment = html.escape(segment if segment else repr(tok))
            is_special = i < len(special_mask) and special_mask[i] == 1
            if is_special:
                style = SPECIAL_TOKEN_STYLE
            else:
                color = _token_color(i, n)
                style = f"background-color:{color};border-radius:2px;padding:0 1px"
            parts.append(f'<span style="{style}" title="id={ids[i]}">{segment}</span>')
    if not parts:
        return "<p><em>No tokens</em></p>"
    return (
        '<div style="font-family:ui-monospace,monospace;font-size:1rem;'
        'line-height:1.6;word-break:break-word;white-space:pre-wrap;">'
        + "".join(parts)
        + "</div>"
    )

# Default tokenizer and preset list
DEFAULT_MODEL = "speakleash/Bielik-11B-v3.0-Instruct"
TOKENIZER_CHOICES = [
    "speakleash/Bielik-11B-v3.0-Instruct",
    "speakleash/Bielik-4.5B-v3.0-Instruct",
    "speakleash/Bielik-1.5B-v3.0-Instruct",
    "speakleash/Bielik-Guard-0.1B-v1.1",
    "mistralai/Mistral-7B-Instruct-v0.2",
    "Qwen/Qwen2.5-14B-Instruct",
    "google/gemma-3-12b-it",
    "openai/gpt-oss-20b",
    "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
    "utter-project/EuroLLM-9B-Instruct",
    "utter-project/EuroLLM-22B-Instruct-2512",
    "swiss-ai/Apertus-8B-Instruct-2509",
]

# Cache tokenizer in module for reuse (per process)
_tokenizer_cache = {}


def get_tokenizer(model_id: str, use_fast: bool = True):
    """Load tokenizer from Hugging Face, with simple in-process cache."""
    key = (model_id or DEFAULT_MODEL, use_fast)
    if key not in _tokenizer_cache:
        try:
            kwargs = dict(
                trust_remote_code=True,
                use_fast=use_fast,
            )
            if HF_TOKEN:
                kwargs["token"] = HF_TOKEN
            _tokenizer_cache[key] = AutoTokenizer.from_pretrained(
                model_id or DEFAULT_MODEL,
                **kwargs,
            )
        except Exception as e:
            raise RuntimeError(f"Failed to load tokenizer for '{model_id}': {e}") from e
    return _tokenizer_cache[key]


def tokenize_text(model_id: str, text: str, use_fast: bool):
    """Tokenize input text and return tokens, ids, and stats."""
    if not (model_id or "").strip():
        model_id = DEFAULT_MODEL
    if not text:
        return (
            "",
            "",
            "",
            "Enter text and click **Tokenize**.",
            None,
            "",
            "",
        )

    try:
        tokenizer = get_tokenizer(model_id.strip(), use_fast)
    except Exception as e:
        return "", "", "", f"**Error loading tokenizer:** {str(e)}", None, "", ""

    # Transformers 5: use tokenizer as callable; get offsets for viz
    encoded = tokenizer(
        text,
        return_attention_mask=False,
        return_special_tokens_mask=True,
        return_offsets_mapping=True,
        add_special_tokens=True,
    )
    ids = encoded["input_ids"]
    if hasattr(ids, "tolist"):
        ids = ids.tolist()
    else:
        ids = list(ids)
    special_mask = encoded.get("special_tokens_mask", [0] * len(ids))
    if hasattr(special_mask, "tolist"):
        special_mask = special_mask.tolist()
    else:
        special_mask = list(special_mask)
    offset_mapping = encoded.get("offset_mapping")
    if offset_mapping is not None:
        if hasattr(offset_mapping, "tolist"):
            offset_mapping = offset_mapping.tolist()
        # Unwrap batch dim when tokenizer returned shape (1, n, 2)
        if offset_mapping and len(offset_mapping) == 1 and isinstance(offset_mapping[0], (list, tuple)):
            offset_mapping = offset_mapping[0]
        offset_mapping = [tuple(p) for p in offset_mapping]

    # Build token display: "token (id)" per token
    tokens = tokenizer.convert_ids_to_tokens(ids)
    parts = []
    for t, i, sp in zip(tokens, ids, special_mask):
        t_display = repr(t) if sp else t
        parts.append(f"{t_display} ({i})")
    token_list_str = " | ".join(parts)
    token_list_md = f"```\n{token_list_str}\n```" if token_list_str else ""

    ids_str = ", ".join(str(i) for i in ids)

    # Round-trip check: decode(encode(text)) vs original
    try:
        decoded = tokenizer.decode(ids, skip_special_tokens=True)
        round_trip_ok = decoded.strip() == text.strip()
        round_trip_row = f"| decode(encode(text)) | {'✓ Match' if round_trip_ok else '✗ Diff'} |\n"
        round_trip_extra = ""
        if not round_trip_ok:
            dec_snippet = decoded[:200] + ("…" if len(decoded) > 200 else "")
            # Avoid breaking markdown code block
            dec_snippet = dec_snippet.replace("`", "'").replace("\n", " ")
            round_trip_extra = f"\n*Decoded (first 200 chars):* `{dec_snippet}`\n"
    except Exception:
        round_trip_ok = False
        round_trip_row = "| decode(encode(text)) | — |\n"
        round_trip_extra = ""

    # Stats
    num_chars = len(text)
    num_tokens = len(ids)
    num_words = len(text.split())
    chars_per_token = num_chars / num_tokens if num_tokens else 0
    tokens_per_word = num_tokens / num_words if num_words else 0
    stats_md = (
        f"| Metric | Value |\n|--------|-------|\n"
        f"| Characters | {num_chars} |\n"
        f"| Words | {num_words} |\n"
        f"| Tokens | {num_tokens} |\n"
        f"| Chars per token | {chars_per_token:.2f} |\n"
        f"| Tokens per word | {tokens_per_word:.2f} |\n"
        f"{round_trip_row}"
        f"{round_trip_extra}"
    )

    # Colored token visualization (same as input text, each token different bg)
    viz_html = _build_token_viz_html(text, tokenizer, ids, special_mask, offset_mapping)

    # Config summary for this tokenizer
    config_md = _format_tokenizer_config(tokenizer)
    chat_md = _format_chat_template(tokenizer)

    return (
        viz_html,
        token_list_md,
        ids_str,
        stats_md,
        tokenizer,
        config_md,
        chat_md,
    )


def _format_tokenizer_config(tokenizer) -> str:
    """Format main tokenizer config fields as markdown."""
    lines = ["| Setting | Value |", "|---------|-------|"]
    attrs = [
        ("model_max_length", getattr(tokenizer, "model_max_length", None)),
        ("vocab_size", getattr(tokenizer, "vocab_size", None)),
        ("pad_token", _repr_token(getattr(tokenizer, "pad_token", None))),
        ("pad_token_id", getattr(tokenizer, "pad_token_id", None)),
        ("eos_token", _repr_token(getattr(tokenizer, "eos_token", None))),
        ("eos_token_id", getattr(tokenizer, "eos_token_id", None)),
        ("bos_token", _repr_token(getattr(tokenizer, "bos_token", None))),
        ("bos_token_id", getattr(tokenizer, "bos_token_id", None)),
        ("unk_token", _repr_token(getattr(tokenizer, "unk_token", None))),
        ("unk_token_id", getattr(tokenizer, "unk_token_id", None)),
        ("sep_token", _repr_token(getattr(tokenizer, "sep_token", None))),
        ("sep_token_id", getattr(tokenizer, "sep_token_id", None)),
        ("cls_token", _repr_token(getattr(tokenizer, "cls_token", None))),
        ("cls_token_id", getattr(tokenizer, "cls_token_id", None)),
    ]
    for name, val in attrs:
        if val is None:
            val = "—"
        elif val == 1000000000000000019884624838656:  # default in some tokenizers
            val = "— (default)"
        lines.append(f"| {name} | {val} |")
    return "\n".join(lines)


def _repr_token(t):
    if t is None:
        return "—"
    return repr(t)


def _format_chat_template(tokenizer) -> str:
    """Show chat template if available and a short example."""
    out = []
    chat_template = getattr(tokenizer, "chat_template", None)
    if chat_template is None:
        # Try tokenizer_config
        config = getattr(tokenizer, "tokenizer_config", None) or {}
        chat_template = config.get("chat_template")
    if chat_template:
        out.append("**Chat template (raw):**")
        out.append("```")
        out.append(chat_template if isinstance(chat_template, str) else str(chat_template)[:2000])
        out.append("```")
        # Try apply_chat_template with example messages (system + user + assistant)
        try:
            example = [
                {"role": "system", "content": "You are a helpful assistant. Answer concisely."},
                {"role": "user", "content": "Hello!"},
                {"role": "assistant", "content": "Hi there."},
            ]
            applied = tokenizer.apply_chat_template(
                example,
                tokenize=False,
                add_generation_prompt=True,
            )
            out.append("\n**Example (system + user + assistant, add_generation_prompt=True):**")
            out.append("```")
            out.append(applied[:1500] if len(applied) > 1500 else applied)
            out.append("```")
            # Colorized token view of the example
            try:
                enc = tokenizer(
                    applied,
                    return_attention_mask=False,
                    return_special_tokens_mask=True,
                    return_offsets_mapping=True,
                    add_special_tokens=False,
                )
                ids = enc["input_ids"]
                ids = ids.tolist() if hasattr(ids, "tolist") else list(ids)
                sp_mask = enc.get("special_tokens_mask", [0] * len(ids))
                sp_mask = sp_mask.tolist() if hasattr(sp_mask, "tolist") else list(sp_mask)
                off = enc.get("offset_mapping")
                if off is not None:
                    if hasattr(off, "tolist"):
                        off = off.tolist()
                    if off and len(off) == 1 and isinstance(off[0], (list, tuple)):
                        off = off[0]
                    off = [tuple(p) for p in off]
                viz_html = _build_token_viz_html(applied, tokenizer, ids, sp_mask, off)
                out.append("\n**Colorized:**")
                out.append(viz_html)
            except Exception as viz_e:
                out.append(f"\n*Colorized view failed: {viz_e}*")
        except Exception as e:
            out.append(f"\n*Example apply failed: {e}*")
    else:
        out.append("No chat template defined for this tokenizer.")
    return "\n".join(out)


def vocab_peek(model_id: str, use_fast: bool, count: int) -> str:
    """Return a list of random tokens from the tokenizer vocab for exploration."""
    if not (model_id or "").strip():
        model_id = DEFAULT_MODEL
    try:
        tokenizer = get_tokenizer(model_id.strip(), use_fast)
    except Exception as e:
        return f"**Error loading tokenizer:** {e}"
    vocab = getattr(tokenizer, "get_vocab", None)
    if not vocab:
        return "This tokenizer does not expose a vocab list."
    vocab_dict = vocab()
    if not vocab_dict:
        return "Vocabulary is empty."
    ids = list(vocab_dict.values())
    k = min(max(1, int(count)), len(ids))
    chosen = random.sample(ids, k)
    tokens = tokenizer.convert_ids_to_tokens(chosen)
    lines = [f"- {repr(t)} → **{tid}**" for t, tid in zip(tokens, chosen)]
    return "**Random tokens from vocab:**\n\n" + "\n".join(lines)


# ---- UI ----
with gr.Blocks(title="Tokenizer Playground") as demo:
    gr.Markdown("# Tokenizer Playground")
    gr.Markdown(
        "Choose a tokenizer and enter text to see how it is tokenized. "
        "View tokenizer config (e.g. EOS, PAD) and chat template below."
    )

    with gr.Row():
        model_id = gr.Dropdown(
            label="Tokenizer",
            choices=TOKENIZER_CHOICES,
            value=DEFAULT_MODEL,
            scale=3,
        )
        use_fast = gr.Checkbox(label="Use fast tokenizer", value=True, scale=1)

    input_text = gr.Textbox(
        label="Input text",
        placeholder="Type or paste text here to tokenize...",
        lines=8,
        value=(
            "Bielik.AI – europejska rodzina otwartych modeli językowych, stworzona w Polsce.\n"
            "Otwarte, bezpłatne i bezpieczne. Rozmawiaj z Bielikiem,\n"
            "eksperymentuj i buduj własne rozwiązania"
        ),
    )
    with gr.Row():
        tokenize_btn = gr.Button("Tokenize", variant="primary")
        vocab_count = gr.Number(label="Vocab peek (n random tokens)", value=20, minimum=1, maximum=100, step=1)

    gr.Markdown("### Tokenized text")
    token_viz = gr.HTML()

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Tokens (token (id))")
            tokens_display = gr.Markdown()
        with gr.Column():
            gr.Markdown("### Token IDs")
            ids_display = gr.Textbox(lines=6, max_lines=12, interactive=False)

    gr.Markdown("### Stats")
    stats_display = gr.Markdown()

    gr.Markdown("---")
    gr.Markdown("### Tokenizer config")
    config_display = gr.Markdown()

    gr.Markdown("### Chat template")
    chat_display = gr.Markdown()

    gr.Markdown("---")
    gr.Markdown("### Vocab peek")
    vocab_display = gr.Markdown()

    tokenizer_state = gr.State(value=None)

    def on_tokenize(mid, text, fast, vcount):
        viz, t_str, i_str, stats, tok, cfg, chat = tokenize_text(mid, text, fast)
        vocab_md = vocab_peek(mid, fast, vcount)
        return viz, t_str, i_str, stats, tok, cfg, chat, vocab_md

    tokenize_btn.click(
        fn=on_tokenize,
        inputs=[model_id, input_text, use_fast, vocab_count],
        outputs=[
            token_viz,
            tokens_display,
            ids_display,
            stats_display,
            tokenizer_state,
            config_display,
            chat_display,
            vocab_display,
        ],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())