from __future__ import annotations

from functools import lru_cache
from html import escape
from typing import Any

import gradio as gr
from transformers import AutoTokenizer


DEFAULT_MODEL_ID = "Qwen/Qwen3.5-2B"
DEFAULT_ADDED_TOKENS = [
    "<|graph|>",
    "<|nodes|>",
    "<|node|>",
    "<|id|>",
    "<|label|>",
    "<|bbox|>",
    "<|conf|>",
    "<|attrs|>",
    "<|edges|>",
    "<|edge|>",
    "<|src|>",
    "<|pred|>",
    "<|tgt|>",
    "<|end_graph|>",
]

DEFAULT_ADDED_TOKEN_TEXT = "\n".join(DEFAULT_ADDED_TOKENS)
SAMPLE_TEXT = "<|graph|><|nodes|><|node|><|id|>person1<|label|>man<|bbox|>0.0 0.0 0.5 1.0<|conf|>0.9<|attrs|>appearance=white t-shirt<|attrs|>size=medium<|node|><|id|>person2<|label|>woman<|bbox|>0.45 0.0 1.0 1.0<|conf|>0.9<|attrs|>appearance=patterned sleeveless top<|attrs|>size=medium<|node|><|id|>table1<|label|>table<|bbox|>0.0 0.7 1.0 1.0<|conf|>0.8<|attrs|>appearance=brown wood<|attrs|>size=large<|edges|><|edge|><|src|>person1<|pred|>next_to<|tgt|>person2<|edge|><|src|>person1<|pred|>sitting_on<|tgt|>table1<|edge|><|src|>person2<|pred|>sitting_on<|tgt|>table1<|end_graph|>"

MAX_DETAIL_ROWS = 12000


@lru_cache(maxsize=4)
def load_tokeniser(model_id: str) -> Any:
    model_id = model_id.strip()
    if not model_id:
        raise ValueError("Tokeniser model id is required.")
    return AutoTokenizer.from_pretrained(model_id, use_fast=True)


def parse_added_tokens(raw_tokens: str) -> list[str]:
    seen: set[str] = set()
    parsed: list[str] = []

    for line in raw_tokens.splitlines():
        token = line.strip()
        if not token or token in seen:
            continue
        seen.add(token)
        parsed.append(token)

    return parsed


def split_by_added_tokens(text: str, added_tokens: list[str]) -> list[tuple[str, str]]:
    if not text:
        return []
    if not added_tokens:
        return [("base", text)]

    sorted_added = sorted(added_tokens, key=len, reverse=True)
    pieces: list[tuple[str, str]] = []
    buffer: list[str] = []
    index = 0

    while index < len(text):
        match = next(
            (token for token in sorted_added if text.startswith(token, index)),
            None,
        )

        if match is not None:
            if buffer:
                pieces.append(("base", "".join(buffer)))
                buffer = []
            pieces.append(("added", match))
            index += len(match)
            continue

        buffer.append(text[index])
        index += 1

    if buffer:
        pieces.append(("base", "".join(buffer)))

    return pieces


def tokenise_base_piece(tokeniser: Any, text: str) -> list[dict[str, str]]:
    if not text:
        return []

    offsets: list[tuple[int, int]] = []

    try:
        encoded = tokeniser(
            text,
            add_special_tokens=False,
            return_offsets_mapping=True,
        )
        input_ids = encoded["input_ids"]
        offsets = encoded.get("offset_mapping") or []
    except Exception:
        input_ids = tokeniser.encode(text, add_special_tokens=False)

    if input_ids and isinstance(input_ids[0], list):
        input_ids = input_ids[0]
    if offsets and isinstance(offsets[0], list) and offsets and offsets[0] and isinstance(offsets[0][0], (list, tuple)):
        offsets = offsets[0]

    token_strings = tokeniser.convert_ids_to_tokens(input_ids)
    tokens: list[dict[str, str]] = []

    for index, token_id in enumerate(input_ids):
        display = ""
        if index < len(offsets):
            start, end = offsets[index]
            if end > start:
                display = text[start:end]

        if not display:
            try:
                display = tokeniser.decode(
                    [token_id],
                    skip_special_tokens=False,
                    clean_up_tokenization_spaces=False,
                )
            except Exception:
                display = token_strings[index] if index < len(token_strings) else ""

        tokens.append(
            {
                "kind": "base",
                "id": str(token_id),
                "raw": token_strings[index] if index < len(token_strings) else display,
                "display": display,
            }
        )

    return tokens


