from __future__ import annotations

import functools
import html
import re
from typing import Optional

import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer

# ── constants ─────────────────────────────────────────────────────────────────

MODEL_ID = "fromziro/JetonCount"
DEFAULT_VOCAB_SIZE = 32_000

PUNCTUATION_CHARS = set(r""".,!?;:'"`~@#$%^&*()-_=+[]{}<>/\|""")
SYMBOL_CHARS      = set(r"""@#$%^&*()-_=+[]{}<>/\|~`""")

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if DEVICE.type == "cuda":
    torch.backends.cudnn.benchmark = True

# ── model / tokenizer loading ─────────────────────────────────────────────────

@functools.lru_cache(maxsize=1)
def load_model():
    model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
    model.eval()
    model.to(DEVICE)
    return model


@functools.lru_cache(maxsize=16)
def load_tokenizer(tokenizer_id: str):
    return AutoTokenizer.from_pretrained(tokenizer_id, use_fast=True)


def get_vocab_size(tokenizer) -> int:
    vs = getattr(tokenizer, "vocab_size", None)
    if isinstance(vs, int) and vs > 0:
        return vs
    try:
        return int(len(tokenizer))
    except Exception:
        return DEFAULT_VOCAB_SIZE

# ── feature computation ───────────────────────────────────────────────────────

def compute_stats(text: str, vocab_size: int) -> dict:
    chars      = len(text)
    words_list = re.findall(r"\b\w+\b", text, flags=re.UNICODE)
    words      = len(words_list)
    avg_cw     = (sum(len(w) for w in words_list) / words) if words else 0.0
    longest    = max((len(w) for w in words_list), default=0)
    if chars:
        punct = sum(1 for ch in text if ch in PUNCTUATION_CHARS) / chars
        sym   = sum(1 for ch in text if ch in SYMBOL_CHARS) / chars
    else:
        punct = sym = 0.0
    return dict(
        chars=float(chars), words=float(words),
        avg_chars_per_word=float(avg_cw), punctuation_ratio=float(punct),
        symbol_ratio=float(sym), longest_word_chars=float(longest),
        vocab_size=float(vocab_size),
    )

# ── inference ─────────────────────────────────────────────────────────────────

@torch.inference_mode()
def predict(stats: dict) -> float:
    x = torch.tensor(
        [[stats["chars"], stats["words"], stats["avg_chars_per_word"],
          stats["punctuation_ratio"], stats["symbol_ratio"],
          stats["longest_word_chars"], stats["vocab_size"]]],
        dtype=torch.float32, device=DEVICE,
    )
    out = load_model()(input_features=x)
    return max(0.0, float(out.logits.squeeze().item()))

# ── event handlers ────────────────────────────────────────────────────────────

def on_text_change(text: str):
    text = text or ""
    chars = len(text)
    words = len(re.findall(r"\b\w+\b", text, flags=re.UNICODE))
    return f"<div class='live-counter'><span>{chars:,} chars</span><span class='sep'>·</span><span>{words:,} words</span></div>"


def on_tokenizer_change(tokenizer_id: str):
    tid = (tokenizer_id or "").strip()
    if not tid:
        return gr.update(interactive=True), _status(""), None
    try:
        tok = load_tokenizer(tid)
        vs  = get_vocab_size(tok)
        return (
            gr.update(value=vs, interactive=False),
            _status(f"Locked to <b>{html.escape(tid)}</b> — vocab {vs:,}", kind="ok"),
            vs,
        )
    except Exception as exc:
        return (
            gr.update(interactive=True),
            _status(f"Could not load <b>{html.escape(tid)}</b>: {html.escape(str(exc)[:120])}", kind="err"),
            None,
        )


def on_clear(current_vocab):
    return gr.update(value=""), gr.update(interactive=True), None, _status("")


