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://api.stepfun.com/v1"
MODEL_NAME = "step-3"

# --------------------

def _get_api_key() -> str:
    """获取 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:
    """将 PIL 图片转换成 base64 Data URL"""
    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:
        raise RuntimeError(
            "API Key 未设置\n请到 Space 的 Settings - Variables and secrets 添加：\n"
            "Name=OPENAI_API_KEY, Value=你的 StepFun API Key（或用 STEPFUN_KEY 也可）。"
        )
    # 创建客户端时指定 StepFun 的 base_url
    return OpenAI(api_key=key, base_url=STEPFUN_ENDPOINT)


client = _make_client()


def chat_with_step3(image: Image.Image, question: str) -> str:
    """调用 StepFun 的 step-3 模型进行推理"""
    if image is None:
        return "请先上传图片。"
    if not question:
        question = "请描述这张图片。"

    # 将 PIL 图片转成 Data URL，传到 StepFun
    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:
        # 调用 StepFun 的 OpenAI 兼容接口，不设置 max_tokens
        resp = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
        )
        return resp.choices[0].message.content
    except Exception as e:
        # 返回异常信息，方便调试
        return f"调用失败: {e!r}"


# 构建 Gradio 界面
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 (API) Demo",
    description="使用 StepFun 的 OpenAI 兼容 API (/v1) 来调用 step-3 模型。",
)

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