import socket
import json
from openai import OpenAI

# 1. 初始化 DeepSeek 客户端
# 注意：务必确保已经执行过 pip install openai
client = OpenAI(
    api_key="sk-420190f448fe41158c4e2ccff90e35ce", 
    base_url="https://api.deepseek.com"  # 核心修改：将网关指向 DeepSeek 服务器
) 

# 核心修改：使用 DeepSeek 的主模型。
# (注: DeepSeek-V4/V4.1 的视觉支持已集成，具体模型名请以你 DeepSeek 后台显示的可用模型为准)
MODEL_NAME = "deepseek-flash" 

def call_cad_mcp_server(tool_name, arguments):
    """封装与 AutoCAD MCP 基座的 TCP 通信"""
    try:
        cad_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cad_socket.connect(("127.0.0.1", 8080))
        
        req = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        cad_socket.sendall(json.dumps(req).encode('utf-8'))
        
        response_data = b""
        while True:
            chunk = cad_socket.recv(8192)
            response_data += chunk
            if len(chunk) < 8192:
                break
                
        cad_socket.close()
        return json.loads(response_data.decode('utf-8'))
    except Exception as e:
        print(f"[Error] 连接 CAD 基座失败: {e}")
        return None

def main():
    print("=== DeepSeek CAD Vision Agent 启动 ===")
    
    # 2. 定义 MCP 工具 (Tool Calling)
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_viewport_screenshot",
                "description": "获取当前 AutoCAD 视口的实时截图。当用户询问图纸上的视觉特征、数量或位置时，必须先调用此工具获取画面。"
            }
        }
    ]
    
    # 3. 初始化历史
    messages = [
        {"role": "system", "content": "你是一个专业的 AutoCAD 视觉审查助手。必须先使用截图工具观察图纸，再回答用户问题。"},
        {"role": "user", "content": "请看看当前 CAD 屏幕上，我一共画了几个圆？"}
    ]
    
    print("\n[DeepSeek] 正在思考如何完成任务...")
    
    # === 第一回合：DeepSeek 思考并决定调用工具 ===
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=messages,
        tools=tools
    )
    
    assistant_message = response.choices[0].message
    messages.append(assistant_message)
    
    # 检查 DeepSeek 是否发起了 Tool Call
    if assistant_message.tool_calls:
        tool_call = assistant_message.tool_calls[0]
        tool_name = tool_call.function.name
        
        print(f"[DeepSeek] 决定调用 CAD 工具: {tool_name}")
        print(f"[CAD] 正在执行截图，请稍候...")
        
        # === 与 CAD 基座通信 ===
        cad_res = call_cad_mcp_server(tool_name, {})
        
        if cad_res and "result" in cad_res:
            content_array = cad_res["result"]["content"]
            base64_img = ""
            mime_type = "image/png"
            for item in content_array:
                if item["type"] == "image":
                    base64_img = item["data"]
                    mime_type = item.get("mimeType", "image/png")
            
            print("[CAD] 截图成功！正在将视觉数据传回大模型神经中枢...")
            
            # 4. 将 Base64 图片按标准多模态格式塞回历史
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": [
                    {"type": "text", "text": "这是 AutoCAD 当前界面的实时截图。请根据图片回答用户的问题。"},
                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_img}"}}
                ]
            })
            
            # === 第二回合：DeepSeek “看”图并回答 ===
            print("[DeepSeek] 正在进行视觉分析...")
            final_response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=messages
            )
            
            print("\n==================================")
            print(f"[DeepSeek 最终回答]:\n{final_response.choices[0].message.content}")
            print("==================================")
            
        else:
            print("[Error] CAD 工具执行失败。")
    else:
        print(f"[DeepSeek 盲猜]: {assistant_message.content}")

if __name__ == "__main__":
    main()