def run(text: str, vocab_size_val, tokenizer_id: str, locked_vocab: Optional[int]):
    text = (text or "").strip()
    tid  = (tokenizer_id or "").strip()

    actual_count: Optional[int] = None
    tok_error:    Optional[str] = None

    if tid:
        try:
            tok            = load_tokenizer(tid)
            resolved_vocab = locked_vocab if locked_vocab is not None else get_vocab_size(tok)
            ids            = tok(text, add_special_tokens=False).input_ids
            actual_count   = len(ids)
        except Exception as exc:
            tok_error      = str(exc)
            resolved_vocab = _safe_int(vocab_size_val, DEFAULT_VOCAB_SIZE)
    else:
        resolved_vocab = _safe_int(vocab_size_val, DEFAULT_VOCAB_SIZE)

    stats = compute_stats(text, resolved_vocab)
    try:
        pred = predict(stats)
    except Exception as exc:
        return _render_error(str(exc)), None

    result_data = dict(
        prediction=pred, actual_count=actual_count,
        vocab_size=resolved_vocab, tokenizer_id=tid,
        stats=stats, tok_error=tok_error,
    )
    return _render_results(result_data), result_data


def _safe_int(val, default: int) -> int:
    try:
        return int(float(val))
    except Exception:
        return default


def _status(msg: str, kind: str = "") -> str:
    if not msg:
        return ""
    icon = "✓" if kind == "ok" else ("✗" if kind == "err" else "ℹ")
    cls  = f"status-{kind}" if kind else ""
    return f"<div class='status-pill {cls}'><span class='status-icon'>{icon}</span>{msg}</div>"

# ── HTML rendering ─────────────────────────────────────────────────────────────

def _render_error(msg: str) -> str:
    return f"""
<div class="result-wrap">
  <div class="error-banner">
    <span class="err-icon">⚠</span>
    <div>
      <div class="err-title">Something went wrong</div>
      <div class="err-body">{html.escape(msg)}</div>
    </div>
  </div>
</div>"""


def _render_results(r: dict) -> str:
    pred         = r["prediction"]
    actual       = r["actual_count"]
    vocab        = r["vocab_size"]
    tid_raw      = r["tokenizer_id"]
    stats        = r["stats"]
    tok_error    = r["tok_error"]
    tid          = html.escape(tid_raw) if tid_raw else None

    pred_int  = round(pred)
    chars_int = int(stats["chars"])
    words_int = int(stats["words"])

    # ── comparison section ──
    if actual is not None:
        diff     = pred_int - actual
        abs_diff = abs(diff)
        pct      = abs_diff / max(actual, 1) * 100
        accuracy = max(0.0, 100.0 - pct)
        bar_w    = min(100, round(accuracy))

        if abs_diff == 0:
            diff_label = "exact match"
            diff_cls   = "diff-exact"
            diff_sign  = ""
        else:
            sign       = "+" if diff > 0 else "−"
            diff_sign  = f"{sign}{abs_diff:,}"
            diff_label = f"{pct:.1f}% off"
            diff_cls   = "diff-over" if diff > 0 else "diff-under"

        bar_color = "#4ade80" if accuracy >= 95 else ("#facc15" if accuracy >= 80 else "#f87171")

        comparison = f"""
    <div class="compare-block">
      <div class="compare-cards">
        <div class="ccard predicted">
          <div class="ccard-label">Predicted</div>
          <div class="ccard-num">{pred_int:,}</div>
        </div>
        <div class="ccard-divider">vs</div>
        <div class="ccard actual">
          <div class="ccard-label">Actual</div>
          <div class="ccard-num">{actual:,}</div>
        </div>
      </div>
      <div class="accuracy-row">
        <div class="accuracy-bar-bg">
          <div class="accuracy-bar-fill" style="width:{bar_w}%; background:{bar_color};"></div>
        </div>
        <div class="accuracy-meta">
          <span class="accuracy-pct">{accuracy:.1f}% accuracy</span>
          <span class="{diff_cls}">{diff_sign and diff_sign + " · "}{diff_label}</span>
        </div>
      </div>
    </div>"""

    elif tid_raw and tok_error:
        comparison = f"""
    <div class="tok-error">
      <span class="tok-err-icon">⚠</span>
      Tokenizer error: {html.escape((tok_error or "")[:120])}
    </div>"""
    else:
        comparison = ""

    # ── feature chips ──
    def chip(label, value):
        return f'<div class="chip"><span class="chip-label">{label}</span><span class="chip-val">{value}</span></div>'

    chips = "".join([
        chip("chars",        f"{chars_int:,}"),
        chip("words",        f"{words_int:,}"),
        chip("avg chars/wd", f"{stats['avg_chars_per_word']:.2f}"),
        chip("longest word", f"{int(stats['longest_word_chars'])}"),
        chip("punct ratio",  f"{stats['punctuation_ratio']:.4f}"),
        chip("symbol ratio", f"{stats['symbol_ratio']:.4f}"),
    ])

    meta_tok  = f'<span class="meta-tag">{tid}</span>' if tid else '<span class="meta-tag muted">no tokenizer</span>'
    meta_vocab = f'<span class="meta-tag">{vocab:,} vocab</span>'
    device_tag = f'<span class="meta-tag">{DEVICE.type.upper()}</span>'

    return f"""
<div class="result-wrap">
  <div class="result-header">
    <div class="result-meta">{meta_tok}{meta_vocab}{device_tag}</div>
  </div>

  <div class="pred-hero">
    <div class="pred-label">Estimated token count</div>
    <div class="pred-number">{pred_int:,}</div>
    <div class="pred-raw">{pred:.5f} · {chars_int / max(pred_int,1):.2f} chars/token</div>
  </div>

  {comparison}

  <div class="features-section">
    <div class="features-title">Text features</div>
    <div class="chips">{chips}</div>
  </div>
</div>"""

