import io
import os
import base64
import gradio as gr
from PIL import Image
import httpx  # Use httpx for direct API calls instead of openai SDK

# ------- 配置区 -------
# 推荐在 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，如果没有设置则返回 None。
    首先尝试读取环境变量 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。
    接收一个 PIL.Image 对象和输出格式（默认为 PNG），
    返回可用于 StepFun OpenAI 兼容接口的 data:image/...;base64,... 字符串。
    """
    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 _post_chat(messages: list, temperature: float = 0.7) -> str:
    """
    调用 StepFun 的 chat/completions 接口并返回模型回复。
    使用 httpx 库向 StepFun 的 OpenAI 兼容接口发送请求，
    避免使用 openai SDK 导致的 "No API found" 错误。
    messages 参数应符合 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 也可）。"
        )
    url = f"{STEPFUN_ENDPOINT}/chat/completions"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": MODEL_NAME,
        "messages": messages,
        "temperature": temperature,
    }
    resp = httpx.post(url, headers=headers, json=payload, timeout=60)
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"]

def chat_with_step3(image: Image.Image, question: str) -> str:
    """
    调用 StepFun 的 step-3 模型进行推理。
    首先检查上传的图像和问题文本是否有效，将图像编码为 data URL，
    构造符合 OpenAI 接口规范的 messages 数组，然后通过 `_post_chat` 发送请求。
    如遇异常则返回错误信息。
    """
    if image is None:
        return "请先上传图片。"
    if not question:
        question = "请描述这张图片。"
    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:
        return _post_chat(messages)
    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()