import gradio as gr
from huggingface_hub import InferenceClient


MODEL_ID = "openai/gpt-oss-20b"


def respond(
    message: str,
    history: list[dict],
    system_message: str,
    reasoning_effort: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    hf_token: gr.OAuthToken | None,
):
    if hf_token is None or not getattr(hf_token, "token", None):
        yield "请先点击左侧的“Sign in with Hugging Face”登录，然后再发送消息。"
        return

    client = InferenceClient(
        model=MODEL_ID,
        provider="auto",
        token=hf_token.token,
    )

    messages = []

    if system_message.strip():
        messages.append(
            {
                "role": "system",
                "content": system_message.strip(),
            }
        )

    # Gradio messages 模式下，history 通常是：
    # [{"role": "user", "content": "..."}, ...]
    for item in history:
        role = item.get("role")
        content = item.get("content")

        if role in {"user", "assistant", "system"} and isinstance(content, str):
            messages.append(
                {
                    "role": role,
                    "content": content,
                }
            )

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

    response_text = ""

    try:
        stream = client.chat_completion(
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            stream=True,

            # 转交给兼容 OpenAI 风格的推理后端。
            # 若当前 Provider 不接受该字段，删除 extra_body 即可。
            extra_body={
                "reasoning_effort": reasoning_effort,
            },
        )

        for chunk in stream:
            if not chunk.choices:
                continue

            delta = chunk.choices[0].delta
            token = getattr(delta, "content", None)

            if token:
                response_text += token
                yield response_text

        if not response_text:
            yield "模型没有返回可显示的文本。请稍后重试或更换推理参数。"

    except Exception as error:
        error_text = str(error)

        if "401" in error_text or "unauthorized" in error_text.lower():
            yield (
                "身份验证失败。请退出后重新登录 Hugging Face，"
                "并确认账户允许调用 Inference Providers。"
            )
        elif "402" in error_text or "payment" in error_text.lower():
            yield (
                "当前 Hugging Face 账户的推理额度不足，"
                "请检查 Inference Providers 余额或计费设置。"
            )
        elif "429" in error_text or "rate limit" in error_text.lower():
            yield "请求过于频繁或免费额度已达到限制，请稍后再试。"
        else:
            yield f"调用模型时发生错误：{error_text}"


with gr.Blocks(title="GPT-OSS 20B Chat") as demo:
    gr.Markdown(
        """
        # GPT-OSS 20B Chat

        使用 Hugging Face Inference Providers 调用 `openai/gpt-oss-20b`。

        请先登录 Hugging Face。推理请求将使用登录用户自己的 HF 账户和额度。
        """
    )

    with gr.Sidebar():
        gr.Markdown("### Hugging Face 账户")
        gr.LoginButton()

        system_message = gr.Textbox(
            value="You are a helpful and friendly assistant.",
            label="System message",
            lines=4,
        )

        reasoning_effort = gr.Radio(
            choices=["low", "medium", "high"],
            value="medium",
            label="Reasoning effort",
        )

        max_tokens = gr.Slider(
            minimum=64,
            maximum=8192,
            value=1024,
            step=64,
            label="Max output tokens",
        )

        temperature = gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=0.7,
            step=0.05,
            label="Temperature",
        )

        top_p = gr.Slider(
            minimum=0.05,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        )

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            system_message,
            reasoning_effort,
            max_tokens,
            temperature,
            top_p,
        ],
        examples=[
            ["请解释 MoE 模型中的总参数和激活参数有什么区别。"],
            ["用 Python 写一个并发批量请求 API 的示例。"],
        ],
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=8,
        max_size=32,
    ).launch()