import os
import gradio as gr
from huggingface_hub import InferenceClient

MODEL_ID = "MiniMaxAI/MiniMax-M2.5"
SYSTEM_PROMPT = (
    "You are a helpful assistant. "
    "Your name is MiniMax-M2.5 and is built by MiniMax."
)

client = InferenceClient(
    provider="novita",
    api_key=os.environ.get("HF_TOKEN"),
)


def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}]

    # history در نسخه‌های جدید Gradio به صورت لیست دیکشنری است
    for msg in history:
        messages.append(
            {
                "role": msg["role"],
                "content": msg["content"],
            }
        )

    messages.append({"role": "user", "content": message})

    response = ""
    for chunk in client.chat_completion(
        model=MODEL_ID,
        messages=messages,
        max_tokens=int(max_tokens),
        temperature=float(temperature),
        top_p=float(top_p),
        stream=True,
    ):
        if chunk.choices:
            delta = chunk.choices[0].delta
            if delta and delta.content:
                response += delta.content
                yield response


demo = gr.ChatInterface(
    fn=respond,
    type="messages",  # مهم
    title="MiniMax M2.5 Chat",
    description=(
        "Chat with MiniMax M2.5 — "
        "a 230B MoE model (10B active) that is SOTA in coding, "
        "agentic tool use, and more."
    ),
    additional_inputs=[
        gr.Textbox(
            value=SYSTEM_PROMPT,
            label="System message",
        ),
        gr.Slider(
            minimum=1,
            maximum=4096,
            value=2048,
            step=1,
            label="Max new tokens",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.05,
            label="Temperature",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        ),
    ],
    examples=[
        ["Write a Python function to check if a number is prime."],
        ["Explain the difference between TCP and UDP in simple terms."],
        ["Help me write a bash script that monitors disk usage and sends an alert."],
    ],
    cache_examples=False,
)


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