import os
from openai import OpenAI
import gradio as gr

BASE_URL = "https://modelflare.dev/v1"
DEFAULT_MODELS = [
    "gpt-4o-mini",
    "gpt-4o",
    "claude-sonnet-4-5",
    "grok-4",
]


def get_client() -> OpenAI:
    api_key = os.environ.get("MODELFLARE_API_KEY", "").strip()
    if not api_key:
        raise gr.Error(
            "缺少 Space Secret：MODELFLARE_API_KEY。"
            "请在 Space Settings → Variables and secrets 里添加后重启。"
        )
    return OpenAI(api_key=api_key, base_url=BASE_URL)


def chat(message: str, history: list, model: str, custom_model: str):
    model_id = (custom_model or "").strip() or (model or "").strip() or DEFAULT_MODELS[0]
    client = get_client()
    messages = []
    for turn in history:
        if isinstance(turn, dict):
            role, content = turn.get("role"), turn.get("content")
            if role in ("user", "assistant") and content:
                messages.append({"role": role, "content": content})
        elif isinstance(turn, (list, tuple)) and len(turn) == 2:
            u, a = turn
            if u:
                messages.append({"role": "user", "content": u})
            if a:
                messages.append({"role": "assistant", "content": a})
    messages.append({"role": "user", "content": message})

    stream = client.chat.completions.create(
        model=model_id,
        messages=messages,
        stream=True,
    )
    partial = ""
    for chunk in stream:
        delta = chunk.choices[0].delta.content or ""
        if delta:
            partial += delta
            yield partial


with gr.Blocks(title="Modelflare Chat Demo") as demo:
    gr.Markdown(
        """
# Modelflare Chat Demo
通过 [Modelflare](https://modelflare.dev/) 统一网关试用 OpenAI 兼容对话。
文档：[docs.modelflare.dev](https://docs.modelflare.dev/) · Base URL：`https://modelflare.dev/v1`
"""
    )
    with gr.Row():
        model = gr.Dropdown(
            choices=DEFAULT_MODELS,
            value=DEFAULT_MODELS[0],
            label="Model",
            scale=2,
        )
        custom_model = gr.Textbox(
            label="Custom model id (可选，填写则覆盖上方)",
            placeholder="例如你的分组里实际可用的模型名",
            scale=3,
        )
    chatbot = gr.ChatInterface(
        fn=chat,
        additional_inputs=[model, custom_model],
        examples=[
            ["用一句话介绍 Modelflare"],
            ["Write a haiku about API gateways"],
        ],
    )

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