"""
Pocket Bestiary
===============
Name any ordinary place or thing. Discover the tiny creature hiding in it.

Built for the "Small Models Big Adventures" hackathon (Track 2). Every discovery
is written up as a field-guide card (name, species, habitat, temperament, a note,
a rarity) and added to your collectible field guide for the session.

Model: openbmb/MiniCPM5-1B  (1B params, well under the 32B ceiling, on-device).
A small model has a confident, slightly feral imagination, which is exactly the
right instrument for inventing creatures. Same proven inference path as the rest
of this project.

Modes:
  * model mode  -> MiniCPM5-1B invents each creature
  * keeper mode -> a charming, input-seeded generator so the app always boots and
                   demos even with no GPU and no network
"""

import os
import re
import html
import random
import inspect

import gradio as gr

# --------------------------------------------------------------------------- #
#  Gradio version compatibility (works on 4 / 5 / 6).
# --------------------------------------------------------------------------- #
_BLOCKS_HAS_CSS = "css" in inspect.signature(gr.Blocks.__init__).parameters
_LAUNCH_HAS_SSR = "ssr_mode" in inspect.signature(gr.Blocks.launch).parameters

# --------------------------------------------------------------------------- #
#  Config.
# --------------------------------------------------------------------------- #
MODEL_ID = os.environ.get("BESTIARY_MODEL", "openbmb/MiniCPM5-1B")
DEBUG = os.environ.get("BESTIARY_DEBUG", "").strip().lower() in {"1", "true", "yes"}
MAX_NEW_TOKENS = 170

SYSTEM_PROMPT = (
    "You are the Keeper of the Pocket Bestiary, a gentle naturalist who discovers "
    "tiny imaginary creatures hiding in ordinary places. Given a place or object, "
    "invent ONE small creature that secretly lives there. Reply using EXACTLY these "
    "five labels, each on its own line, and nothing else:\n"
    "Name: <a whimsical creature name, 1 to 3 words>\n"
    "Type: <a fanciful species or order, 2 to 4 words>\n"
    "Habitat: <one short phrase tying it to the given place>\n"
    "Temperament: <2 to 4 words>\n"
    "Note: <one or two warm, surprising sentences>\n"
    "Make the creature's name, look, and habits grow clearly out of the specific "
    "details of the place: its mood, time of day, textures, sounds, and what it is "
    "for. A different place should give a clearly different creature. "
    "Keep everything gentle and suitable for all ages. Never describe or identify "
    "real people; if the place involves a person, invent a creature that lives "
    "near them instead."
)

RARITIES = [
    ("Common", "#6f8456", 50),
    ("Uncommon", "#3a7ca5", 30),
    ("Rare", "#a9641c", 15),
    ("Mythic", "#7a4fa3", 5),
]

EXAMPLES = [
    "the gap behind the fridge",
    "your coat pocket",
    "a bus stop at night",
    "the bottom of a teacup",
]

# --------------------------------------------------------------------------- #
#  ZeroGPU decorator (no-op off ZeroGPU, so the same file runs anywhere).
# --------------------------------------------------------------------------- #
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

# --------------------------------------------------------------------------- #
#  Model loading (guarded; falls back to keeper mode).
# --------------------------------------------------------------------------- #
_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, trust_remote_code=True, torch_dtype=torch.float32,
        )
        _model.eval()
        MODE = "model"
        print(f"[Bestiary] Loaded {MODEL_ID} -- model mode.")
    except Exception as exc:                # noqa: BLE001
        MODE = "keeper"
        print(f"[Bestiary] Could not load {MODEL_ID} ({exc}). Keeper mode active.")


load_model()


@GPU(duration=40)
def _model_generate(user_text):
    import torch
    if torch.cuda.is_available() and next(_model.parameters()).device.type != "cuda":
        _model.to("cuda", dtype=torch.bfloat16)
    # One worked example dramatically improves format adherence for a 1B model.
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Place or object: the dust under a piano"},
        {"role": "assistant", "content":
            "Name: Pedalfluff\n"
            "Type: Lesser Resonance-mote\n"
            "Habitat: nestles in the dust beneath the piano's pedals\n"
            "Temperament: drowsy, music-loving\n"
            "Note: It wakes only for the low notes, and hums along a half-second behind."},
        {"role": "user", "content": f"Place or object: {user_text}"},
    ]

    def _encode(extra):
        return _tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt",
            return_dict=True, **extra
        )

    # Turn off the model's <think> reasoning trace if the template supports it.
    try:
        enc = _encode({"enable_thinking": False})
    except TypeError:
        enc = _encode({})
    enc = enc.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, repetition_penalty=1.1,
            pad_token_id=_tokenizer.eos_token_id,
        )
    return _tokenizer.decode(out[0][input_len:], skip_special_tokens=True).strip()


