"""
Riddle Sprite
=============
Whisper a secret word. The sprite spins a little riddle for it that you can pose
to a friend.

Track 2 toy for the "Small Models Big Adventures" hackathon.
Model: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 (4B, well under the 32B ceiling).

Nemotron is a reasoning model, so it likes to think out loud; we turn that off and
strip any stray reasoning, keeping only the riddle.

Modes: model (Nemotron) / keeper (a templated riddle, so it always demos).
"""

import os
import re
import html
import random
import inspect

import gradio as gr

_BLOCKS_HAS_CSS = "css" in inspect.signature(gr.Blocks.__init__).parameters
_LAUNCH_HAS_SSR = "ssr_mode" in inspect.signature(gr.Blocks.launch).parameters

MODEL_ID = os.environ.get("RIDDLE_MODEL", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
DEBUG = os.environ.get("RIDDLE_DEBUG", "").strip().lower() in {"1", "true", "yes"}
MAX_NEW_TOKENS = 160

SYSTEM_PROMPT = (
    "/no_think\n"
    "You are a playful riddle sprite. Given a secret word, write ONE short riddle "
    "(2 to 4 lines) whose answer is that word. Never use the word itself or any plural "
    "or obvious form of it. Be whimsical and fair. End with the line: What am I? "
    "Output only the riddle, no preamble, no explanation, no answer."
)

EXAMPLES = ["moon", "umbrella", "shadow", "library"]

def _noop_gpu(*a, **k):
    def wrap(fn):
        return fn
    return wrap(a[0]) if a and callable(a[0]) else wrap

if os.environ.get("SPACES_ZERO_GPU", "").lower() in {"true", "1"}:
    try:
        import spaces
        GPU = spaces.GPU
    except Exception:                          # noqa: BLE001
        GPU = _noop_gpu
else:
    GPU = _noop_gpu

_tokenizer = None
_model = None
MODE = "keeper"
_warmed = False


def load_model():
    global _tokenizer, _model, MODE
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto",
        )
        _model.eval()
        MODE = "model"
        print(f"[Riddle] Loaded {MODEL_ID} -- model mode.")
    except Exception as exc:                    # noqa: BLE001
        MODE = "keeper"
        print(f"[Riddle] Could not load {MODEL_ID} ({exc}). Keeper mode active.")


load_model()


@GPU(duration=40)
def _model_generate(word):
    import torch
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Secret word: {word}"},
    ]
    enc = _tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
    ).to(_model.device)
    input_len = enc["input_ids"].shape[1]
    with torch.no_grad():
        out = _model.generate(
            **enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=True,
            temperature=0.8, top_p=0.9, pad_token_id=_tokenizer.eos_token_id,
        )
    return _tokenizer.decode(out[0][input_len:], skip_special_tokens=True).strip()


_K_TEMPLATES = [
    "I have no voice, yet I am known,\nI come to you when light has flown.\nLook for me where things are hidden.\nWhat am I?",
    "Small and patient, I wait unseen,\nuntil you need me, then I lean.\nI guard you from what falls above.\nWhat am I?",
    "I follow close but make no sound,\nI stretch and shrink along the ground.\nWithout the light I disappear.\nWhat am I?",
]


def _keeper_riddle(word):
    rng = random.Random(word.strip().lower())
    return rng.choice(_K_TEMPLATES)


def _clean(raw, word):
    # strip reasoning traces in any of the forms a reasoning model might emit
    raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.S | re.I)
    raw = re.sub(r"<think>.*$", "", raw, flags=re.S | re.I)
    raw = re.sub(r"^.*?</think>", "", raw, flags=re.S | re.I)
    raw = raw.strip().strip('"').strip()
    # reject if empty, leaks the answer, or rambles like reasoning
    low = raw.lower()
    w = word.strip().lower()
    if (not raw or w and w in low or len(raw) > 320
            or re.search(r"\b(I should|let me|the answer|the word)\b", raw, re.I)):
        return ""
    return raw


def make_riddle(word):
    raw = ""
    if MODE == "model":
        try:
            raw = _model_generate(word)
        except Exception as exc:                # noqa: BLE001
            print(f"[Riddle] generation error: {exc}")
            raw = ""
    riddle = _clean(raw, word) if raw else ""
    if not riddle:
        riddle = _keeper_riddle(word)
    return riddle


def esc(s):
    return html.escape(str(s))


def render(riddle, word, revealed):
    if not riddle:
        return ('<div class="scroll empty">Whisper a secret word below, '
                'then ask for a riddle.</div>')
    body = esc(riddle).replace("\n", "<br>")
    ans = (f'<div class="answer">answer: <b>{esc(word)}</b></div>' if revealed
           else '<div class="answer hidden">answer hidden</div>')
    return f'<div class="scroll"><div class="riddle">{body}</div>{ans}</div>'


def on_make(word, _state):
    global _warmed
    word = (word or "").strip()
    if not word:
        yield ('<div class="scroll empty">The sprite needs a secret word first.</div>',
               {"riddle": "", "word": "", "revealed": False}, "")
        return
    if MODE == "model" and not _warmed:
        yield ('<div class="scroll waking">🧚 <i>The sprite is thinking up a riddle '
               '(the first one takes a moment)...</i></div>',
               {"riddle": "", "word": word, "revealed": False}, "")
    riddle = make_riddle(word)
    if MODE == "model":
        _warmed = True
    state = {"riddle": riddle, "word": word, "revealed": False}
    yield render(riddle, word, False), state, ""


