import gradio as gr
from gradio import Server
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ─────────────────────────────────────────────
# 1. MODEL SETUP
# ─────────────────────────────────────────────
MODEL_ID = "huihui-ai/Huihui-Qwen3.5-9B-abliterated"

print(f"Loading tokenizer: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

print(f"Loading model: {MODEL_ID}")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",       # ZeroGPU manages CUDA device
    trust_remote_code=True,  # Needed for Qwen-based custom archs
)
model.eval()

# ─────────────────────────────────────────────
# 1b. THINKING ENFORCEMENT
#
#  This model isn't guaranteed to emit <think>...</think> on its own
#  (especially as an abliterated finetune). To make it 100% reliable:
#
#    1. Try passing enable_thinking=True to apply_chat_template
#       (native support on Qwen3-style templates, if present).
#    2. ALWAYS force-append "<think>\n" onto the templated prompt
#       so generation is physically forced to begin inside a think
#       block, regardless of whether enable_thinking worked.
#    3. Append a short mandatory instruction onto whatever system
#       prompt is passed in, telling the model to close the tag.
#    4. Manually re-prepend "<think>\n" onto the decoded output,
#       since it was part of the forced prompt and gets stripped
#       out by skip_prompt / prompt-slicing.
# ─────────────────────────────────────────────
THINK_INSTRUCTION = (
    "\n\nAlways reason through the problem step by step inside <think> "
    "and </think> tags first. After the closing </think> tag, give your "
    "final answer. Never skip the opening or closing think tags."
)


def build_prompt(prompt: str, system_prompt: str) -> str:
    full_system = (system_prompt or "You are a helpful assistant.") + THINK_INSTRUCTION

    messages = [
        {"role": "system", "content": full_system},
        {"role": "user", "content": prompt},
    ]

    try:
        # Native Qwen3-style thinking toggle, if the tokenizer supports it
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=True,
        )
    except TypeError:
        # Tokenizer/template doesn't accept enable_thinking — fall back
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

    # Force the assistant turn to start inside a <think> block,
    # guaranteeing the tag appears no matter what the model would
    # have done on its own.
    if not text.rstrip().endswith("<think>"):
        text = text + "<think>\n"

    return text


# ─────────────────────────────────────────────
# 2. INFERENCE FUNCTION (ZeroGPU decorated)
# ─────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate(
    prompt: str,
    system_prompt: str = "You are a helpful assistant.",
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
    do_sample: bool = True,
) -> str:
    """
    Generate a text response from the model.
    Response is guaranteed to start with a <think> block.
    """
    text = build_prompt(prompt, system_prompt)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=do_sample,
            pad_token_id=tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    # Re-add the <think> tag we force-injected into the prompt —
    # it was stripped out because it's technically part of the input.
    response = "<think>\n" + response
    return response


# ─────────────────────────────────────────────
# 3. STREAMING INFERENCE (SSE / token-by-token)
# ─────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate_stream(
    prompt: str,
    system_prompt: str = "You are a helpful assistant.",
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """
    Stream a text response token-by-token via SSE.
    Guaranteed to start with a <think> block.
    """
    from transformers import TextIteratorStreamer
    from threading import Thread

    text = build_prompt(prompt, system_prompt)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        streamer=streamer,
        pad_token_id=tokenizer.eos_token_id,
    )

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    # Re-inject the forced <think> prefix as the very first chunk,
    # since it was part of the prompt and streamer.skip_prompt=True
    # will not emit it on its own.
    partial = "<think>\n"
    yield partial

    for new_text in streamer:
        partial += new_text
        yield partial


# ─────────────────────────────────────────────
# 4. gr.Server — REST API + OPTIONAL SWAGGER UI
# ─────────────────────────────────────────────
app = Server(
    title="Mira-1-Large API",
    summary="ZeroGPU-backed REST API for Smilyai-labs/Mira-1-large (Qwen arch)",
    version="1.0.0",
)

app.api(generate,        name="generate")
app.api(generate_stream, name="generate_stream")


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID}


# ─────────────────────────────────────────────
# 5. LAUNCH
# ─────────────────────────────────────────────
app.launch()