def tokenise_with_added_tokens(
    text: str,
    added_tokens: list[str],
    tokeniser: Any,
) -> list[dict[str, str]]:
    output: list[dict[str, str]] = []

    for kind, value in split_by_added_tokens(text, added_tokens):
        if kind == "added":
            output.append(
                {
                    "kind": "added",
                    "id": "added",
                    "raw": value,
                    "display": value,
                }
            )
        else:
            output.extend(tokenise_base_piece(tokeniser, value))

    return output


def printable(value: str) -> str:
    return (
        value.replace("\\", "\\\\")
        .replace("\n", "\\n")
        .replace("\r", "\\r")
        .replace("\t", "\\t")
        .replace(" ", "·")
    )


def render_metrics(total: int, base: int, added: int, chars: int) -> str:
    return f"""
    <div class="metrics-grid">
      <div class="metric-card primary"><span>Total tokens</span><strong>{total:,}</strong></div>
      <div class="metric-card"><span>Base tokeniser</span><strong>{base:,}</strong></div>
      <div class="metric-card"><span>Added tokens</span><strong>{added:,}</strong></div>
      <div class="metric-card"><span>Characters</span><strong>{chars:,}</strong></div>
    </div>
    """


def render_added_token_chips(added_tokens: list[str]) -> str:
    if not added_tokens:
        return '<div class="empty-note">No added tokens configured.</div>'

    chips = "".join(
        f'<span class="added-chip">{escape(token)}</span>' for token in added_tokens
    )
    return f'<div class="chip-wrap">{chips}</div>'


def render_highlights(tokens: list[dict[str, str]], max_rendered: int) -> str:
    if not tokens:
        return '<div class="empty-note">No tokens to show.</div>'

    limited = tokens[:max_rendered]
    spans: list[str] = []

    for index, token in enumerate(limited):
        if token["kind"] == "added":
            class_name = "token added"
            token_type = "added"
        else:
            class_name = f"token base base-{index % 3}"
            token_type = "base"

        display = token["display"]
        title = escape(
            f"#{index + 1} | {token_type} | id: {token['id']} | {printable(token['raw'])}"
        )
        body = escape(display) if display else "&nbsp;"
        spans.append(f'<span class="{class_name}" title="{title}">{body}</span>')

    note = (
        f"Showing first {len(limited):,} of {len(tokens):,} tokens"
        if len(tokens) > len(limited)
        else f"{len(tokens):,} tokens shown"
    )

    return f"""
    <div class="token-note">{note}</div>
    <div class="token-output">{''.join(spans)}</div>
    """


def build_rows(tokens: list[dict[str, str]], max_rendered: int) -> list[list[str]]:
    rows: list[list[str]] = []

    for index, token in enumerate(tokens[:max_rendered]):
        rows.append(
            [
                str(index + 1),
                token["kind"],
                token["id"],
                printable(token["display"] or token["raw"]),
            ]
        )

    return rows


def analyze(
    text: str,
    added_token_text: str,
    model_id: str,
    max_rendered: int,
) -> tuple[str, str, str, list[list[str]], str]:
    added_tokens = parse_added_tokens(added_token_text or "")
    max_rendered = max(1, min(int(max_rendered or 2000), MAX_DETAIL_ROWS))

    try:
        tokeniser = load_tokeniser(model_id)
        tokens = tokenise_with_added_tokens(text or "", added_tokens, tokeniser)
    except Exception as exc:
        status = f'<div class="status error">Tokeniser error: {escape(str(exc))}</div>'
        return (
            render_metrics(0, 0, 0, len(text or "")),
            render_added_token_chips(added_tokens),
            status,
            [],
            status,
        )

    added_count = sum(1 for token in tokens if token["kind"] == "added")
    base_count = len(tokens) - added_count
    status = f'<div class="status ready">Loaded {escape(model_id.strip())}</div>'

    return (
        render_metrics(len(tokens), base_count, added_count, len(text or "")),
        render_added_token_chips(added_tokens),
        render_highlights(tokens, max_rendered),
        build_rows(tokens, max_rendered),
        status,
    )