# --------------------------------------------------------------------------- #
#  Keeper-mode generator: input-seeded so the same place yields the same
#  creature, which feels intentional rather than broken.
# --------------------------------------------------------------------------- #
_PREFIX = ["Dust", "Moss", "Hollow", "Lint", "Ember", "Glimmer", "Brass", "Gloam",
           "Thistle", "Murk", "Candle", "Hush", "Drowse", "Pocket", "Cinder", "Fen"]
_SUFFIX = ["wyrm", "sprite", "mote", "kin", "ling", "snout", "wick", "boggle",
           "fluff", "grub", "nibble", "whisker", "tuft", "snuffle"]
_ORDERS = ["Lesser Hearth-sprite", "Common Pocket-wyrm", "Drowsy Threshold-kin",
           "Crumb-tending Mote", "Shy Corner-boggle", "Lint-spinning Grub",
           "Dust-shy Whiskerling", "Twilight Sill-wick"]
_TEMPERS = ["shy but loyal", "grumbly, fond of quiet", "endlessly curious",
            "drowsy and ceremonious", "mischievous, harmless", "gentle, easily startled"]
_NOTES = [
    "It tidies one small thing each night and takes no credit for it.",
    "It hums a tune only forgotten objects can hear.",
    "It collects the warmth left behind by hands and saves it for winter.",
    "It is mostly polite, and deeply offended by vacuum cleaners.",
    "It keeps the lost buttons company so they feel less alone.",
    "It naps through the day and rearranges the dust into small constellations at night.",
]


def _keeper_creature(user_text):
    rng = random.Random(user_text.strip().lower())
    name = rng.choice(_PREFIX) + rng.choice(_SUFFIX)
    place = user_text.strip()
    return {
        "name": name,
        "type": rng.choice(_ORDERS),
        "habitat": f"makes its home in {place}",
        "temperament": rng.choice(_TEMPERS),
        "note": rng.choice(_NOTES),
    }


# --------------------------------------------------------------------------- #
#  Parse the model's labelled output, degrading gracefully if a 1B model does
#  not follow the format perfectly.
# --------------------------------------------------------------------------- #
_THINK = re.compile(r"<think>.*?</think>", re.S | re.I)
# values that are actually reasoning or echoed template placeholders, not answers
_BAD = re.compile(
    r"<|whimsical creature name|fanciful species|short phrase|one or two|"
    r"\bI'll\b|\bI will\b|\blet me\b|\bmaybe\b|\bI should\b|\bI need\b|"
    r"\bthe response\b|\blabels?\b|\bsomething similar\b|\bfits\b", re.I)


def _strip_think(t):
    t = _THINK.sub("", t)                       # remove closed <think>...</think>
    t = re.sub(r"^.*?</think>", "", t, flags=re.S | re.I)   # leading, up to a close
    t = re.sub(r"<think>.*$", "", t, flags=re.S | re.I)     # unclosed think tail
    return t.strip()


def _grab(label, text):
    # anchor to the start of a line so we never pick up a label named mid-sentence
    m = re.search(rf"^\s*{label}\s*[:\-]\s*(.+)$", text, re.IGNORECASE | re.MULTILINE)
    if not m:
        return ""
    val = m.group(1).strip().strip('"').strip()
    return "" if _BAD.search(val) else val


