import openai
import json
from tools import Tools

with open('key.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

openai_url = config.get("openai_url")
api_key = config.get("api_key")
model = config.get("model")


client = openai.OpenAI(
    base_url=openai_url,
    api_key=api_key,
)

tool = Tools()

# 流式输出处理函数
def stream_response(messages, tools=None):
    full_response = ""
    tool_calls = []

    # 使用client对象调用API
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        tools=tools,
        tool_choice="auto",
        stream=True
    )

    # 处理流式响应
    for chunk in response:
        # 提取delta内容
        delta = chunk.choices[0].delta

        # 处理文本内容（最终回答）
        if hasattr(delta, 'content') and delta.content:
            content = delta.content
            full_response += content
            print(content, end="", flush=True)  # 实时输出到控制台

        # 处理工具调用（tool_calls）
        if hasattr(delta, 'tool_calls') and delta.tool_calls:
            for tool_call in delta.tool_calls:
                # 初始化tool_calls数组（如果需要）
                while len(tool_calls) <= tool_call.index:
                    tool_calls.append({
                        "id": "",
                        "type": "function",
                        "function": {
                            "name": "",
                            "arguments": ""
                        }
                    })

                # 更新工具调用信息
                if tool_call.id:
                    tool_calls[tool_call.index]["id"] = tool_call.id

                if tool_call.function:
                    if tool_call.function.name:
                        tool_calls[tool_call.index]["function"]["name"] = tool_call.function.name
                    if tool_call.function.arguments:
                        tool_calls[tool_call.index]["function"]["arguments"] += tool_call.function.arguments

    # 如果有工具调用，处理并返回
    if tool_calls:
        # 解析工具调用参数并执行
        executed_tools = []
        for tool_call in tool_calls:
            try:
                arguments = json.loads(tool_call["function"]["arguments"])
            except json.JSONDecodeError:
                arguments = {}

            tool_name = tool_call["function"]["name"]
            tool_id = tool_call["id"]

            # 执行工具
            tool_result = tool.call_tool(tool_name, arguments)

            executed_tools.append({
                "name": tool_name,
                "arguments": arguments,
                "result": tool_result,
                "id": tool_id
            })

        return {
            "content": full_response,
            "tool_calls": executed_tools
        }

    # 返回完整响应和工具调用信息
    return {
        "content": full_response,
        "tool_calls": None
    }


# 主函数：处理多轮工具调用的流式输出
def run_conversation(messages):
    # 系统提示
    system_prompt = """
你是一个可以调用工具的AI助手。请根据用户需求，判断是否需要调用工具：
1. 若需要，可调用工具列表中的工具（可多次调用，直到获取足够信息）；
2. 调用工具后，需等待工具返回结果，再决定是否继续调用或整理结果回答用户；
3. 工具返回结果会以函数消息的形式提供给你，请基于所有信息生成最终回答。
你的回答需要保持严谨和专业，假如你无法从工具调用结果和系统提示中获取足够的信息，请不要随意回答，请直接告诉用户暂无详情
"""

    # 添加系统提示到消息历史（如果还没有）
    if not any(msg.get("role") == "system" for msg in messages):
        messages.insert(0, {"role": "system", "content": system_prompt})

    print("\n🤖 助手: ", end="", flush=True)

    # 多轮工具调用循环
    while True:
        # 流式获取模型响应
        response = stream_response(messages, tool.get_tool_list())

        # 检查是否需要调用工具
        if response["tool_calls"]:
            tool_calls = response["tool_calls"]

            # 添加助手的工具调用到消息历史
            tool_call_list = []
            for tool_call in tool_calls:
                tool_call_list.append({
                    "id": tool_call["id"],
                    "type": "function",
                    "function": {
                        "name": tool_call["name"],
                        "arguments": json.dumps(tool_call["arguments"])
                    }
                })

            messages.append({
                "role": "assistant",
                "tool_calls": tool_call_list
            })

            # 执行所有工具并将结果添加到对话历史
            for tool_call in tool_calls:
                tool_name = tool_call["name"]
                tool_result = tool_call["result"]
                tool_id = tool_call["id"]

                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_id,
                    "name": tool_name,
                    "content": tool_result
                })

                print(f"\n\n🔍 工具结果 ({tool_name}): {tool_result}")

            print("\n🤖 助手: ", end="", flush=True)
        else:
            # 没有工具调用，对话结束
            messages.append({
                "role": "assistant",
                "content": response["content"]
            })
            break

    return response["content"]


# 多轮对话主循环
def main():
    # 初始化对话历史
    messages = []
    
    print("开始多轮对话（输入'exit'结束对话）")
    
    while True:
        user_input = input("\n👨‍💻 用户: ").strip()
        
        if user_input.lower() in ['退出', 'quit', 'exit']:
            print("对话结束")
            break
            
        if user_input:
            messages.append({"role": "user", "content": user_input})
            run_conversation(messages)
        else:
            print("请输入有效内容")


# 示例：运行对话
if __name__ == "__main__":
    main()