"""
TokenScope · Build Small Hackathon · Thousand Token Wood track.

Two modes, both teaching how an LLM works from the inside:
  1 · Guess the token  — next-token prediction (Qwen2.5-3B logits, via Modal).
  2 · Semantle         — semantic similarity / embeddings (NVIDIA Nemotron
                         Embed VL, via Modal). The same math RAG uses.

The models run on Modal endpoints; this Space is pure UI (no GPU). See model.py.
"""

import html
import json
import random
from pathlib import Path

import gradio as gr

import model as M

TEXTS = json.loads((Path(__file__).parent / "texts.json").read_text(encoding="utf-8"))
MAX_ROUNDS = 6

SECRETS = [
    "ocean", "music", "mountain", "coffee", "robot", "dragon", "library", "winter",
    "garden", "planet", "river", "castle", "forest", "thunder", "diamond", "volcano",
    "piano", "galaxy", "desert", "whale", "engine", "festival", "shadow", "harvest",
    "compass", "lantern", "mirror", "bridge", "glacier", "rocket",
]

SEEDS_UI = ["animal", "nature", "machine", "place", "music", "feeling"]

CSS = """
.gradio-container { max-width: 1020px !important; margin: 0 auto !important;
  background: radial-gradient(1100px 520px at 50% -14%, #e7e9ff 0%, rgba(248,249,255,0) 55%) !important; }
#hero { text-align: center; padding: 20px 0 6px; }
#hero h1 { font-size: 46px !important; font-weight: 800 !important; letter-spacing: -1.5px;
  color: #4f46e5 !important; margin: 0 0 4px !important; }
#hero p { color: #64748b !important; font-size: 16px !important; margin: 0 !important; }
button.primary { background: linear-gradient(135deg,#6366f1,#8b5cf6) !important; border: 0 !important;
  color: #fff !important; box-shadow: 0 8px 22px -8px rgba(99,102,241,.65) !important; font-weight: 600 !important; }
button.primary:hover { filter: brightness(1.08); }
button { border-radius: 11px !important; transition: filter .15s ease; }
.block, .tabitem { border-radius: 14px !important; }
"""


# =========================================================================== #
# MODE 1 — Guess the token                                                    #
# =========================================================================== #
def analyze(context, guess):
    return M.analyze(context, guess)


def bars_html(dist, player_rank, real_word):
    real_norm = M.normalize(real_word)
    rows = []
    maxp = max((p for _, p in dist), default=1.0) or 1.0
    for i, (tok, p) in enumerate(dist, start=1):
        label = html.escape(tok.strip() or "·(space)")
        width = max(2, int(p / maxp * 100))
        is_player = player_rank == i
        is_real = M.normalize(tok) == real_norm and real_norm != ""
        color, tag = "#6366f1", ""
        if is_player and is_real:
            color, tag = "#22c55e", " ← your word · and the real one ✓"
        elif is_player:
            color, tag = "#f59e0b", " ← your word"
        elif is_real:
            color, tag = "#22c55e", " ← the real word ✓"
        rows.append(
            f'<div style="margin:6px 0;font-family:ui-monospace,monospace;font-size:14px">'
            f'<div style="display:flex;justify-content:space-between">'
            f'<span><b>{i}.</b> {label}<span style="color:#888">{tag}</span></span>'
            f'<span style="color:#888">{p*100:.1f}%</span></div>'
            f'<div style="background:{color};height:10px;width:{width}%;'
            f'border-radius:5px;margin-top:2px;transition:width .5s ease"></div></div>'
        )
    return "<div>" + "".join(rows) + "</div>"


def new_game():
    item = random.choice(TEXTS)
    words = item["full"][len(item["seed"]):].strip().split()[:MAX_ROUNDS]
    state = {"context": item["seed"], "remaining": words, "score": 0,
             "round": 1, "total": len(words), "categoria": item["categoria"], "done": False}
    return (state, context_md(state), "", "<i>Type your word and hit Guess.</i>",
            score_md(state), gr.update(interactive=True))


def context_md(state):
    return (
        f"<span style='color:#888;font-size:13px'>{state['categoria']} · round "
        f"{state['round']}/{state['total']}</span>\n\n"
        f"## {state['context']} <span style='color:#6366f1'>___</span>"
    )


def score_md(state):
    return f"### 🏆 {state['score']} points"


