"""A ZeroGPU Gradio code-assistant demo for XHToken/Spark-X2.5-4B.

Spark-X2.5-4B is a compact general-purpose model with strong coding and
reasoning ability. This Space wraps it in a streaming chat UI tuned for
programming tasks: write, explain, debug, refactor, test, and translate code,
with an optional collapsible reasoning (<think>) trace.
"""

import os
import threading
import time

import spaces  # noqa: F401  (must be imported before torch / transformers)

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


MODEL_ID = "XHToken/Spark-X2.5-4B"
MAX_CONTEXT_TOKENS = 32_768
MIN_NEW_TOKENS = 256
MAX_NEW_TOKENS = 4096

# The model ships a custom `spark2_5` architecture via `trust_remote_code`.
# Its custom attention path currently requires the eager implementation.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
    attn_implementation="eager",
).to("cuda").eval()

BASE_SYSTEM_PROMPT = (
    "You are Spark Code, an expert programming assistant. "
    "Give correct, runnable code in fenced Markdown blocks with the language tag. "
    "Prefer clear, idiomatic, production-quality solutions. "
    "Explain briefly, call out edge cases, and state assumptions when the request is ambiguous."
)

TASK_PRESETS = {
    "General coding": "",
    "Write new code": "Focus on writing a complete, self-contained implementation.",
    "Explain code": "Explain what the code does step by step, then summarize the key ideas.",
    "Debug & fix": "Identify the bug(s), explain the root cause, and provide a corrected version.",
    "Refactor": "Improve readability, structure, and performance without changing behavior.",
    "Write tests": "Produce thorough unit tests, including edge cases and failure modes.",
    "Translate language": "Port the code to the language the user requests, preserving behavior and idioms.",
}


def split_reasoning(text: str, enable_thinking: bool) -> tuple[str, str]:
    """Split a partial/complete generation into (reasoning, answer).

    With thinking enabled the chat template appends `<think>` to the prompt, so
    the generated text is pure reasoning until it emits `</think>` and then the
    answer. With thinking disabled the model answers directly.
    """
    if not enable_thinking:
        return "", text.strip()
    if "</think>" in text:
        reasoning, answer = text.split("</think>", 1)
        return reasoning.replace("<think>", "").strip(), answer.strip()
    return text.replace("<think>", "").strip(), ""


def history_to_messages(history: list[object] | None) -> list[dict[str, str]]:
    """Convert Gradio's chat history into Spark chat-template messages."""
    messages: list[dict[str, str]] = []
    for item in history or []:
        if isinstance(item, dict):
            content = item.get("content", "")
            if isinstance(content, list):
                content = "".join(
                    block.get("text", "") for block in content if isinstance(block, dict)
                )
            messages.append({"role": item.get("role", "user"), "content": str(content)})
        else:
            user_text, assistant_text = item
            messages.extend(
                [
                    {"role": "user", "content": user_text},
                    {"role": "assistant", "content": assistant_text},
                ]
            )
    return messages


def _estimate_duration(
    message=None,
    history=None,
    system_prompt=None,
    enable_thinking=None,
    max_new_tokens=2048,
    temperature=None,
    top_p=None,
    *args,
    **kwargs,
):
    """ZeroGPU duration callable: scale the reservation with the token budget."""
    try:
        budget = int(max_new_tokens)
    except (TypeError, ValueError):
        budget = 1024
    return min(240, 40 + budget // 8)


@spaces.GPU(duration=_estimate_duration)
def respond(
    message: str,
    history: list[object] | None,
    system_prompt: str,
    task: str,
    enable_thinking: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
):
    """Stream a coding answer from Spark-X2.5-4B, revealing its reasoning trace."""
    history = history or []
    if not message or not message.strip():
        yield history, ""
        return

    task_hint = TASK_PRESETS.get(task, "")
    system_parts = [system_prompt.strip()] if system_prompt and system_prompt.strip() else []
    if task_hint:
        system_parts.append(task_hint)
    system_text = "\n\n".join(system_parts)

    history_messages = history_to_messages(history)
    messages = []
    if system_text:
        messages.append({"role": "system", "content": system_text})
    messages.extend(history_messages)
    messages.append({"role": "user", "content": message.strip()})

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

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        do_sample=temperature > 0,
        temperature=max(float(temperature), 1e-5),
        top_p=float(top_p),
        # Spark's generation config uses -1 for disabled top-k; Transformers
        # expects 0 for the same behavior.
        top_k=0,
        use_cache=True,
        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs, daemon=True)
    thread.start()

    base = history_messages + [{"role": "user", "content": message.strip()}]
    completion = ""
    last_emit = 0.0
    for new_text in streamer:
        completion += new_text
        reasoning, answer = split_reasoning(completion, bool(enable_thinking))
        display = answer or ("_Thinking…_" if reasoning else "")
        now = time.perf_counter()
        if now - last_emit >= 0.1:
            last_emit = now
            yield base + [{"role": "assistant", "content": display}], (
                reasoning if enable_thinking else ""
            )
    thread.join()

    reasoning, answer = split_reasoning(completion, bool(enable_thinking))
    if answer:
        final_text = answer
    elif reasoning:
        final_text = (
            reasoning
            + "\n\n> ⚠️ The model stopped before finishing its reasoning — "
            "increase **Max new tokens** and try again."
        )
    else:
        final_text = ""
    yield base + [{"role": "assistant", "content": final_text}], (
        reasoning if enable_thinking else ""
    )


