app.py

"""
Nyra Chat — Hugging Face Space (Gradio + ZeroGPU).

Local: mock streaming only (no GGUF download).
Spaces: loads HauhauCS Qwen3.6-35B-A3B Aggressive GGUF under @spaces.GPU.
"""

from __future__ import annotations

import inspect
import os
from pathlib import Path

import gradio as gr

from model_engine import (
    DEFAULT_MAX_TOKENS,
    DEFAULT_TEMPERATURE,
    is_spaces,
    runtime_mode,
    start_background_warmup,
    status_label,
    stream_chat,
)
from theme import build_theme

# On Spaces: download GGUF + ensure llama-cpp on CPU before the first GPU call
# so ZeroGPU duration is not wasted on pip/download.
if is_spaces():
    start_background_warmup()

ROOT = Path(__file__).resolve().parent
CSS_PATH = ROOT / "static" / "app.css"
JS_PATH = ROOT / "static" / "app.js"

try:
    import spaces

    GPU = spaces.GPU
except Exception:  # noqa: BLE001

    def GPU(*_args, **_kwargs):  # type: ignore[misc]
        def deco(fn):
            return fn

        if _args and callable(_args[0]) and not _kwargs:
            return _args[0]
        return deco


UI = {
    "pt": {
        "hello": "Oi, eu sou a Nyra.",
        "subtitle": "Chat imersivo · modelo no ZeroGPU do Space.",
        "tagline": "Pense livre. Responda com clareza.",
        "placeholder": "Pergunte qualquer coisa…",
        "send": "Enviar",
        "new_chat": "Novo chat",
        "clear": "Limpar",
        "settings": "Ajustes",
        "temperature": "Temperatura",
        "max_tokens": "Máx. tokens",
        "thinking": "Thinking (Qwen)",
        "lang_btn": "EN",
        "chips": [
            ("⚡ Quantum em 5 bullets", "Explique computação quântica em 5 bullets claros."),
            ("💻 Snippet Python", "Escreva um snippet Python elegante para retry com backoff exponencial."),
            ("🌍 Viagem barata", "Me dê ideias de viagem barata na América do Sul por 7 dias."),
            ("💡 Ideia de app", "Brainstorme um app original com IA para estudantes."),
        ],
        "banner_mock": "Demo local — o modelo **não** baixa aqui; só no Space.",
        "banner_spaces": "ZeroGPU ativo — primeira resposta pode demorar (download/load do GGUF).",
    },
    "en": {
        "hello": "Hi, I'm Nyra.",
        "subtitle": "Immersive chat · model on Space ZeroGPU.",
        "tagline": "Think freely. Answer clearly.",
        "placeholder": "Ask anything…",
        "send": "Send",
        "new_chat": "New chat",
        "clear": "Clear",
        "settings": "Settings",
        "temperature": "Temperature",
        "max_tokens": "Max tokens",
        "thinking": "Thinking (Qwen)",
        "lang_btn": "PT",
        "chips": [
            ("⚡ Quantum in 5 bullets", "Explain quantum computing in 5 clear bullets."),
            ("💻 Python snippet", "Write an elegant Python snippet for retry with exponential backoff."),
            ("🌍 Cheap trip", "Give me cheap 7-day trip ideas in South America."),
            ("💡 App idea", "Brainstorm an original AI app for students."),
        ],
        "banner_mock": "Local demo — model is **not** downloaded here; only on the Space.",
        "banner_spaces": "ZeroGPU on — first reply may take a while (GGUF download/load).",
    },
}


def load_css() -> str:
    return CSS_PATH.read_text(encoding="utf-8") if CSS_PATH.exists() else ""


def load_js() -> str:
    return JS_PATH.read_text(encoding="utf-8") if JS_PATH.exists() else ""


# Critical CSS inlined so HF iframe cannot ignore/stale-cache file CSS
CRITICAL_CSS = """
html,body,.gradio-container,main.app,.wrap,.contain{height:auto!important;min-height:0!important;flex-grow:0!important}
footer,.built-with,button.show-api,button.settings{display:none!important}
#nyra-input-row{height:56px!important;max-height:56px!important;min-height:56px!important;overflow:hidden!important;flex:0 0 56px!important}
#nyra-input textarea,textarea[data-testid=textbox]{height:40px!important;min-height:40px!important;max-height:40px!important;resize:none!important;overflow:hidden!important;line-height:40px!important}
#nyra-chatbot{height:380px!important;max-height:380px!important;min-height:380px!important}
"""