def submit(state, guess):
    if state is None or state.get("done"):
        return state, gr.update(), guess, "Start a new round.", "", gr.update()
    if not guess.strip():
        return state, gr.update(), guess, "✏️ Type a word first.", score_md(state), gr.update()

    context = state["context"]
    real_word = state["remaining"][0]
    dist, rank, prob = analyze(context, guess)
    pts = M.score_for_rank(rank)
    state["score"] += pts

    if rank is None:
        msg = f"😅 <b>'{guess.strip()}'</b> wasn't in the model's 200 most likely words. +0 points."
    else:
        msg = (f"✅ Your word was in <b>position {rank}</b> "
               f"({prob*100:.1f}% probability). <b>+{pts} points.</b>")
    msg += f"<br>The real word was: <b>{real_word}</b>"

    reveal = (
        f"<p style='font-size:15px'>{msg}</p>"
        f"<p style='color:#888;font-size:13px'>Top-10 the model considered after "
        f"«…{' '.join(context.split()[-5:])}»:</p>" + bars_html(dist, rank, real_word)
    )

    state["context"] = context + " " + real_word
    state["remaining"] = state["remaining"][1:]
    state["round"] += 1

    if not state["remaining"]:
        state["done"] = True
        reveal += final_card(state)
        return state, context_md_done(state), "", reveal, score_md(state), gr.update(interactive=False)
    return state, context_md(state), "", reveal, score_md(state), gr.update(interactive=True)


def context_md_done(state):
    return f"## {state['context']}\n\n<span style='color:#22c55e'>✓ sentence complete</span>"


def final_card(state):
    s, mx = state["score"], state["total"] * 100
    pct = (s / mx * 100) if mx else 0
    if pct >= 70:
        verdict = "🧠 You think almost like the model. You predict tokens like a champ."
    elif pct >= 35:
        verdict = "🙂 Decent intuition for what comes next."
    else:
        verdict = "🎲 Language is more unpredictable than it looks, huh?"
    return (
        "<hr><div style='background:#1e1b4b;color:#e0e7ff;padding:16px;border-radius:12px'>"
        f"<h3>Game over — {s}/{mx} points</h3><p>{verdict}</p>"
        "<p style='font-size:14px;opacity:.85'>You just played inside the core mechanic "
        "of an LLM: <b>predicting the next token</b>. The model doesn't \"know\" the "
        "sentence; it computes, word by word, what's most likely to come next. That's "
        "all it does... millions of times.</p>"
        "<p>Hit <b>New round</b> for another sentence.</p></div>"
    )


# =========================================================================== #
# MODE 2 — Semantle (embeddings)                                              #
# =========================================================================== #
def _heat(sim):
    """Rescale the model's compressed word-cosine range (~0.15-0.60) to 0-100,
    so getting closer in meaning clearly warms up. Returns (pct, label)."""
    pct = max(0, min(100, round((sim - 0.15) / 0.45 * 100)))
    if pct >= 78:
        return pct, "🔥 burning hot"
    if pct >= 56:
        return pct, "🥵 hot"
    if pct >= 38:
        return pct, "🙂 warm"
    if pct >= 20:
        return pct, "🧊 cool"
    return pct, "❄️ cold"


def new_semantle():
    state = {"secret": random.choice(SECRETS), "guesses": [], "won": False}
    intro = ("<i>I'm thinking of a secret word (a common English noun). Tap a starter "
             "below or type a word — you'll see how close you are <b>by meaning</b>.</i>")
    return state, "", intro, ""


def semantle_board(state):
    if not state["guesses"]:
        return ""
    rows = sorted(state["guesses"], key=lambda x: -x[1])[:15]
    out = ["<p style='color:#888;font-size:13px'>Your guesses, closest first:</p>"]
    for w, s in rows:
        pct, label = _heat(s)
        width = max(2, pct)
        out.append(
            f'<div style="margin:5px 0;font-family:ui-monospace,monospace;font-size:14px">'
            f'<div style="display:flex;justify-content:space-between">'
            f'<span>{html.escape(w)}</span>'
            f'<span style="color:#888">{pct} · {label}</span></div>'
            f'<div style="background:#6366f1;height:8px;width:{width}%;'
            f'border-radius:4px;margin-top:2px;transition:width .5s ease"></div></div>'
        )
    return "<div>" + "".join(out) + "</div>"