def parse_creature(raw, user_text):
    raw = _strip_think(raw)
    fields = {
        "name": _grab("Name", raw),
        "type": _grab("Type", raw),
        "habitat": _grab("Habitat", raw),
        "temperament": _grab("Temperament", raw),
        "note": _grab("Note", raw),
    }
    found = sum(1 for v in fields.values() if v)
    # If the small model ignored the format, fall back to a seeded creature. Only
    # reuse its prose as the note if it actually reads like a note, not reasoning.
    if found < 2:
        base = _keeper_creature(user_text)
        cleaned = " ".join(raw.split())
        if cleaned and not _BAD.search(cleaned) and 8 < len(cleaned) < 240:
            base["note"] = cleaned
        return base
    base = _keeper_creature(user_text)   # fills any missing fields sensibly
    for k, v in fields.items():
        if v:
            base[k] = v.rstrip(". ")[:200] if k != "note" else v[:300]
    return base


def roll_rarity():
    pool = []
    for label, color, weight in RARITIES:
        pool += [(label, color)] * weight
    return random.choice(pool)


def discover(user_text):
    raw = ""
    if MODE == "model":
        try:
            raw = _model_generate(user_text)
        except Exception as exc:            # noqa: BLE001 -- never crash mid-demo
            print(f"[Bestiary] generation error: {exc}")
            raw = ""
    creature = parse_creature(raw, user_text) if raw else _keeper_creature(user_text)
    label, color = roll_rarity()
    creature["rarity"] = label
    creature["rarity_color"] = color
    creature["found_in"] = user_text.strip()
    return creature


# --------------------------------------------------------------------------- #
#  Rendering.
# --------------------------------------------------------------------------- #
def esc(s):
    return html.escape(str(s))


def sigil_svg(seed_str, color, size=64):
    """A small deterministic 'pressed specimen' emblem, unique per creature.
    No image model: just seeded shapes in the storybook palette."""
    rng = random.Random("sigil:" + seed_str.lower())
    import math
    cx = cy = size / 2
    body_r = size * 0.28
    parts = [f'<circle cx="{cx}" cy="{cy}" r="{body_r}" fill="#fffaf0" stroke="{color}" stroke-width="2"/>']
    # radial limbs
    limbs = rng.randint(4, 8)
    for i in range(limbs):
        ang = (2 * 3.14159 * i / limbs) + rng.uniform(-0.2, 0.2)
        x1 = cx + body_r * math.cos(ang)
        y1 = cy + body_r * math.sin(ang)
        x2 = cx + (body_r + size * 0.16) * math.cos(ang)
        y2 = cy + (body_r + size * 0.16) * math.sin(ang)
        parts.append(f'<line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" '
                     f'stroke="{color}" stroke-width="2" stroke-linecap="round"/>')
        parts.append(f'<circle cx="{x2:.1f}" cy="{y2:.1f}" r="{size*0.03:.1f}" fill="{color}"/>')
    # eyes
    eyes = rng.choice([1, 2, 2, 3])
    spread = body_r * 0.5
    for j in range(eyes):
        ex = cx + (j - (eyes - 1) / 2) * spread
        ey = cy - body_r * 0.1
        parts.append(f'<circle cx="{ex:.1f}" cy="{ey:.1f}" r="{size*0.045:.1f}" fill="var(--ink, #3a2f25)"/>')
    return (f'<svg viewBox="0 0 {size} {size}" width="{size}" height="{size}" '
            f'xmlns="http://www.w3.org/2000/svg" class="sigil">{"".join(parts)}</svg>')


def render_card(c):
    if not c:
        return ('<div class="card empty">Name a place below, then press '
                '<b>Discover</b> to find what lives there.</div>')
    return f"""
    <div class="card">
      <div class="card-top">
        <span class="rarity" style="--r:{c['rarity_color']}">{esc(c['rarity'])}</span>
        <span class="found">found in {esc(c['found_in'])}</span>
      </div>
      <div class="card-head">
        <div class="sigil-wrap">{sigil_svg(c['name'] + c['found_in'], c['rarity_color'])}</div>
        <div>
          <div class="cname">{esc(c['name'])}</div>
          <div class="ctype">{esc(c['type'])}</div>
        </div>
      </div>
      <div class="crow"><span class="lbl">Habitat</span><span>{esc(c['habitat'])}</span></div>
      <div class="crow"><span class="lbl">Temperament</span><span>{esc(c['temperament'])}</span></div>
      <div class="cnote">{esc(c['note'])}</div>
    </div>
    """


