import threading
import time

import gradio as gr
import spaces
import torch

import mfn_hf  # shared look + usage analytics; import before @spaces.GPU

from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer

import spark  # local copy of the model's remote code, ported to transformers 5; registers spark2_5

MODEL_ID = "XHToken/Spark-X2.6-4B"
REVISION = "66e3ae7bee54e0a523d692cb8ef841f35409807b"  # the snapshot the local code was ported from

# Sampling from the model card / generation_config.json: temperature 1.0, top_p 0.95, top_k off.
DEFAULT_TEMPERATURE = 1.0
DEFAULT_TOP_P = 0.95
DEFAULT_MAX_NEW_TOKENS = 4096
MAX_NEW_TOKENS = 16384
MAX_INPUT_TOKENS = 8192  # older turns are dropped past this

THINKING = ["None", "Low", "High"]
DEFAULT_THINKING = "Low"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    dtype=torch.bfloat16,
    attn_implementation="eager",  # the model code only implements eager attention
).to("cuda")
model.eval()

EOS = tokenizer.eos_token  # <｜end▁of▁sentence｜>
SPECIALS = [t for t in {EOS, tokenizer.bos_token, tokenizer.pad_token} if t]


def _text(content):
    """Plain text from a Gradio message content (str or list of blocks)."""
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        return content.get("text") or ""
    if isinstance(content, (list, tuple)):
        return "".join(_text(c) for c in content)
    return ""


def _turns(history, message):
    """Chatbot history -> chat-template turns. Thinking blocks are display-only and are left out."""
    turns = []
    for m in history or []:
        meta = m.get("metadata") or {}
        if meta.get("title"):
            continue
        role = m.get("role")
        text = _text(m.get("content")).strip()
        if role not in ("user", "assistant") or not text:
            continue
        if turns and turns[-1]["role"] == role:
            turns[-1]["content"] += "\n\n" + text
        else:
            turns.append({"role": role, "content": text})
    # The UI shows the new message before generation starts; don't count it twice.
    if turns and turns[-1]["role"] == "user" and turns[-1]["content"] == message:
        turns.pop()
    return turns + [{"role": "user", "content": message}]


def _template_kwargs(thinking):
    level = (thinking or DEFAULT_THINKING).lower()
    if level == "none":
        return {"enable_thinking": False}
    return {"enable_thinking": True, "reasoning_effort": level}


def build_inputs(message, history, thinking, system_prompt):
    turns = _turns(history, message)
    system = [{"role": "system", "content": system_prompt.strip()}] if (system_prompt or "").strip() else []
    kw = _template_kwargs(thinking)
    while True:
        ids = tokenizer.apply_chat_template(
            system + turns, add_generation_prompt=True, return_tensors="pt", return_dict=True, **kw
        )["input_ids"]
        if ids.shape[-1] <= MAX_INPUT_TOKENS or len(turns) <= 1:
            break
        turns = turns[2:] if len(turns) > 2 else turns[1:]  # drop the oldest exchange
    if ids.shape[-1] > MAX_INPUT_TOKENS:
        raise gr.Error("That message is too long for this demo. Try a shorter one.")
    return ids


class _Stop(StoppingCriteria):
    def __init__(self, event):
        self.event = event

    def __call__(self, input_ids, scores, **kwargs):
        return self.event.is_set()


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


def gpu_duration(message, history, thinking="Low", system_prompt="", temperature=1.0, top_p=0.95, max_new_tokens=DEFAULT_MAX_NEW_TOKENS):
    return int(min(30 + int(max_new_tokens) / 35, 480))


