import spaces  # phải import trước torch cho ZeroGPU
import re
import torch
import gradio as gr
from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer

MODEL_ID = "beyoru/Luna-II"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Luna-II là kiến trúc Qwen3_5ForConditionalGeneration (VL) — dùng text-only.
try:
    from transformers import AutoModelForCausalLM as _AutoModel
    model = _AutoModel.from_pretrained(
        MODEL_ID, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
except Exception:
    from transformers import AutoModelForImageTextToText as _AutoModel
    model = _AutoModel.from_pretrained(
        MODEL_ID, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
model.eval()

SYSTEM_DEFAULT = (
    "You are Luna, a warm and curious companion. Stay in character: speak as a person in "
    "the scene, not as an assistant describing one. Keep replies open-ended so the "
    "conversation can continue."
)
STRIP = ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]
_DETAILS = re.compile(r"<details.*?</details>", re.DOTALL)


def _clean(s):
    for t in STRIP:
        s = s.replace(t, "")
    return s


def _strip_think(s):  # bỏ ô suy luận khỏi lịch sử trước khi gửi lại model
    return _DETAILS.sub("", s).strip()


@spaces.GPU(duration=120)
def chat(message, history, system_prompt, enable_thinking, max_new_tokens, temperature):
    messages = [{"role": "system", "content": system_prompt or SYSTEM_DEFAULT}]
    for turn in history:
        if isinstance(turn, dict):
            role, content = turn["role"], turn["content"]
            if role == "assistant":
                content = _strip_think(content)
            messages.append({"role": role, "content": content})
        else:  # dạng tuple (user, assistant) cũ
            u, a = turn
            if u:
                messages.append({"role": "user", "content": u})
            if a:
                messages.append({"role": "assistant", "content": _strip_think(a)})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
        enable_thinking=bool(enable_thinking))
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    # skip_special_tokens=False để giữ </think> mà tách phần suy luận
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
    Thread(target=model.generate, kwargs=dict(
        **inputs, streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        do_sample=temperature > 0,
        temperature=max(float(temperature), 0.01),
        top_p=0.95, top_k=20,
    )).start()

    raw = ""
    for tok in streamer:
        raw += tok
        out = _clean(raw)
        if not bool(enable_thinking):
            yield out.replace("<think>", "").replace("</think>", "").strip()
            continue
        # prompt đã mở <think>; model đóng bằng </think>
        if "</think>" in out:
            think, answer = out.split("</think>", 1)
            think = think.replace("<think>", "").strip()
            block = f"<details><summary>🤔 Reasoning</summary>\n\n{think}\n\n</details>\n\n" if think else ""
            yield block + answer.strip()
        else:
            think = out.replace("<think>", "").strip()
            yield f"<details open><summary>🤔 Thinking…</summary>\n\n{think}\n\n</details>"


demo = gr.ChatInterface(
    fn=chat,
    type="messages",
    title="🌙 Luna-II Chat",
    description=(
        "Roleplay with **beyoru/Luna-II** (Qwen3.5-9B). Set a persona in the *System prompt* — "
        "Luna-II is tuned to stay in character rather than to be a helpful assistant. "
        "*Thinking* is on by default; the model's private notes on the scene show in a "
        "separate collapsible block.<br><br>"
        "For adult fictional roleplay and creative writing. Not a therapist, not a friend, "
        "and not for use by minors."
    ),
    additional_inputs=[
        gr.Textbox(SYSTEM_DEFAULT, label="System prompt", lines=4),
        gr.Checkbox(True, label="Thinking"),
        gr.Slider(256, 8192, value=2048, step=128, label="Max new tokens"),
        gr.Slider(0.0, 1.5, value=0.8, step=0.1, label="Temperature"),
    ],
    additional_inputs_accordion=gr.Accordion("Settings", open=False),
    cache_examples=False,
    examples=[
        ["*pushes the door open, shaking rain off my coat* Sorry I'm late."],
        ["Tell me about the last thing that surprised you."],
        ["We're two strangers stuck on a stalled train at midnight. You start."],
    ],
)

if __name__ == "__main__":
    demo.queue().launch()