def render_collection(items):
    if not items:
        return '<div class="coll-empty">Your field guide is empty. Go find something.</div>'
    cells = ""
    for c in reversed(items):                # newest first
        cells += (
            f'<div class="mini" style="--r:{c["rarity_color"]}">'
            f'<div class="mini-sigil">{sigil_svg(c["name"] + c["found_in"], c["rarity_color"], size=38)}</div>'
            f'<div class="mini-rar">{esc(c["rarity"])}</div>'
            f'<div class="mini-name">{esc(c["name"])}</div>'
            f'<div class="mini-type">{esc(c["type"])}</div>'
            f'<div class="mini-found">{esc(c["found_in"])}</div></div>'
        )
    return (f'<div class="coll-head">Your field guide '
            f'<span class="count">{len(items)} discovered</span></div>'
            f'<div class="coll-grid">{cells}</div>')


# --------------------------------------------------------------------------- #
#  Handlers.
# --------------------------------------------------------------------------- #
def on_discover(user_text, collection):
    global _warmed
    user_text = (user_text or "").strip()
    if not user_text:
        msg = ('<div class="card empty">Tell the Keeper a place first, '
               'a real one or a tiny one.</div>')
        yield msg, render_collection(collection), collection, ""
        return

    if MODE == "model" and not _warmed:
        yield ('<div class="card waking">🔎 <i>The Keeper turns the lens, focusing '
               '(the first discovery takes a moment)...</i></div>',
               render_collection(collection), collection, "")

    creature = discover(user_text)
    if MODE == "model":
        _warmed = True
    collection = collection + [creature]
    yield render_card(creature), render_collection(collection), collection, ""


def on_reset():
    return render_card(None), render_collection([]), [], ""


def use_example(text):
    return text