@spaces.GPU(duration=gpu_duration)
def chat(
    message,
    history,
    thinking=DEFAULT_THINKING,
    system_prompt="",
    temperature=DEFAULT_TEMPERATURE,
    top_p=DEFAULT_TOP_P,
    max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
):
    message = _text(message).strip()
    if not message:
        raise gr.Error("Type a message first.")
    history = list(history or [])
    ids = build_inputs(message, history, thinking, system_prompt).to("cuda")
    base = history[:-1] if history and history[-1].get("role") == "user" and _text(history[-1].get("content")).strip() == message else history
    base = base + [{"role": "user", "content": message}]

    thinks = (thinking or DEFAULT_THINKING).lower() != "none"
    stop = threading.Event()
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False, timeout=120)
    temperature = float(temperature)
    gen = dict(
        input_ids=ids,
        attention_mask=torch.ones_like(ids),
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        stopping_criteria=StoppingCriteriaList([_Stop(stop)]),
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    if temperature > 0:
        gen.update(do_sample=True, temperature=temperature, top_p=float(top_p), top_k=0)
    else:
        gen.update(do_sample=False, temperature=None, top_p=None, top_k=None)

    error = []

    def run():
        try:
            with torch.inference_mode():
                model.generate(**gen)
        except Exception as e:  # surfaced below without echoing inputs
            error.append(e)
            streamer.on_finalized_text("", stream_end=True)

    t0 = time.time()
    worker = threading.Thread(target=run, daemon=True)
    worker.start()

    raw = ""
    t_think = None
    try:
        for piece in streamer:
            raw += piece
            text = _clean(raw)
            if thinks:
                if "</think>" in text:
                    reasoning, answer = text.split("</think>", 1)
                    if t_think is None:
                        t_think = time.time() - t0
                    status, title = "done", f"Thought for {t_think:.1f} s"
                else:
                    reasoning, answer = text, ""
                    status, title = "pending", "Thinking"
                msgs = [{"role": "assistant", "content": reasoning.strip(), "metadata": {"title": title, "status": status}}]
                if answer.strip():
                    msgs.append({"role": "assistant", "content": answer.strip()})
            else:
                msgs = [{"role": "assistant", "content": text.replace("</think>", "").strip()}]
            yield base + msgs, f"Took {time.time() - t0:.1f} s"
    finally:
        stop.set()
    worker.join(timeout=10)
    if error:
        raise gr.Error("Generation failed. Please try again.")

    text = _clean(raw)
    if thinks:
        if "</think>" in text:
            reasoning, answer = text.split("</think>", 1)
            msgs = [{"role": "assistant", "content": reasoning.strip(), "metadata": {"title": f"Thought for {t_think or time.time() - t0:.1f} s", "status": "done"}}]
            if not reasoning.strip():
                msgs = []
        else:  # ran out of tokens while thinking
            reasoning, answer = text, ""
            msgs = [{"role": "assistant", "content": reasoning.strip(), "metadata": {"title": "Thinking (stopped at the token limit)", "status": "done"}}]
        if answer.strip():
            msgs.append({"role": "assistant", "content": answer.strip()})
    else:
        msgs = [{"role": "assistant", "content": text.replace("</think>", "").strip()}]
    yield base + msgs, f"Took {time.time() - t0:.1f} s"


def show_user(message, history):
    message = _text(message).strip()
    if not message:
        return gr.update(), message, history
    return "", message, list(history or []) + [{"role": "user", "content": message}]


EXAMPLES = [
    "Write a Python function that merges overlapping intervals, with a few tests.",
    "A bat and a ball cost $1.10 together. The bat costs $1.00 more than the ball. How much is the ball? Explain.",
    "Explain the difference between a process and a thread, with a short example in Rust.",
    "用三句话解释什么是量子纠缠，让中学生也能听懂。",
    "鸡兔同笼，共有35个头，94只脚。鸡和兔各有多少只？",
    "Plan a three-day trip to Hefei for someone who loves food and museums.",
]

with gr.Blocks(title="Spark-X2.6 4B") as demo:
    mfn_hf.header(
        "Spark-X2.6 4B",
        model=MODEL_ID,
        description="Chat with XHToken's 4B English and Chinese model. Pick how much it thinks before answering.",
        flow_fn="/chat",
    )
    chatbot = gr.Chatbot(
        height=560,
        show_label=False,
        placeholder="Ask anything, in English or Chinese.",
    )
    with gr.Row(equal_height=True):
        msg = gr.Textbox(
            show_label=False,
            placeholder="Message Spark-X2.6",
            lines=1,
            max_lines=8,
            scale=7,
            container=False,
        )
        send = gr.Button("Send", variant="primary", scale=1, min_width=80)
        stop_btn = gr.Button("Stop", scale=1, min_width=80)
    with gr.Row(equal_height=True):
        thinking = gr.Radio(THINKING, value=DEFAULT_THINKING, label="Thinking", scale=3)
        clear = gr.Button("New chat", scale=1, min_width=100)
    took = mfn_hf.took()

    with gr.Accordion("Settings", open=False):
        system_prompt = gr.Textbox(label="System prompt", placeholder="Optional, e.g. Answer in one paragraph.", lines=2)
        with gr.Row():
            temperature = gr.Slider(0.0, 1.5, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature")
            top_p = gr.Slider(0.05, 1.0, value=DEFAULT_TOP_P, step=0.05, label="Top-p")
            max_new_tokens = gr.Slider(256, MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS, step=256, label="Max new tokens")

    gr.Examples(EXAMPLES, [msg], label="Examples")

    gr.Markdown(
        "Spark-X2.6-4B is by [XHToken](https://huggingface.co/XHToken) (Apache-2.0). "
        "Thinking levels follow the model's chat template: None skips reasoning, Low and High set its thinking effort."
    )

    pending = gr.State("")
    settings = [thinking, system_prompt, temperature, top_p, max_new_tokens]

    events = []
    for trigger in (msg.submit, send.click):
        ev = trigger(show_user, [msg, chatbot], [msg, pending, chatbot], queue=False, api_name=False).then(
            chat, [pending, chatbot, *settings], [chatbot, took], api_name=False
        )
        events.append(ev)
    stop_btn.click(None, None, None, cancels=events, api_name=False)
    clear.click(lambda: ([], "", ""), None, [chatbot, msg, took], queue=False, api_name=False)

    # API endpoint: chat(message, history, ...) -> (history, took)
    api_in = gr.Textbox(visible=False)
    api_btn = gr.Button(visible=False)
    api_btn.click(chat, [api_in, chatbot, *settings], [chatbot, took], api_name="chat")

    mfn_hf.footer()


if __name__ == "__main__":
    demo.launch()  # mfn_hf patches launch: theme, CSS, queue and analytics