"""Salience Chat — ZeroGPU backend + static React frontend.

Serves vectionlabs/Salience-27B-R5. The Gradio app (mounted at /gradio) only
exists to expose the streaming `chat` API consumed by the React UI in ./static.

TWO THINGS CHANGED WITH R5, AND BOTH ARE STRUCTURAL.

1. NF4, not bf16. R5 is 27.78B against Nano's 9B: bf16 would be 51.7 GiB and
   there is no ZeroGPU tier that holds it. NF4 quantises the body to ~11.8 GiB
   and leaves the embedding and lm_head in bf16 (bitsandbytes skips nn.Embedding
   and the output head by design), landing around 16.5 GiB.

2. Reasoning is NATIVE now, so the prompt scaffolding is gone. The old code
   forced thinking by pasting instructions into a system prompt -- "reason step
   by step", "write a rough draft, check it" -- and prefilling an opening
   <think> tag, because Nano's template had no thinking support.

   R5's template does. It takes a real `reasoning_effort` parameter and emits
   the opening tag itself. Keeping the old prompts would be actively harmful:
   telling a natively-reasoning model to reason makes it PERFORM reasoning for
   the reader instead of doing it, and the two compound into narrated
   pseudo-thought. All of that text is deleted rather than adapted.
"""

import threading

import gradio as gr
import spaces
import torch
from starlette.responses import FileResponse
from starlette.routing import Mount, Route
from starlette.staticfiles import StaticFiles
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)

MODEL_ID = "vectionlabs/Salience-27B-R5"

# Sampling recommended for this family's reasoning mode.
REASONING_PARAMS = {"temperature": 1.0, "top_p": 0.95, "top_k": 20}

QUANT = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    # bf16 compute. ZeroGPU is Ampere or newer, so this is safe here -- it is
    # NOT on Turing (T4), where bnb needs float16 instead.
    bnb_4bit_compute_dtype=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    quantization_config=QUANT,
    dtype=torch.bfloat16,
    device_map="cuda",
)

# The frontend ships four pills; the template takes three values. Rather than
# invent a fourth setting that does nothing, High and xHigh both send xhigh and
# differ in how long they are allowed to run. Inventing seven names for three
# real settings would be theatre.
EFFORT = {
    "low":    {"kw": {"reasoning_effort": "low"},    "max_tok": 2048},
    "medium": {"kw": {"reasoning_effort": "medium"}, "max_tok": 4096},
    "high":   {"kw": {"reasoning_effort": "xhigh"},  "max_tok": 8192},
    "xhigh":  {"kw": {"reasoning_effort": "xhigh"},  "max_tok": 16384},
}


def _to_messages(history: list | None, message: str) -> list[dict]:
    """No system prompt. R5's template writes its own instructions from
       reasoning_effort, and a hand-rolled one would fight it."""
    messages = []
    for turn in history or []:
        content = turn["content"]
        if turn["role"] == "assistant":
            # Drop prior reasoning, exactly as agent harnesses do: it only
            # inflates the context, and the model does not need to re-read how
            # it reached somewhere it already reached.
            content = content.split("</think>")[-1].strip()
        messages.append(
            {"role": turn["role"], "content": [{"type": "text", "text": content}]}
        )
    messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
    return messages


@spaces.GPU(duration=120)
def chat(message: str, history: list | None, effort: str):
    cfg = EFFORT.get(effort, EFFORT["medium"])

    text = processor.apply_chat_template(
        _to_messages(history, message),
        add_generation_prompt=True,
        tokenize=False,
        **cfg["kw"],
    )
    inputs = processor(text=[text], return_tensors="pt").to(model.device)

    # <think>/</think> are special tokens: skip_special_tokens would strip them
    # and the frontend could no longer split reasoning from answer.
    streamer = TextIteratorStreamer(
        processor.tokenizer, skip_prompt=True, skip_special_tokens=False
    )
    thread = threading.Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=cfg["max_tok"],
            do_sample=True,
            **REASONING_PARAMS,
        ),
    )
    thread.start()

    # THE PARSING TRAP ON THIS FAMILY, HANDLED HERE SO THE CLIENT NEVER SEES IT.
    # The template emits the OPENING <think> as part of the generation prompt,
    # and skip_prompt=True means the stream starts *inside* the block -- the
    # completion carries only the closing tag. The frontend looks for a matching
    # pair, so without this line it would find none, report zero thinking, and
    # render the entire chain of thought as the answer.
    opened = text.rstrip().endswith("<think>")
    accumulated = "<think>\n" if opened else ""
    for chunk in streamer:
        accumulated += chunk
        yield accumulated.replace("<|im_end|>", "").replace("<|endoftext|>", "")


with gr.Blocks() as demo:
    message_in = gr.Textbox(label="message")
    history_in = gr.JSON(label="history")
    effort_in = gr.Textbox(label="effort", value="medium")
    response_out = gr.Textbox(label="response")
    send = gr.Button("send")
    send.click(
        chat,
        inputs=[message_in, history_in, effort_in],
        outputs=response_out,
        api_name="chat",
    )

demo.queue(default_concurrency_limit=2)

# ZeroGPU only registers @spaces.GPU functions from inside demo.launch(), so
# Gradio must own the server. The React SPA is injected into its router
# afterwards: / serves the SPA shell, /app/* its assets. Gradio's own routes
# (/config, /gradio_api/...) keep serving the API.
app, _, _ = demo.launch(prevent_thread_lock=True, ssr_mode=False, show_error=True)


async def spa_index(request):
    return FileResponse("static/index.html")


app.router.routes.insert(0, Route("/", spa_index, methods=["GET", "HEAD"]))
app.router.routes.insert(
    0, Mount("/app", app=StaticFiles(directory="static"), name="spa-assets")
)

demo.block_thread()