# --------------------------------------------------------------------------- #
#  Custom UI (storybook field-guide; design tokens forced light so it is
#  readable regardless of the visitor's browser theme).
# --------------------------------------------------------------------------- #
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:#3a2f25; --ink-soft:#6b5a47;
  --forest:#4a5d3a; --forest-2:#6f8456; --amber:#c98a2e; --amber-deep:#a9641c;
  --line:#c9b48f;
}
/* force the warm light palette onto Gradio's own tokens (readable in dark mode too) */
.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); --border-color-accent:var(--amber);
  --body-text-color:var(--ink); --body-text-color-subdued:var(--ink-soft);
  --block-label-text-color:var(--forest); --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(--amber);
  --button-primary-background-fill-hover:var(--amber-deep);
  --button-primary-text-color:#fffaf0; --button-primary-border-color:var(--amber-deep);
  --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(--amber-deep); --color-accent-soft:#f3e3c4;
}
.gradio-container{
  background:radial-gradient(120% 80% at 80% -10%, #f7efdd 0%, var(--paper) 55%, var(--paper-2) 100%) !important;
  font-family:'Spectral', Georgia, serif !important; color:var(--ink) !important;
  max-width:1100px !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;}

.bz-title{font-family:'Fraunces',serif; font-weight:600; font-size:2.6rem; line-height:1; margin:.2rem 0 0;}
.bz-title em{font-style:italic; color:var(--amber-deep);}
.bz-sub{font-style:italic; color:var(--ink-soft); margin:.35rem 0 1rem; font-size:1.05rem;}
.bz-mode{display:inline-block; font-size:.72rem; letter-spacing:.12em; text-transform:uppercase;
  color:var(--forest); border:1px solid var(--line); border-radius:999px; padding:.15rem .6rem;}

/* the big card */
.card{font-family:'Spectral',serif; color:var(--ink); background:#fffaf0;
  border:1px solid var(--line); border-radius:18px; padding:18px 20px;
  box-shadow:0 14px 36px -22px rgba(58,47,37,.7); position:relative;
  background-image:repeating-linear-gradient(135deg,#fffaf0,#fffaf0 13px,#f9f1de 13px,#f9f1de 26px);}
.card.empty, .card.waking{color:var(--ink-soft); font-style:italic; text-align:center;
  background:#fffaf0; box-shadow:none;}
.card-top{display:flex; justify-content:space-between; align-items:center; margin-bottom:6px;}
.rarity{font-size:.72rem; letter-spacing:.1em; text-transform:uppercase; font-weight:600;
  color:#fffaf0; background:var(--r); border-radius:999px; padding:.18rem .7rem;}
.found{font-size:.82rem; color:var(--ink-soft); font-style:italic;}
.cname{font-family:'Fraunces',serif; font-weight:600; font-size:1.9rem; line-height:1.05; margin:.1rem 0;}
.ctype{color:var(--forest); font-style:italic; margin-bottom:.6rem;}
.card-head{display:flex; gap:14px; align-items:center; margin-bottom:.4rem;}
.sigil-wrap{flex:0 0 auto; background:#fffaf0; border:1px solid var(--line); border-radius:14px; padding:5px;}
.sigil{display:block;}
.mini-sigil{margin-bottom:5px;}
.crow{display:flex; gap:10px; margin:.2rem 0; font-size:.96rem;}
.crow .lbl{flex:0 0 110px; font-size:.7rem; letter-spacing:.08em; text-transform:uppercase;
  color:var(--amber-deep); padding-top:.18rem;}
.cnote{margin-top:.7rem; padding-top:.7rem; border-top:1px dashed var(--line); line-height:1.5;}

/* the collection */
.coll-head{font-family:'Fraunces',serif; font-weight:600; font-size:1.15rem; margin:0 0 10px;
  display:flex; justify-content:space-between; align-items:baseline; color:var(--ink);}
.coll-head .count{font-family:'Spectral',serif; font-size:.85rem; color:var(--ink-soft); font-style:italic;}
.coll-empty{color:var(--ink-soft); font-style:italic; padding:10px 2px;}
.coll-grid{display:grid; grid-template-columns:repeat(auto-fill,minmax(150px,1fr)); gap:10px;}
.mini{background:#fffaf0; border:1px solid var(--line); border-left:5px solid var(--r);
  border-radius:12px; padding:9px 11px;}
.mini-rar{font-size:.62rem; letter-spacing:.1em; text-transform:uppercase; color:var(--r); font-weight:600;}
.mini-name{font-family:'Fraunces',serif; font-weight:600; font-size:1.05rem; line-height:1.1;}
.mini-type{font-size:.8rem; color:var(--forest); font-style:italic;}
.mini-found{font-size:.74rem; color:var(--ink-soft); margin-top:3px;}
.bz-foot{color:var(--ink-soft); font-size:.82rem; font-style:italic; text-align:center; margin-top:10px;}
footer{display:none !important;}
"""

# --------------------------------------------------------------------------- #
#  Layout.
# --------------------------------------------------------------------------- #
_blocks_kwargs = {"title": "Pocket Bestiary"}
if _BLOCKS_HAS_CSS:
    _blocks_kwargs["css"] = CSS
    _blocks_kwargs["theme"] = gr.themes.Soft()

with gr.Blocks(**_blocks_kwargs) as demo:
    collection = gr.State([])

    public_label = "MiniCPM5-1B · on-device" if MODE == "model" else "Pocket Bestiary"
    mode_label = f"{public_label}  ·  [{MODE}]" if DEBUG else public_label
    gr.HTML(
        f"""
        <div>
          <div class="bz-title">Pocket <em>Bestiary</em></div>
          <div class="bz-sub">Name an ordinary place. Discover the creature hiding in it.</div>
          <span class="bz-mode">{mode_label}</span>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                place = gr.Textbox(
                    placeholder="a place or a thing... (the gap behind the fridge, a bus stop at night)",
                    show_label=False, scale=8, autofocus=True,
                )
                go = gr.Button("Discover", variant="primary", scale=1)
            with gr.Row():
                ex_btns = [gr.Button(e, size="sm") for e in EXAMPLES]
            card = gr.HTML(render_card(None))
            reset = gr.Button("Start a new field guide", size="sm")
        with gr.Column(scale=2):
            coll = gr.HTML(render_collection([]))

    gr.HTML('<div class="bz-foot">Every small place has something living in it. '
            'Rarer creatures turn up when you least expect them.</div>')

    go.click(on_discover, [place, collection], [card, coll, collection, place])
    place.submit(on_discover, [place, collection], [card, coll, collection, place])
    reset.click(on_reset, None, [card, coll, collection, place])
    for b, e in zip(ex_btns, EXAMPLES):
        b.click(use_example, gr.State(e), place)


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