with gr.Blocks(title="Spark-X2.5-4B Code Assistant") as demo:
    gr.Markdown(
        "# 💻 Spark-X2.5-4B Code Assistant\n"
        "A streaming coding assistant built on "
        "[XHToken/Spark-X2.5-4B](https://huggingface.co/XHToken/Spark-X2.5-4B) — "
        "a 4B model with strong coding, reasoning, and agentic ability. "
        "Write, explain, debug, refactor, test, and translate code."
    )

    chatbot = gr.Chatbot(
        height=520,
        label="Conversation",
        render_markdown=True,
    )
    with gr.Row():
        message = gr.Textbox(
            label="Your request",
            placeholder="e.g. Write a Python LRU cache with O(1) get/put and unit tests…",
            lines=3,
            scale=8,
        )
        send = gr.Button("Send", variant="primary", scale=1)

    with gr.Accordion("Task & generation settings", open=False):
        with gr.Row():
            task = gr.Dropdown(
                choices=list(TASK_PRESETS.keys()),
                value="General coding",
                label="Task preset",
            )
            enable_thinking = gr.Checkbox(label="Show reasoning trace", value=True)
        system_prompt = gr.Textbox(
            label="System prompt", value=BASE_SYSTEM_PROMPT, lines=3
        )
        with gr.Row():
            max_new_tokens = gr.Slider(
                MIN_NEW_TOKENS, MAX_NEW_TOKENS, value=2048, step=256, label="Max new tokens"
            )
            temperature = gr.Slider(0, 1.5, value=1.0, step=0.05, label="Temperature")
            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
    reasoning = gr.Textbox(label="Reasoning trace", lines=8, visible=True)
    clear = gr.ClearButton([message, chatbot, reasoning], value="Clear conversation")

    inputs = [
        message,
        chatbot,
        system_prompt,
        task,
        enable_thinking,
        max_new_tokens,
        temperature,
        top_p,
    ]
    outputs = [chatbot, reasoning]
    send.click(respond, inputs=inputs, outputs=outputs).then(lambda: "", outputs=message)
    message.submit(respond, inputs=inputs, outputs=outputs).then(
        lambda: "", outputs=message
    )

    gr.Examples(
        examples=[
            ["Write a Python function that merges two sorted lists in O(n+m) and add pytest tests covering empty inputs and duplicates."],
            ["Explain what this does and its time complexity:\n\nfrom functools import lru_cache\n@lru_cache(maxsize=None)\ndef fib(n):\n    return n if n < 2 else fib(n-1) + fib(n-2)"],
            ["This async Python snippet sometimes hangs. Find the bug and fix it:\n\nasync def main():\n    results = [await fetch(u) for u in urls]\n    return results"],
            ["Refactor this JavaScript into a clean, tested ES module:\n\nfunction p(a){var r=[];for(var i=0;i<a.length;i++){if(a[i]%2==0)r.push(a[i]*a[i]);}return r;}"],
        ],
        inputs=[message],
        label="Try a coding example",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        mcp_server=True, theme=gr.themes.Soft()
    )