# pip install openai python-dotenv
import os
import json
from openai import OpenAI
from dotenv import load_dotenv
from httpx import Client, Proxy

load_dotenv()


class KIMI:
    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.getenv(
            "KIMI_API_KEY", os.getenv("KIMI_API_KEY")
        )
        if not self.key:
            raise RuntimeError("找不到 KIMI_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://api.moonshot.cn/v1",  # 已去掉多余空格
            http_client=httpx_client,
        )

    def chat(self, question_json: dict) -> str:
        resp = self.client.chat.completions.create(
            model="kimi-k2-turbo-preview",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {
                    "role": "user",
                    "content": json.dumps(question_json, ensure_ascii=False),
                },
            ],
            temperature=0,
        )
        return resp.choices[0].message.content or ""


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