from openai import OpenAI
import json
from event_trigger import event_trigger
import os

# 设置OpenAI API密钥和基础URL
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-3.5-turbo")

tools = [
    {
        "type": "function",
        "function": {
            'name': 'generate_complaint_chain',
            'description': '根据角色信息和近期遭遇的事件，生成一个患者的主诉请求认知变化链',
            'parameters': {
                "type": "object",
                "properties": {
                    "chain": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "stage": {
                                    "type": "integer"
                                },
                                "content": {
                                    "type": "string"
                                }
                            },
                            "additionalProperties": False,
                            "required": [
                                "stage",
                                "content"
                            ]
                        },
                        "minItems": 3,
                        "maxItems": 7
                        }
                },
                "required": ["chain"]
            },
        }
    }
]

# 根据profile和event生成主诉启发链
def gen_complaint_chain(profile):
    # 提取患者信息
    patient_info = f"### 患者信息\n年龄：{profile['age']}\n性别：{profile['gender']}\n职业：{profile['occupation']}\n婚姻状况：{profile['marital_status']}\n症状：{profile['symptoms']}"

    event = event_trigger(profile)

    client = OpenAI(
        api_key=api_key,
        base_url=base_url
    )

    response = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "user", "content": f"### 任务\n根据患者情况及近期遭遇事件生成患者的主诉认知变化链。请注意，事件可能与患者信息冲突，如果发生这种情况，以患者的信息为准。\n{patient_info}\n### 近期遭遇事件\n{event}"}
        ],
        tools=tools,
        tool_choice={"type": "function", "function": {"name": "generate_complaint_chain"}}
    )

    chain = json.loads(response.choices[0].message.tool_calls[0].function.arguments)["chain"]

    return chain

# unit test
# while True:
#     # 模拟患者信息
#     profile = {
#         "drisk": 3,
#         "srisk": 2,
#         "age": "42",
#         "gender": "女",
#         "marital_status": "离婚",
#         "occupation": "教师",
#         "symptoms": "缺乏自信心，自我价值感低，有自罪感，无望感；体重剧烈增加；精神运动性激越；有自杀想法"
#         }

#     print(gen_complaint_chain(profile))