def on_reveal(state):
    if not state or not state.get("riddle"):
        return render("", "", False), state
    state = {**state, "revealed": not state.get("revealed")}
    return render(state["riddle"], state["word"], state["revealed"]), state


def use_example(text):
    return text


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Fraunces:ital,opsz,wght@0,9..144,400;0,9..144,600;1,9..144,400&family=Spectral:ital,wght@0,400;0,500;1,400&display=swap');
:root{--paper:#f3ead7;--paper-2:#efe3c9;--ink:#33301f;--ink-soft:#6b5a47;
  --moss:#5a6b3a;--plum:#6d4b86;--line:#c9b48f;}
.gradio-container,.gradio-container.dark,.dark{
  --body-background-fill:transparent;--background-fill-primary:#fffaf0;--background-fill-secondary:#f6edd8;
  --block-background-fill:#fffaf0;--block-border-color:var(--line);--border-color-primary:var(--line);
  --body-text-color:var(--ink);--body-text-color-subdued:var(--ink-soft);
  --block-label-text-color:var(--moss);--block-title-text-color:var(--ink);
  --block-label-background-fill:#efe3c9;--block-title-background-fill:transparent;
  --input-background-fill:#fffaf0;--input-border-color:var(--line);--input-placeholder-color:var(--ink-soft);
  --button-primary-background-fill:var(--plum);--button-primary-background-fill-hover:#583a6e;
  --button-primary-text-color:#fbf7ec;--button-primary-border-color:#583a6e;
  --button-secondary-background-fill:#efe3c9;--button-secondary-background-fill-hover:#e7d6b3;
  --button-secondary-text-color:var(--ink);--button-secondary-border-color:var(--line);
  --color-accent:var(--plum);--color-accent-soft:#f0e6f3;}
.gradio-container{background:radial-gradient(120% 80% at 80% -10%,#f7efdd,var(--paper) 55%,var(--paper-2));
  font-family:'Spectral',Georgia,serif !important;color:var(--ink) !important;max-width:860px !important;}
.gradio-container textarea,.gradio-container input[type="text"],.gradio-container input:not([type]){
  background:#fffaf0 !important;color:var(--ink) !important;-webkit-text-fill-color:var(--ink) !important;border-color:var(--line) !important;}
.gradio-container textarea::placeholder,.gradio-container input::placeholder{color:var(--ink-soft) !important;-webkit-text-fill-color:var(--ink-soft) !important;opacity:1;}
.rd-title{font-family:'Fraunces',serif;font-weight:600;font-size:2.5rem;line-height:1;margin:.2rem 0 0;}
.rd-title em{font-style:italic;color:var(--plum);}
.rd-sub{font-style:italic;color:var(--ink-soft);margin:.35rem 0 1rem;font-size:1.05rem;}
.rd-mode{display:inline-block;font-size:.72rem;letter-spacing:.12em;text-transform:uppercase;color:var(--moss);border:1px solid var(--line);border-radius:999px;padding:.15rem .6rem;}
.scroll{background:#fffaf0;border:1px solid var(--line);border-radius:16px;padding:24px 26px;
  box-shadow:0 14px 36px -22px rgba(51,48,31,.7);text-align:center;}
.scroll.empty,.scroll.waking{color:var(--ink-soft);font-style:italic;}
.riddle{font-family:'Fraunces',serif;font-size:1.45rem;line-height:1.5;color:var(--ink);font-style:italic;}
.answer{margin-top:16px;padding-top:12px;border-top:1px dashed var(--line);color:var(--moss);}
.answer.hidden{color:var(--ink-soft);font-style:italic;}
.rd-foot{color:var(--ink-soft);font-size:.82rem;font-style:italic;text-align:center;margin-top:10px;}
footer{display:none !important;}
"""

_bk = {"title": "Riddle Sprite"}
if _BLOCKS_HAS_CSS:
    _bk["css"] = CSS
    _bk["theme"] = gr.themes.Soft()

with gr.Blocks(**_bk) as demo:
    state = gr.State({"riddle": "", "word": "", "revealed": False})
    public = "Nemotron-3-Nano-4B · on-device" if MODE == "model" else "Riddle Sprite"
    mode_label = f"{public}  ·  [{MODE}]" if DEBUG else public
    gr.HTML(f"""
      <div><div class="rd-title">Riddle <em>Sprite</em></div>
      <div class="rd-sub">Whisper a secret word. Get a riddle to puzzle a friend.</div>
      <span class="rd-mode">{mode_label}</span></div>""")

    with gr.Row():
        word = gr.Textbox(placeholder="a secret word... (moon, umbrella, library)",
                          show_label=False, scale=8, autofocus=True)
        go = gr.Button("Spin a riddle", variant="primary", scale=1)
    with gr.Row():
        ex_btns = [gr.Button(e, size="sm") for e in EXAMPLES]
    scroll = gr.HTML(render("", "", False))
    reveal = gr.Button("Reveal the answer", size="sm")
    gr.HTML('<div class="rd-foot">Pose the riddle to a friend before you reveal it.</div>')

    go.click(on_make, [word, state], [scroll, state, word])
    word.submit(on_make, [word, state], [scroll, state, word])
    reveal.click(on_reveal, state, [scroll, state])
    for b, e in zip(ex_btns, EXAMPLES):
        b.click(use_example, gr.State(e), word)

if __name__ == "__main__":
    _lk = {}
    if not _BLOCKS_HAS_CSS:
        _lk["css"] = CSS
        _lk["theme"] = gr.themes.Soft()
    if _LAUNCH_HAS_SSR:
        _lk["ssr_mode"] = False
    demo.queue(max_size=24).launch(**_lk)