HEAD = f"""
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&family=IBM+Plex+Sans:ital,wght@0,400;0,500;0,600;1,400&family=Syne:wght@600;700;800&display=swap" rel="stylesheet">
<script src="https://unpkg.com/lucide@0.469.0/dist/umd/lucide.min.js"></script>
<meta name="theme-color" content="#07070A">
<meta name="color-scheme" content="dark">
<style id="nyra-critical">{CRITICAL_CSS}</style>
<script>
window.iFrameResizer = {{ autoResize: false, sizeWidth: false }};
</script>
"""


def hero_html(lang: str) -> str:
    t = UI[lang]
    return f"""
<div id="nyra-hero">
  <div class="nyra-hero-icon" aria-hidden="true"><i data-lucide="sparkles"></i></div>
  <h1>{t["hello"]}</h1>
  <p>{t["subtitle"]}<br/><span style="opacity:.85">{t["tagline"]}</span></p>
</div>
"""


def header_brand_html(lang: str) -> str:
    mode = runtime_mode()
    dot = "mock" if mode == "mock" else ""
    return f"""
<div class="nyra-brand-row">
  <div class="nyra-logo-mark" aria-hidden="true"><i data-lucide="sparkles"></i></div>
  <div class="nyra-wordmark">Nyra<span>Chat</span></div>
  <div class="nyra-status">
    <span class="nyra-status-dot {dot}"></span>
    <span>{status_label(lang)}</span>
  </div>
</div>
"""


def banner_text(lang: str) -> str:
    t = UI[lang]
    return t["banner_spaces"] if is_spaces() else t["banner_mock"]


@GPU(duration=90)
def gpu_stream(
    history: list,
    message: str,
    temperature: float,
    max_tokens: float,
    thinking: bool,
):
    """GPU-bound generation — keep duration low for ZeroGPU quota/limits."""
    from model_engine import ensure_weights_ready

    ensure_weights_ready()
    max_tokens = min(int(max_tokens), 512)
    yield from stream_chat(
        history=history or [],
        user_text=message,
        temperature=float(temperature),
        max_tokens=max_tokens,
        thinking=bool(thinking),
    )


def respond(
    message: str,
    history: list,
    temperature: float,
    max_tokens: float,
    thinking: bool,
    lang: str,
):
    history = list(history or [])
    message = (message or "").strip()
    lang = lang if lang in UI else "pt"

    if not message:
        yield history, gr.update()
        return

    prior = list(history)
    history = prior + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": "…"},
    ]
    yield history, ""

    try:
        for partial in gpu_stream(
            prior, message, temperature, max_tokens, thinking
        ):
            history = prior + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": partial},
            ]
            yield history, ""
    except Exception as exc:  # noqa: BLE001
        err = str(exc)
        if "quota" in err.lower() or "ZeroGPU" in err:
            if lang == "pt":
                msg = (
                    "⚠️ **Cota ZeroGPU esgotada** neste período.\n\n"
                    "Entre com sua conta Hugging Face no Space "
                    "para mais cota, ou tente de novo depois do reset.\n\n"
                    f"`{err}`"
                )
            else:
                msg = (
                    "⚠️ **ZeroGPU quota exceeded** for this period.\n\n"
                    "Sign in with your Hugging Face account on the Space "
                    "for more quota, or retry after the quota resets.\n\n"
                    f"`{err}`"
                )
        elif "aborted" in err.lower():
            msg = (
                "⚠️ Tarefa GPU interrompida (timeout/abort). "
                "Tente de novo com uma pergunta mais curta."
                if lang == "pt"
                else "⚠️ GPU task aborted (timeout). "
                "Retry with a shorter prompt."
            )
        else:
            msg = (
                f"Erro ao gerar: `{exc}`"
                if lang == "pt"
                else f"Generation error: `{exc}`"
            )
        history = prior + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": msg},
        ]
        yield history, ""


