import io
import os
import base64
import gradio as gr
from PIL import Image
from openai import OpenAI

# ------- 配置区 -------
# 推荐在 HF Space 的 Settings → Variables and secrets 里设置：
# Name: OPENAI_API_KEY   Value: 你的 StepFun API Key
# 如需自定义变量名（比如 STEPFUN_KEY），下面会做兼容处理。
STEPFUN_ENDPOINT = "https://platform.stepfun.com/v1"
MODEL_NAME = "step3-fp8"
MAX_TOKENS = 1024
# ---------------------

def _get_api_key():
    # 优先用 OPENAI_API_KEY（OpenAI SDK 的默认约定），否则回退 STEPFUN_KEY
    return os.getenv("OPENAI_API_KEY") or os.getenv("STEPFUN_KEY")

def _pil_to_data_url(img: Image.Image, fmt: str = "PNG") -> str:
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    mime = "image/png" if fmt.upper() == "PNG" else "image/jpeg"
    return f"data:{mime};base64,{b64}"

def _make_client() -> OpenAI:
    key = _get_api_key()
    if not key:
        # 给出更友好的报错，避免默认的 OpenAIError 让人摸不着头脑
        raise RuntimeError(
            "API Key 未设置。请到 Space 的 Settings → Variables and secrets 添加：\n"
            "Name=OPENAI_API_KEY，Value=你的 StepFun API Key（或用 STEPFUN_KEY 也可）。"
        )
    return OpenAI(api_key=key, base_url=STEPFUN_ENDPOINT)

client = _make_client()

def chat_with_step3(image: Image.Image, question: str):
    if image is None:
        return "请先上传图片。"
    if not question:
        question = "请描述这张图片。"

    # 将 PIL 图转成 data URL，传到 OpenAI 兼容接口
    data_url = _pil_to_data_url(image, fmt="PNG")

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": data_url}},
                {"type": "text", "text": question},
            ],
        }
    ]

    try:
        resp = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=MAX_TOKENS,
        )
        return resp.choices[0].message.content
    except Exception as e:
        # 返回更清晰的错误信息到前端，方便你在 Space 控制台外定位问题
        return f"调用失败：{e}"

# 用 Interface 足够简单，也可改 Blocks
iface = gr.Interface(
    fn=chat_with_step3,
    inputs=[
        gr.Image(type="pil", label="Upload image"),
        gr.Textbox(label="Question", placeholder="问点什么…"),
    ],
    outputs=gr.Textbox(label="Answer"),
    title="Step3 FP8 (API) Demo",
    description="使用 StepFun 的 OpenAI 兼容 API（/v1）来调用 step3-fp8 模型。",
)

if __name__ == "__main__":
    # 在 HF Spaces 里不需要 share=True
    iface.launch()