import os, json
import spaces
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen3-0.6B")
TRC = os.environ.get("TRUST_REMOTE_CODE", "0") == "1"
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=TRC)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype="auto", trust_remote_code=TRC
).to("cuda")

THINK = os.environ.get("ENABLE_THINKING", "0") == "1"

# Vendor-recommended sampling for the served model, set as Space variables
# (each Space serves ONE model, so its generation defaults live here).
# e.g. Qwen thinking: GEN_TOP_P=0.95 · LFM2.5: GEN_MIN_P=0.15 GEN_REP_PENALTY=1.05
GEN_KW = {}
if os.environ.get("GEN_TOP_P"):
    GEN_KW["top_p"] = float(os.environ["GEN_TOP_P"])
if os.environ.get("GEN_MIN_P"):
    GEN_KW["min_p"] = float(os.environ["GEN_MIN_P"])
if os.environ.get("GEN_REP_PENALTY"):
    GEN_KW["repetition_penalty"] = float(os.environ["GEN_REP_PENALTY"])

def _render(messages, tools, **kw):
    # enable_thinking is Qwen/SmolLM-specific; fall back for families without it
    if tools:
        kw["tools"] = tools
    try:
        return tok.apply_chat_template(messages, enable_thinking=THINK, **kw)
    except TypeError:
        return tok.apply_chat_template(messages, **kw)

def build_inputs(messages, tools):
    kw = dict(add_generation_prompt=True, return_tensors="pt", return_dict=True)
    if tools:
        # Some templates silently ignore the standard `tools=` kwarg and instead
        # read tools off the system MESSAGE (e.g. Phi-4-mini renders a
        # <|tool|>...</|tool|> block from message['tools']). Detect by comparing
        # the render WITH tools against the render WITHOUT: if identical, the
        # template dropped them — re-inject via the system message. (Never match
        # on tool NAMES: the user's own message text may mention them.)
        with_tools = _render(messages, tools, add_generation_prompt=True, tokenize=False)
        without_tools = _render(messages, None, add_generation_prompt=True, tokenize=False)
        if with_tools == without_tools:
            tjson = json.dumps([t.get("function", t) for t in tools])
            msgs = [dict(m) for m in messages]
            if msgs and msgs[0].get("role") == "system":
                msgs[0]["tools"] = tjson
            else:
                msgs.insert(0, {"role": "system",
                                "content": "You are a helpful assistant with access to the following tools.",
                                "tools": tjson})
            return _render(msgs, None, **kw)
    return _render(messages, tools, **kw)

@spaces.GPU(duration=90)
def run(messages_json, tools_json="", max_new_tokens=512, temperature=0.7):
    messages = json.loads(messages_json)
    # OpenAI-standard assistant tool-call messages carry content: null; templates
    # that concatenate content raw (Phi-4-mini) crash on None. Coerce to "".
    for m in messages:
        if m.get("content") is None:
            m["content"] = ""
    tools = json.loads(tools_json) if (tools_json and tools_json.strip()) else None
    inputs = build_inputs(messages, tools).to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=int(max_new_tokens),
            do_sample=(float(temperature) > 0),
            temperature=max(float(temperature), 0.01),
            pad_token_id=(tok.eos_token_id or tok.pad_token_id),
            **GEN_KW,
        )
    n = inputs["input_ids"].shape[1]
    return f"[MODEL={MODEL_ID}]\n" + tok.decode(out[0][n:], skip_special_tokens=False)

demo = gr.Interface(
    fn=run,
    inputs=[gr.Textbox(label="messages_json"), gr.Textbox(label="tools_json"),
            gr.Number(label="max_new_tokens", value=512), gr.Number(label="temperature", value=0.7)],
    outputs=gr.Textbox(label="output"),
    api_name="run",
)
demo.launch()