# Please install OpenAI SDK first: `pip3 install openai`
import os
import json
from openai import OpenAI
from dotenv import load_dotenv
from httpx import Client, Proxy

load_dotenv()


class DOUBAO:
    system_prompt = (
        "You are a professional QA answer assistant.\n"
        "The user will provide a JSON object containing a question.\n"
        "Your job:\n"
        "1. Read the 'title' and the 'chooies'.\n"
        "2. Determine the correct answer(s).\n"
        "3. Output ONLY the 'num' of the correct choices.\n"
        "4. If multiple answers exist, join them with English commas, e.g., 'A,C,D'.\n"
        "5. If you cannot determine the answer confidently, return an empty string.\n"
        "6. Do not output explanations or any extra text."
    )

    def __init__(self) -> None:
        self.key = os.environ.get("DOUBAO_API_KEY")
        if not self.key:
            raise Exception("找不到 DOUBAO_API_KEY 请设置环境变量在运行代码")

        httpx_client = Client(
            # proxy=Proxy("socks5://127.0.0.1:9000"),  # 传入SOCKS代理地址
            timeout=30.0,  # 可选：设置超时时间
            # verify=False,
        )

        self.client = OpenAI(
            api_key=self.key,
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            http_client=httpx_client,
        )

    def chat(self, question_json: dict) -> str:
        response = self.client.chat.completions.create(
            # model="deepseek-chat",
            model="doubao-seed-code-preview-251028",
            # model="deepseek-reasoner",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": json.dumps(question_json)},
            ],
            stream=False,
        )

        # print(response.choices[0].message.content)
        return (
            response.choices[0].message.content
            if response.choices[0].message.content
            else ""
        )


if __name__ == "__main__":
    question = {
        "id": "25594265",
        "index": "1",
        "type": "单选",
        "title": "下列哪个属于AI视觉识别应用？",
        "chooies": [
            {"num": "A", "txt": "音乐合成"},
            {"num": "B", "txt": "植物识别"},
            {"num": "C", "txt": "文案创作"},
            {"num": "D", "txt": "机器翻译"},
        ],
    }
    DOUBAO().chat(question)