"""
Hello, Stranger
===============
Write a small, kind message. A wandering creature carries it to a faraway place
and brings back how a friend there might say it.

Track 2 toy for the "Small Models Big Adventures" hackathon.
Model: CohereLabs/aya-expanse-8b (8B, well under 32B, 23 languages).

NOTE: Aya Expanse is gated and CC-BY-NC. To run in model mode you must
  1) accept the licence at https://huggingface.co/CohereLabs/aya-expanse-8b , and
  2) add your read token as a Space secret named HF_TOKEN.
The hackathon is non-commercial, so CC-BY-NC is fine. Without access the app runs
in keeper mode (a small built-in phrasebook) so it still 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("STRANGER_MODEL", "CohereLabs/aya-expanse-8b")
DEBUG = os.environ.get("STRANGER_DEBUG", "").strip().lower() in {"1", "true", "yes"}
MAX_NEW_TOKENS = 150

LANGUAGES = ["Surprise me", "Arabic", "Chinese", "Czech", "Dutch", "French", "German",
             "Greek", "Hebrew", "Hindi", "Indonesian", "Italian", "Japanese", "Korean",
             "Persian", "Polish", "Portuguese", "Romanian", "Russian", "Spanish",
             "Turkish", "Ukrainian", "Vietnamese"]
_PICKABLE = [l for l in LANGUAGES if l != "Surprise me"]

SYSTEM_PROMPT = (
    "You are a warm, well-travelled messenger. The user gives a short, kind message. "
    "Render it into the target language the way a friendly local would actually say "
    "it, then help the user say it back. Reply using EXACTLY these three labels, each "
    "on its own line, and nothing else:\n"
    "Phrase: <the message in the target language, in that language's own script>\n"
    "Sounds like: <a simple pronunciation using English letters>\n"
    "Note: <one short, warm sentence about the phrase or the place>"
)

EXAMPLES = ["I'm proud of you", "I missed you", "Safe travels, friend", "Thank you for everything"]

# tiny offline phrasebook for keeper mode (a handful of common warm phrases)
_PHRASEBOOK = {
    "French": ("Je suis fier de toi", "zhuh swee fee-AIR duh twah", "Said with a hand on the heart in Lyon."),
    "Spanish": ("Estoy orgulloso de ti", "es-TOY or-goo-YOH-so deh tee", "A phrase that warms any plaza at dusk."),
    "Japanese": ("あなたを誇りに思います", "a-na-ta wo ho-ko-ri ni o-moy-mas", "Often paired with a small, sincere bow."),
    "Italian": ("Sono orgoglioso di te", "SOH-no or-go-LYOH-zo dee teh", "Best delivered over an unhurried espresso."),
    "Swahili-ish": ("Ninajivunia wewe", "nee-na-jee-voo-NEE-a WEH-weh", "Carried on a warm coastal breeze."),
}


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
        tok_kwargs = {}
        token = os.environ.get("HF_TOKEN")
        if token:
            tok_kwargs["token"] = token
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, **tok_kwargs)
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", **tok_kwargs,
        )
        _model.eval()
        MODE = "model"
        print(f"[Stranger] Loaded {MODEL_ID} -- model mode.")
    except Exception as exc:                    # noqa: BLE001
        MODE = "keeper"
        print(f"[Stranger] Could not load {MODEL_ID} ({exc}). Keeper mode active.")


load_model()


@GPU(duration=40)
def _model_generate(message, language):
    import torch
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Target language: {language}\nMessage: {message}"},
    ]
    enc = _tokenizer.apply_chat_template(
        messages, tokenize=True, 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.3, top_p=0.9, pad_token_id=_tokenizer.eos_token_id,
        )
    return _tokenizer.decode(out[0][input_len:], skip_special_tokens=True).strip()


_BAD = re.compile(r"<|target language|the message in|pronunciation using|short, warm sentence", re.I)


def _grab(label, text):
    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 _keeper(message, language):
    if language == "Surprise me" or language not in _PHRASEBOOK:
        language = random.choice(list(_PHRASEBOOK.keys()))
    phrase, sounds, note = _PHRASEBOOK[language]
    return {"language": language, "phrase": phrase, "sounds": sounds,
            "note": note + " (The messenger is resting; this is a phrase from the satchel.)"}


def carry(message, language):
    if language == "Surprise me":
        target = random.choice(_PICKABLE)
    else:
        target = language
    raw = ""
    if MODE == "model":
        try:
            raw = _model_generate(message, target)
        except Exception as exc:                # noqa: BLE001
            print(f"[Stranger] generation error: {exc}")
            raw = ""
    if raw:
        phrase = _grab("Phrase", raw)
        sounds = _grab("Sounds like", raw)
        note = _grab("Note", raw)
        if phrase:
            return {"language": target, "phrase": phrase[:200],
                    "sounds": sounds[:200], "note": (note or "Carried home with care.")[:240]}
    return _keeper(message, language)


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


def render(c):
    if not c:
        return ('<div class="card empty">Write a small kind message, choose a place, '
                'and send the messenger.</div>')
    sounds = (f'<div class="crow"><span class="lbl">Sounds like</span>'
              f'<span>{esc(c["sounds"])}</span></div>' if c.get("sounds") else "")
    return f"""
    <div class="card">
      <div class="lang">a friend in <b>{esc(c['language'])}</b> might say</div>
      <div class="phrase">{esc(c['phrase'])}</div>
      {sounds}
      <div class="note">{esc(c['note'])}</div>
    </div>
    """


def on_send(message, language):
    global _warmed
    message = (message or "").strip()
    if not message:
        yield ('<div class="card empty">Give the messenger a few warm words to carry.</div>', "")
        return
    if MODE == "model" and not _warmed:
        yield ('<div class="card waking">✉ <i>The messenger sets off down the road '
               '(the first journey takes a moment)...</i></div>', "")
    c = carry(message, language)
    if MODE == "model":
        _warmed = True
    yield render(c), ""


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:#2f2a22;--ink-soft:#6b5a47;
  --sea:#2f6d6a;--coral:#c2693e;--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(--sea);--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(--sea);--button-primary-background-fill-hover:#245451;
  --button-primary-text-color:#fbf7ec;--button-primary-border-color:#245451;
  --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(--coral);--color-accent-soft:#f3e3c4;}
.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:880px !important;}
.gradio-container textarea,.gradio-container input[type="text"],.gradio-container input:not([type]),.gradio-container select{
  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;}
.hs-title{font-family:'Fraunces',serif;font-weight:600;font-size:2.5rem;line-height:1;margin:.2rem 0 0;}
.hs-title em{font-style:italic;color:var(--coral);}
.hs-sub{font-style:italic;color:var(--ink-soft);margin:.35rem 0 1rem;font-size:1.05rem;}
.hs-mode{display:inline-block;font-size:.72rem;letter-spacing:.12em;text-transform:uppercase;color:var(--sea);border:1px solid var(--line);border-radius:999px;padding:.15rem .6rem;}
.card{background:#fffaf0;border:1px solid var(--line);border-radius:16px;padding:22px 24px;
  box-shadow:0 14px 36px -22px rgba(47,42,34,.7);text-align:center;}
.card.empty,.card.waking{color:var(--ink-soft);font-style:italic;}
.lang{color:var(--sea);font-style:italic;margin-bottom:.5rem;}
.phrase{font-family:'Fraunces',serif;font-size:2rem;line-height:1.25;color:var(--ink);margin:.2rem 0 .6rem;}
.crow{margin:.3rem 0;}.crow .lbl{font-size:.7rem;letter-spacing:.08em;text-transform:uppercase;color:var(--coral);margin-right:8px;}
.note{margin-top:.8rem;padding-top:.7rem;border-top:1px dashed var(--line);color:var(--ink);font-style:italic;}
.hs-foot{color:var(--ink-soft);font-size:.82rem;font-style:italic;text-align:center;margin-top:10px;}
footer{display:none !important;}
"""

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

with gr.Blocks(**_bk) as demo:
    public = "Aya Expanse 8B · on-device" if MODE == "model" else "Hello, Stranger"
    mode_label = f"{public}  ·  [{MODE}]" if DEBUG else public
    gr.HTML(f"""
      <div><div class="hs-title">Hello, <em>Stranger</em></div>
      <div class="hs-sub">Write a kind message. A messenger carries it somewhere far.</div>
      <span class="hs-mode">{mode_label}</span></div>""")

    with gr.Row():
        message = gr.Textbox(placeholder="a small, kind message... (I'm proud of you)",
                             show_label=False, scale=6, autofocus=True)
        lang = gr.Dropdown(LANGUAGES, value="Surprise me", show_label=False, scale=3)
        go = gr.Button("Send it", variant="primary", scale=2)
    with gr.Row():
        ex_btns = [gr.Button(e, size="sm") for e in EXAMPLES]
    card = gr.HTML(render(None))
    gr.HTML('<div class="hs-foot">Phrases are written the way a friendly local might really say them.</div>')

    go.click(on_send, [message, lang], [card, message])
    message.submit(on_send, [message, lang], [card, message])
    for b, e in zip(ex_btns, EXAMPLES):
        b.click(use_example, gr.State(e), message)

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)