# ── CSS ───────────────────────────────────────────────────────────────────────

CSS = """
:root {
  --bg:       #070707;
  --surf:     #0f0f0f;
  --surf2:    #161616;
  --border:   #222;
  --border2:  #2e2e2e;
  --text:     #efefef;
  --muted:    #777;
  --muted2:   #555;
  --green:    #4ade80;
  --yellow:   #facc15;
  --red:      #f87171;
  --blue:     #60a5fa;
  --r:        12px;
  --r-sm:     8px;
}

*, *::before, *::after { box-sizing: border-box; }

body, .gradio-container {
  background: var(--bg) !important;
  color: var(--text) !important;
  font-family: ui-sans-serif, system-ui, -apple-system, sans-serif !important;
}
.gradio-container {
  max-width: 1060px !important;
  margin: 0 auto !important;
  padding: 24px 20px !important;
}

h1,h2,h3,h4,p,span,label,div,textarea,input,select,button {
  color: var(--text) !important;
}

.mono { font-family: ui-monospace, "SF Mono", Menlo, monospace !important; }

/* ── hero ── */
.hero {
  padding: 22px 26px 20px;
  border: 1px solid var(--border2);
  border-radius: 18px;
  background: linear-gradient(145deg, rgba(255,255,255,0.038) 0%, rgba(255,255,255,0.008) 100%);
  margin-bottom: 22px;
}
.hero-eyebrow {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  color: var(--muted) !important;
  margin-bottom: 8px;
}
.hero-title {
  font-size: 28px;
  font-weight: 800;
  letter-spacing: -0.03em;
  line-height: 1;
  margin-bottom: 10px;
}
.hero-desc {
  font-size: 13.5px;
  color: var(--muted) !important;
  line-height: 1.6;
  max-width: 680px;
}
.hero-badges {
  display: flex;
  gap: 6px;
  margin-top: 14px;
  flex-wrap: wrap;
}
.badge {
  font-size: 11px;
  font-weight: 600;
  padding: 3px 10px;
  border-radius: 99px;
  border: 1px solid var(--border2);
  color: var(--muted) !important;
  background: var(--surf);
  letter-spacing: 0.04em;
}

/* ── gradio internals ── */
.block, .block-container, .group, .wrap, .panel, .form {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
}

/* ── inputs ── */
textarea, input[type=text], input[type=number], select {
  background: var(--surf) !important;
  border: 1px solid var(--border2) !important;
  border-radius: var(--r) !important;
  color: var(--text) !important;
  box-shadow: none !important;
  transition: border-color 0.15s !important;
}
textarea:focus, input:focus {
  border-color: #3a3a3a !important;
  outline: none !important;
}
textarea::placeholder, input::placeholder {
  color: var(--muted2) !important;
}
.label-wrap label, .svelte-1gfkn6j {
  font-size: 12px !important;
  font-weight: 600 !important;
  letter-spacing: 0.04em !important;
  text-transform: uppercase !important;
  color: var(--muted) !important;
  margin-bottom: 6px !important;
}

/* ── buttons ── */
button {
  border-radius: var(--r) !important;
  border: 1px solid var(--border2) !important;
  font-weight: 600 !important;
  transition: opacity 0.15s, transform 0.1s !important;
}
button.primary {
  background: #fff !important;
  color: #000 !important;
  border-color: #fff !important;
  letter-spacing: 0.01em !important;
}
button.primary:hover { opacity: 0.88 !important; }
button.primary:active { transform: scale(0.98) !important; }
button.secondary {
  background: var(--surf) !important;
  color: var(--muted) !important;
}
button.secondary:hover { border-color: #3a3a3a !important; color: var(--text) !important; }

/* ── live counter ── */
.live-counter {
  font-size: 12px;
  color: var(--muted) !important;
  padding: 6px 2px 0;
  display: flex;
  gap: 0;
  align-items: center;
}
.live-counter .sep { margin: 0 8px; opacity: 0.35; }

/* ── status pill ── */
.status-pill {
  font-size: 12.5px;
  line-height: 1.5;
  padding: 8px 12px;
  border-radius: var(--r-sm);
  border: 1px solid var(--border);
  background: var(--surf);
  color: var(--muted) !important;
  display: flex;
  align-items: flex-start;
  gap: 8px;
}
.status-pill.status-ok  { border-color: rgba(74,222,128,0.25); background: rgba(74,222,128,0.06); }
.status-pill.status-err { border-color: rgba(248,113,113,0.25); background: rgba(248,113,113,0.06); }
.status-icon { opacity: 0.7; flex-shrink: 0; margin-top: 1px; }
.status-pill b { font-weight: 600; color: inherit !important; }

/* ── accordion ── */
details, .accordion {
  border: 1px solid var(--border) !important;
  border-radius: var(--r) !important;
  background: var(--surf) !important;
}

/* ═══════════════════════════════
   RESULT PANEL
═══════════════════════════════ */
.result-wrap {
  display: flex;
  flex-direction: column;
  gap: 14px;
}

/* ── header meta ── */
.result-header { display: flex; align-items: center; justify-content: space-between; }
.result-meta { display: flex; gap: 6px; flex-wrap: wrap; }
.meta-tag {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.05em;
  padding: 3px 10px;
  border-radius: 99px;
  border: 1px solid var(--border2);
  background: var(--surf2);
  color: var(--muted) !important;
}
.meta-tag.muted { opacity: 0.5; }

/* ── prediction hero ── */
.pred-hero {
  padding: 28px 26px 24px;
  border: 1px solid var(--border2);
  border-radius: 16px;
  background: linear-gradient(145deg, rgba(255,255,255,0.042) 0%, rgba(255,255,255,0.008) 100%);
  text-align: center;
}
.pred-label {
  font-size: 11px;
  font-weight: 700;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  color: var(--muted) !important;
  margin-bottom: 12px;
}
.pred-number {
  font-size: 64px;
  font-weight: 900;
  line-height: 1;
  letter-spacing: -0.04em;
  font-variant-numeric: tabular-nums;
  background: linear-gradient(135deg, #ffffff 0%, rgba(255,255,255,0.6) 100%);
  -webkit-background-clip: text;
  -webkit-text-fill-color: transparent;
  background-clip: text;
  margin-bottom: 10px;
}
.pred-raw {
  font-size: 12px;
  color: var(--muted2) !important;
  font-family: ui-monospace, monospace;
  letter-spacing: 0.02em;
}

/* ── comparison ── */
.compare-block {
  border: 1px solid var(--border2);
  border-radius: 16px;
  background: var(--surf);
  padding: 18px 20px 16px;
}
.compare-cards {
  display: flex;
  align-items: center;
  gap: 12px;
  margin-bottom: 16px;
}
.ccard {
  flex: 1;
  padding: 14px 16px;
  border-radius: var(--r);
  border: 1px solid var(--border);
  background: var(--surf2);
  text-align: center;
}
.ccard.predicted { border-color: rgba(255,255,255,0.1); }
.ccard.actual    { border-color: rgba(255,255,255,0.06); }
.ccard-label {
  font-size: 10px;
  font-weight: 700;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--muted) !important;
  margin-bottom: 6px;
}
.ccard-num {
  font-size: 28px;
  font-weight: 800;
  letter-spacing: -0.03em;
  font-variant-numeric: tabular-nums;
}
.ccard-divider {
  font-size: 12px;
  font-weight: 600;
  color: var(--muted2) !important;
  letter-spacing: 0.08em;
  flex-shrink: 0;
}
.accuracy-row { display: flex; flex-direction: column; gap: 6px; }
.accuracy-bar-bg {
  height: 5px;
  border-radius: 99px;
  background: var(--border);
  overflow: hidden;
}
.accuracy-bar-fill {
  height: 100%;
  border-radius: 99px;
  transition: width 0.4s ease;
}
.accuracy-meta {
  display: flex;
  justify-content: space-between;
  font-size: 12px;
}
.accuracy-pct { font-weight: 700; color: var(--text) !important; }
.diff-exact { color: var(--green) !important; font-weight: 600; }
.diff-over  { color: var(--red)   !important; }
.diff-under { color: var(--blue)  !important; }

/* ── tokenizer error ── */
.tok-error {
  padding: 12px 14px;
  border-radius: var(--r);
  border: 1px solid rgba(248,113,113,0.25);
  background: rgba(248,113,113,0.05);
  font-size: 12.5px;
  color: var(--muted) !important;
  display: flex;
  gap: 10px;
  align-items: flex-start;
}
.tok-err-icon { color: #f87171 !important; flex-shrink: 0; font-size: 14px; }

/* ── features ── */
.features-section {
  border: 1px solid var(--border);
  border-radius: 16px;
  background: var(--surf);
  padding: 16px 18px;
}
.features-title {
  font-size: 11px;
  font-weight: 700;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--muted) !important;
  margin-bottom: 12px;
}
.chips {
  display: grid;
  grid-template-columns: repeat(3, 1fr);
  gap: 7px;
}
@media (max-width: 560px) { .chips { grid-template-columns: repeat(2, 1fr); } }
.chip {
  display: flex;
  justify-content: space-between;
  align-items: center;
  padding: 8px 11px;
  border: 1px solid var(--border);
  border-radius: var(--r-sm);
  background: var(--surf2);
  gap: 8px;
  min-width: 0;
}
.chip-label {
  font-size: 11px;
  color: var(--muted) !important;
  white-space: nowrap;
  overflow: hidden;
  text-overflow: ellipsis;
}
.chip-val {
  font-size: 12px;
  font-weight: 700;
  font-family: ui-monospace, monospace;
  flex-shrink: 0;
}

/* ── error banner ── */
.error-banner {
  display: flex;
  gap: 14px;
  align-items: flex-start;
  padding: 18px 20px;
  border: 1px solid rgba(248,113,113,0.3);
  border-radius: 16px;
  background: rgba(248,113,113,0.06);
}
.err-icon  { font-size: 18px; color: #f87171 !important; flex-shrink: 0; margin-top: 2px; }
.err-title { font-size: 14px; font-weight: 700; margin-bottom: 4px; }
.err-body  { font-size: 13px; color: var(--muted) !important; line-height: 1.5; }

/* ── empty state ── */
.empty-state {
  padding: 40px 24px;
  text-align: center;
  border: 1px dashed var(--border2);
  border-radius: 16px;
}
.empty-icon  { font-size: 28px; margin-bottom: 12px; opacity: 0.3; }
.empty-title { font-size: 15px; font-weight: 700; margin-bottom: 6px; }
.empty-desc  { font-size: 13px; color: var(--muted) !important; line-height: 1.6; }
"""

