"""HTML rendering for the demo: coloured tokens by source, legend, stats, replay payload."""

from __future__ import annotations

import html
import json

# token source -> (legend label, tooltip)
SOURCES = {
    "copy": ("Copied from context", "copied from the prompt or earlier output by n-gram lookup, then verified"),
    "head": ("Exit head (layer 4)", "drafted by the shallow exit head, then verified"),
    "star": ("Head top-4 candidate", "an alternative from the head's top-4, then verified"),
    "draft": ("Draft model", "drafted by a separate model, then verified"),
    "model": ("Full model", "produced by the full model (correction or bonus token of a verification)"),
}


class Detokenizer:
    """Incremental decoding into (text, source) pieces. Tokens that end inside a multi-byte
    character are held back until the character is complete; a piece takes the source of its
    first token."""

    def __init__(self, tokenizer):
        self.tok = tokenizer
        self.ids: list[int] = []
        self.text = ""
        self.pending: list[str] = []

    def _emit(self, final: bool) -> list[tuple[str, str]]:
        cur = self.tok.decode(self.ids)
        if cur.endswith("�") and not final:
            return []
        piece = cur[len(self.text):]
        src = self.pending[0] if self.pending else "model"
        self.text, self.pending = cur, []
        return [(piece, src)] if piece else []

    def feed(self, ids: list[int], sources: list[str]) -> list[tuple[str, str]]:
        out = []
        for t, s in zip(ids, sources):
            self.ids.append(int(t))
            self.pending.append(s)
            out += self._emit(final=False)
        return out

    def flush(self) -> list[tuple[str, str]]:
        return self._emit(final=True) if self.pending else []


def tokens_html(pieces, colored: bool = True, cursor: bool = False) -> str:
    out = []
    for text, src in pieces:
        t = html.escape(text)
        if colored and src != "model" and src in SOURCES:
            out.append(f'<span class="tok tok-{src}" title="{html.escape(SOURCES[src][1])}">{t}</span>')
        else:
            out.append(t)
    if cursor:
        out.append('<span class="cursor"></span>')
    return f'<div class="dssd-text">{"".join(out)}</div>'


def legend_html() -> str:
    chips = "".join(
        f'<span class="legend-item"><span class="tok tok-{s} swatch">Aa</span>'
        f"{html.escape(SOURCES[s][0])}</span>"
        for s in ("copy", "head", "star", "model"))
    return f'<div class="legend">{chips}</div>'


def _card(value: str, label: str, cls: str = "") -> str:
    return f'<div class="stat {cls}"><div class="stat-value">{value}</div><div class="stat-label">{label}</div></div>'


def stats_html(g) -> str:
    """Stat cards and the source-share bar for a DSSD Generation (with g.ar when compared)."""
    n = max(len(g.sources), 1)
    counts = {s: g.sources.count(s) for s in ("copy", "head", "star", "draft", "model")}
    cards = []
    if g.ar is not None and g.speedup is not None:
        cards.append(_card(f"{g.speedup:.2f}×", "decoding speedup", "stat-main"))
    cards.append(_card(f"{g.decode_tokens_per_s:.1f}", "DSSD tokens/s"))
    if g.ar is not None:
        cards.append(_card(f"{g.ar.decode_tokens_per_s:.1f}", "autoregressive tokens/s"))
    cards.append(_card(f"{g.stats['tau']:.2f}", "tokens per verification"))
    if g.ar is not None:
        ok = bool(g.identical)
        cards.append(_card("✓ identical" if ok else "✗ outputs differ", "DSSD vs autoregressive",
                           "stat-ok" if ok else "stat-bad"))
    bar = "".join(
        f'<div class="share tok-{s}" style="width:{100 * c / n:.2f}%" title="{SOURCES[s][0]}: {c} tokens"></div>'
        for s, c in counts.items() if c)
    labels = " · ".join(f"{SOURCES[s][0]} {round(100 * c / n)}%" for s, c in counts.items() if c)
    return (f'<div class="stats">{"".join(cards)}</div>'
            f'<div class="share-bar">{bar}</div><div class="share-labels">{html.escape(labels)}</div>')


def replay_html(timelines: dict) -> str:
    """Hidden element holding {"dssd": [[t, text, source], ...], "ar": [...]} for the client-side
    side-by-side replay at the measured speed."""
    data = html.escape(json.dumps(timelines, ensure_ascii=False), quote=True)
    return f"<div class=\"replay-data\" data-replay='{data}'></div>"