def semantle_win(state):
    n = len(state["guesses"])
    return (
        "<div style='background:#1e1b4b;color:#e0e7ff;padding:16px;border-radius:12px'>"
        f"<h3>🎉 You found it: {html.escape(state['secret'])} — in {n} guesses!</h3>"
        "<p style='font-size:14px;opacity:.9'>You just played with <b>embeddings</b>: "
        "every word became a vector, and \"closeness\" meant similar <i>meaning</i>, not "
        "similar letters. This is exactly how a <b>RAG</b> works — it turns your question "
        "into a vector and retrieves the nearest chunks of text by distance. The same math "
        "you just used powers semantic search and retrieval-augmented generation.</p>"
        "<p>Hit <b>New word</b> to play again.</p></div>"
    )


def semantle_guess(state, word):
    if state is None:
        state = {"secret": random.choice(SECRETS), "guesses": [], "won": False}
    word = (word or "").strip()
    if state.get("won"):
        return state, "", semantle_win(state), semantle_board(state)
    if not word:
        return state, "", "✏️ Type a word first.", semantle_board(state)

    sim = M.embed_similarity(word, state["secret"])
    state["guesses"].append([word, sim])

    if M.normalize(word) == M.normalize(state["secret"]):
        state["won"] = True
        last = semantle_win(state)
    else:
        pct, label = _heat(sim)
        last = (f"<p style='font-size:16px'>「<b>{html.escape(word)}</b>」 → "
                f"<b>{pct}</b>/100 · {label}</p>")
    return state, "", last, semantle_board(state)


# =========================================================================== #
# UI                                                                          #
# =========================================================================== #
with gr.Blocks(title="TokenScope", theme=gr.themes.Soft(primary_hue="indigo"), css=CSS) as demo:
    gr.Markdown("# 🔭 TokenScope\nLearn how a language model works from the inside — by playing.",
                elem_id="hero")

    with gr.Tabs():
        with gr.Tab("1 · Guess the token"):
            gr.Markdown("Read the cut-off sentence and type the next word. The model shows "
                        "you its top 10 predictions and where yours landed.")
            m1_state = gr.State()
            with gr.Row():
                with gr.Column(scale=3):
                    context_view = gr.Markdown()
                    with gr.Row():
                        guess_box = gr.Textbox(placeholder="What word comes next?",
                                               show_label=False, scale=4, autofocus=True)
                        guess_btn = gr.Button("Guess", variant="primary", scale=1)
                    new_btn = gr.Button("🔄 New round")
                with gr.Column(scale=2):
                    score_view = gr.Markdown()
                    reveal_view = gr.HTML()
            m1_out = [m1_state, context_view, guess_box, reveal_view, score_view, guess_btn]

        with gr.Tab("2 · Semantle"):
            gr.Markdown("Guess my secret word. After each guess you'll see how close you are "
                        "**by meaning** — the foundation of how RAG retrieves.\n\n"
                        "*New here? Tap a starter word, then follow the heat — warmer = closer.*")
            m2_state = gr.State()
            with gr.Row():
                with gr.Column(scale=3):
                    with gr.Row():
                        seed_btns = [gr.Button(s, size="sm") for s in SEEDS_UI]
                    s2_intro = gr.HTML()
                    with gr.Row():
                        s2_box = gr.Textbox(placeholder="Guess a word...",
                                            show_label=False, scale=4, autofocus=True)
                        s2_btn = gr.Button("Guess", variant="primary", scale=1)
                    s2_new = gr.Button("🔄 New word")
                with gr.Column(scale=2):
                    s2_board = gr.HTML()
            m2_out = [m2_state, s2_box, s2_intro, s2_board]

    demo.load(new_game, outputs=m1_out)
    new_btn.click(new_game, outputs=m1_out)
    guess_btn.click(submit, [m1_state, guess_box], m1_out)
    guess_box.submit(submit, [m1_state, guess_box], m1_out)

    demo.load(new_semantle, outputs=m2_out)
    s2_new.click(new_semantle, outputs=m2_out)
    s2_btn.click(semantle_guess, [m2_state, s2_box], m2_out)
    s2_box.submit(semantle_guess, [m2_state, s2_box], m2_out)
    for _b, _s in zip(seed_btns, SEEDS_UI):
        _b.click(lambda st, w=_s: semantle_guess(st, w), [m2_state], m2_out)


if __name__ == "__main__":
    demo.launch()