EMPTY_HTML = """
<div class="empty-state">
  <div class="empty-icon">⬡</div>
  <div class="empty-title">Ready to predict</div>
  <div class="empty-desc">
    Paste text above and press <b>Predict</b>.<br>
    Add a tokenizer repo ID to compare against ground-truth token count.
  </div>
</div>
"""

# ── UI layout ──────────────────────────────────────────────────────────────────

with gr.Blocks(title="JetonCount") as demo:

    gr.HTML(f"""
    <div class="hero">
      <div class="hero-eyebrow">Token Count Estimator</div>
      <div class="hero-title">JetonCount</div>
      <div class="hero-desc">
        Predict how many tokens a text will produce — without running a full tokenizer.
        Optionally compare against any Hugging Face tokenizer for accuracy metrics.
      </div>
      <div class="hero-badges">
        <span class="badge">MLP regressor</span>
        <span class="badge">fromziro/JetonCount</span>
        <span class="badge">{DEVICE.type.upper()}</span>
      </div>
    </div>
    """)

    locked_vocab = gr.State(None)

    with gr.Row(equal_height=False):

        # ── left ────────────────────────────────────────────────────────────
        with gr.Column(scale=5, min_width=360):
            text_in = gr.Textbox(
                label="Text",
                lines=14,
                placeholder="Paste your text here…",
                container=True,
            )
            counter_html = gr.HTML(value="<div class='live-counter'><span>0 chars</span><span class='sep'>·</span><span>0 words</span></div>")
            predict_btn  = gr.Button("⬡  Predict tokens", variant="primary", size="lg")

        # ── right ───────────────────────────────────────────────────────────
        with gr.Column(scale=3, min_width=260):
            tokenizer_in = gr.Textbox(
                label="Tokenizer repo (optional)",
                placeholder="e.g. openai-community/gpt2",
            )
            vocab_in = gr.Number(
                label="Vocab size",
                value=DEFAULT_VOCAB_SIZE,
                precision=0,
                interactive=True,
            )
            status_html = gr.HTML(value="")
            clear_btn = gr.Button("Clear tokenizer", variant="secondary", size="sm")

    results_html = gr.HTML(value=EMPTY_HTML)

    with gr.Accordion("Raw JSON", open=False):
        raw_json = gr.JSON(label="")

    # ── wiring ────────────────────────────────────────────────────────────────

    text_in.change(fn=on_text_change, inputs=[text_in], outputs=[counter_html])

    tokenizer_in.blur(
        fn=on_tokenizer_change, inputs=[tokenizer_in],
        outputs=[vocab_in, status_html, locked_vocab],
    )
    tokenizer_in.submit(
        fn=on_tokenizer_change, inputs=[tokenizer_in],
        outputs=[vocab_in, status_html, locked_vocab],
    )
    clear_btn.click(
        fn=on_clear, inputs=[vocab_in],
        outputs=[tokenizer_in, vocab_in, locked_vocab, status_html],
    )
    predict_btn.click(
        fn=run, inputs=[text_in, vocab_in, tokenizer_in, locked_vocab],
        outputs=[results_html, raw_json],
    )
    text_in.submit(
        fn=run, inputs=[text_in, vocab_in, tokenizer_in, locked_vocab],
        outputs=[results_html, raw_json],
    )


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Monochrome(), css=CSS)