def chip_prompt(index: int, lang: str) -> str:
    lang = lang if lang in UI else "pt"
    return UI[lang]["chips"][index][1]


def on_chip(index: int, history, temp, max_tok, thinking, lang):
    yield from respond(
        chip_prompt(index, lang), history, temp, max_tok, thinking, lang
    )


def toggle_lang(lang: str):
    new_lang = "en" if lang == "pt" else "pt"
    t = UI[new_lang]
    return (
        new_lang,
        header_brand_html(new_lang),
        hero_html(new_lang),
        gr.update(placeholder=t["placeholder"]),
        gr.update(value=t["send"]),
        gr.update(value=t["new_chat"]),
        gr.update(value=t["clear"]),
        gr.update(value=t["lang_btn"]),
        banner_text(new_lang),
        gr.update(value=t["chips"][0][0]),
        gr.update(value=t["chips"][1][0]),
        gr.update(value=t["chips"][2][0]),
        gr.update(value=t["chips"][3][0]),
    )


def clear_chat():
    return [], ""


def make_chatbot():
    """Gradio 5 needs type='messages'; Gradio 6 dropped it."""
    common = dict(
        value=[],
        elem_id="nyra-chatbot",
        height=380,
        render_markdown=True,
        show_label=False,
        layout="bubble",
    )
    # Prefer Gradio 5 API (Space pins 5.49.1)
    try:
        return gr.Chatbot(**common, type="messages", show_copy_button=True)
    except TypeError:
        pass
    try:
        return gr.Chatbot(**common, buttons=["copy"])
    except TypeError:
        return gr.Chatbot(**common)