CSS = """
.gradio-container {
  max-width: 1440px !important;
}
.app-title h1 {
  font-size: 1.7rem;
  letter-spacing: 0;
  margin-bottom: 0.35rem;
}
.app-title p {
  color: #53636d;
  margin-top: 0;
}
.metrics-grid {
  display: grid;
  grid-template-columns: repeat(4, minmax(130px, 1fr));
  gap: 10px;
}
.metric-card {
  border: 1px solid #d6dee2;
  border-radius: 8px;
  background: #fff;
  padding: 12px 14px;
}
.metric-card.primary {
  border-color: #8ccbd2;
}
.metric-card span {
  display: block;
  color: #53636d;
  font-size: 0.84rem;
  font-weight: 700;
}
.metric-card strong {
  display: block;
  color: #172026;
  font-size: 2rem;
  line-height: 1;
  margin-top: 6px;
}
.metric-card.primary strong {
  color: #006873;
}
.chip-wrap {
  display: flex;
  flex-wrap: wrap;
  gap: 7px;
}
.added-chip,
.token {
  border-radius: 5px;
  font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, "Liberation Mono", monospace;
}
.added-chip {
  border: 1px solid #d78c00;
  background: #ffe8b3;
  color: #4f3200;
  padding: 4px 7px;
  font-size: 0.86rem;
}
.token-note {
  color: #53636d;
  font-size: 0.86rem;
  margin-bottom: 8px;
  text-align: right;
}
.token-output {
  min-height: 300px;
  max-height: 620px;
  overflow: auto;
  border: 1px solid #d6dee2;
  border-radius: 8px;
  background: #fff;
  padding: 14px;
  line-height: 1.9;
  white-space: pre-wrap;
  overflow-wrap: anywhere;
}
.token {
  display: inline;
  border: 1px solid #cbd7db;
  margin: 0 1px 2px 0;
  padding: 2px 3px;
  cursor: default;
  white-space: pre-wrap;
}
.token.added {
  border-color: #d78c00;
  background: #ffe8b3;
  color: #4f3200;
  font-weight: 750;
}
.token.base-0 {
  background: #e4f4ef;
  color: #14231e;
}
.token.base-1 {
  background: #e9eefb;
  color: #182139;
}
.token.base-2 {
  background: #f2e8f7;
  color: #2e2135;
}
.status {
  border-radius: 999px;
  display: inline-block;
  padding: 6px 11px;
  font-size: 0.88rem;
  border: 1px solid #d6dee2;
}
.status.ready {
  border-color: #9bd2aa;
  color: #137333;
}
.status.error {
  border-color: #f1a6a0;
  color: #b42318;
  border-radius: 8px;
  max-width: 100%;
  overflow-wrap: anywhere;
}
.empty-note {
  color: #53636d;
}
@media (max-width: 820px) {
  .metrics-grid {
    grid-template-columns: 1fr 1fr;
  }
}
@media (max-width: 520px) {
  .metrics-grid {
    grid-template-columns: 1fr;
  }
}
"""


with gr.Blocks(title="Custom Added Token Counter") as demo:
    gr.Markdown(
        """
        # Custom added token counter
        Count raw tokeniser pieces with editable added tokens.
        """,
        elem_classes=["app-title"],
    )

    with gr.Row():
        model_id = gr.Textbox(
            value=DEFAULT_MODEL_ID,
            label="Tokeniser model",
            scale=4,
            container=True,
        )
        max_rendered = gr.Slider(
            minimum=100,
            maximum=MAX_DETAIL_ROWS,
            value=2000,
            step=100,
            label="Rendered token limit",
            scale=2,
        )

    with gr.Row():
        with gr.Column(scale=3):
            input_text = gr.Textbox(
                value=SAMPLE_TEXT,
                label="Input text",
                lines=14,
                max_lines=28,
                autoscroll=False,
            )
        with gr.Column(scale=2):
            added_tokens = gr.Textbox(
                value=DEFAULT_ADDED_TOKEN_TEXT,
                label="Added tokens",
                lines=14,
                max_lines=28,
                autoscroll=False,
            )

    with gr.Row():
        count_button = gr.Button("Count tokens", variant="primary")
        clear_button = gr.Button("Clear text")

    status_html = gr.HTML()
    metrics_html = gr.HTML()
    added_token_html = gr.HTML(label="Active added tokens")
    highlighted_html = gr.HTML(label="Highlighted tokens")
    token_table = gr.Dataframe(
        headers=["#", "type", "id", "token"],
        datatype=["str", "str", "str", "str"],
        interactive=False,
        wrap=True,
        label="Token details",
    )

    analyze_inputs = [input_text, added_tokens, model_id, max_rendered]
    analyze_outputs = [
        metrics_html,
        added_token_html,
        highlighted_html,
        token_table,
        status_html,
    ]

    demo.load(analyze, inputs=analyze_inputs, outputs=analyze_outputs)
    count_button.click(analyze, inputs=analyze_inputs, outputs=analyze_outputs)
    input_text.submit(analyze, inputs=analyze_inputs, outputs=analyze_outputs)
    model_id.submit(analyze, inputs=analyze_inputs, outputs=analyze_outputs)
    added_tokens.submit(analyze, inputs=analyze_inputs, outputs=analyze_outputs)
    clear_button.click(lambda: "", outputs=input_text).then(
        analyze,
        inputs=analyze_inputs,
        outputs=analyze_outputs,
    )


if __name__ == "__main__":
    demo.launch(css=CSS, ssr_mode=False)