def build_app() -> gr.Blocks:
    default_lang = "pt"
    t0 = UI[default_lang]

    # fill_height=False: avoids flex feedback loop that makes the input grow forever
    blocks_kwargs = dict(
        title="Nyra Chat",
        fill_height=False,
        analytics_enabled=False,
        elem_id="nyra-root",
        theme=build_theme(),
        css=load_css(),
        js=load_js(),
        head=HEAD,
    )
    # Gradio 6 may reject theme/css/js/head on Blocks
    try:
        ctx = gr.Blocks(**blocks_kwargs)
    except TypeError:
        ctx = gr.Blocks(
            title="Nyra Chat",
            fill_height=False,
            analytics_enabled=False,
            elem_id="nyra-root",
        )

    with ctx as demo:
        lang_state = gr.State(default_lang)

        with gr.Row(elem_id="nyra-header"):
            brand = gr.HTML(header_brand_html(default_lang))
            with gr.Row():
                lang_btn = gr.Button(
                    t0["lang_btn"],
                    elem_id="nyra-lang",
                    elem_classes=["nyra-ghost"],
                    scale=0,
                    min_width=56,
                )
                new_btn = gr.Button(
                    t0["new_chat"],
                    elem_id="nyra-new-chat",
                    elem_classes=["nyra-ghost"],
                    scale=0,
                    min_width=110,
                )

        with gr.Row(equal_height=False):
            with gr.Column(scale=1, min_width=180, elem_id="nyra-sidebar"):
                clear_btn = gr.Button(
                    t0["clear"],
                    elem_id="nyra-clear",
                    elem_classes=["nyra-ghost"],
                )
                with gr.Accordion(t0["settings"], open=False, elem_id="nyra-settings"):
                    temperature = gr.Slider(
                        0.1,
                        1.5,
                        value=DEFAULT_TEMPERATURE,
                        step=0.05,
                        label=t0["temperature"],
                    )
                    max_tokens = gr.Slider(
                        128,
                        2048,
                        value=DEFAULT_MAX_TOKENS,
                        step=64,
                        label=t0["max_tokens"],
                    )
                    thinking = gr.Checkbox(value=False, label=t0["thinking"])

            with gr.Column(scale=5, elem_id="nyra-main"):
                banner = gr.Markdown(
                    banner_text(default_lang), elem_id="nyra-banner"
                )
                hero = gr.HTML(hero_html(default_lang))

                with gr.Row(elem_id="nyra-chips"):
                    chip_btns = [
                        gr.Button(
                            t0["chips"][i][0],
                            elem_classes=["nyra-chip"],
                            scale=1,
                        )
                        for i in range(4)
                    ]

                chatbot = make_chatbot()

                # Single-line input only — multi-line auto-resize causes HF iframe grow
                with gr.Row(elem_id="nyra-input-row"):
                    msg = gr.Textbox(
                        elem_id="nyra-input",
                        placeholder=t0["placeholder"],
                        show_label=False,
                        lines=1,
                        max_lines=1,
                        scale=6,
                        container=False,
                        autofocus=False,
                        interactive=True,
                        elem_classes=["nyra-fixed-input"],
                    )
                    send_btn = gr.Button(
                        t0["send"],
                        elem_id="nyra-send",
                        variant="primary",
                        scale=0,
                        min_width=88,
                    )

                # Always-on pin script (Gradio js= is unreliable on Spaces iframe)
                gr.HTML(
                    """
<script>
(function(){
  if (window.__nyraPinInstalled) return;
  window.__nyraPinInstalled = true;
  window.iFrameResizer = { autoResize: false, sizeWidth: false };
  function pin(){
    var list = document.querySelectorAll('#nyra-input textarea, textarea[data-testid="textbox"]');
    for (var i=0;i<list.length;i++){
      var el = list[i];
      el.setAttribute('rows','1');
      el.style.setProperty('height','40px','important');
      el.style.setProperty('min-height','40px','important');
      el.style.setProperty('max-height','40px','important');
      el.style.setProperty('resize','none','important');
      el.style.setProperty('overflow','hidden','important');
      el.style.setProperty('line-height','40px','important');
      el.oninput = function(e){
        e.target.style.setProperty('height','40px','important');
      };
    }
    var row = document.getElementById('nyra-input-row');
    if(row){
      row.style.setProperty('height','56px','important');
      row.style.setProperty('max-height','56px','important');
      row.style.setProperty('min-height','56px','important');
      row.style.setProperty('overflow','hidden','important');
    }
    var nodes = document.querySelectorAll('.gradio-container, main, .wrap, .contain, .fillable');
    for (var j=0;j<nodes.length;j++){
      nodes[j].style.setProperty('min-height','0','important');
      nodes[j].style.setProperty('height','auto','important');
      nodes[j].style.setProperty('flex-grow','0','important');
    }
    document.documentElement.style.setProperty('height','auto','important');
    document.body.style.setProperty('height','auto','important');
    document.body.style.setProperty('min-height','0','important');
  }
  setInterval(pin, 100);
  document.addEventListener('DOMContentLoaded', pin);
  window.addEventListener('load', pin);
  pin();
})();
</script>
                    """,
                    visible=True,
                )

        inputs = [msg, chatbot, temperature, max_tokens, thinking, lang_state]
        outputs = [chatbot, msg]

        send_btn.click(respond, inputs=inputs, outputs=outputs)
        msg.submit(respond, inputs=inputs, outputs=outputs)

        def make_chip(idx: int):
            def _run(history, temp, max_tok, th, lang):
                yield from on_chip(idx, history, temp, max_tok, th, lang)

            return _run

        for i, btn in enumerate(chip_btns):
            btn.click(
                make_chip(i),
                inputs=[chatbot, temperature, max_tokens, thinking, lang_state],
                outputs=outputs,
            )

        new_btn.click(clear_chat, outputs=[chatbot, msg])
        clear_btn.click(clear_chat, outputs=[chatbot, msg])

        lang_btn.click(
            toggle_lang,
            inputs=[lang_state],
            outputs=[
                lang_state,
                brand,
                hero,
                msg,
                send_btn,
                new_btn,
                clear_btn,
                lang_btn,
                banner,
                *chip_btns,
            ],
        )

    return demo


demo = build_app()


def _launch_kwargs() -> dict:
    kw = {
        "server_name": os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        "server_port": int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        "show_error": True,
    }
    try:
        params = inspect.signature(gr.Blocks.launch).parameters
    except (TypeError, ValueError):
        params = {}
    for key, value in (
        ("theme", build_theme()),
        ("css", load_css()),
        ("js", load_js()),
        ("head", HEAD),
    ):
        if key in params:
            kw[key] = value
    return kw


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch(**_launch_kwargs())