import socket
import json
import time
from openai import OpenAI

# 初始化 DeepSeek 客户端
client = OpenAI(
    api_key="sk-420190f448fe41158c4e2ccff90e35ce", 
    base_url="https://api.deepseek.com" 
) 
MODEL_NAME = "deepseek-chat"

def fetch_and_translate_tools():
    """向 CAD 基座发送 tools/list，并转换为大模型认识的格式"""
    try:
        cad_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cad_socket.connect(("127.0.0.1", 8080))
        
        # 1. 向 MCP Server 发送标准的 tools/list 请求
        req = {
            "jsonrpc": "2.0",
            "id": 100,
            "method": "tools/list"
        }
        cad_socket.sendall(json.dumps(req).encode('utf-8'))
        
        response_data = cad_socket.recv(8192)
        cad_socket.close()
        
        mcp_response = json.loads(response_data.decode('utf-8'))
        
        if "result" not in mcp_response or "tools" not in mcp_response["result"]:
            print("[Error] 无法从 CAD 获取工具列表")
            return []
            
        llm_tools = []
        # 2. 遍历 CAD 返回的工具，将其翻译为 DeepSeek/OpenAI 格式
        for mcp_tool in mcp_response["result"]["tools"]:
            llm_tools.append({
                "type": "function",
                "function": {
                    "name": mcp_tool["name"],
                    "description": mcp_tool["description"],
                    # 核心转换：MCP 的 inputSchema 等价于 LLM 的 parameters
                    "parameters": mcp_tool["inputSchema"] 
                }
            })
            
        print(f"[Init] 成功从 CAD 动态加载了 {len(llm_tools)} 个工具！")
        return llm_tools
        
    except Exception as e:
        print(f"[Error] 获取工具列表失败: {e}")
        return []

def call_cad_mcp_server(tool_name, arguments):
    """底层 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 run_agent(user_prompt):
    print("=== 🚀 CAD 具身智能 Agent 启动 ===")

    # 动态向 CAD 请求可用工具，彻底解耦！
    tools = fetch_and_translate_tools()
    
    if not tools:
        print("没有可用的工具，Agent 退出。")
        return

    # 核心心智设定：教它怎么形成视觉闭环
    system_prompt = """你是一个高阶 AutoCAD 视觉检查专家。
请严格遵循以下工作流：
1. 第一步永远是调用 `get_viewport_screenshot` 观察当前画面。
2. 仔细评估图像。如果目标物体（如图元、文字）太小导致无法确信细节，请调用 `zoom_window_normalized` 放大该局部区域。
3. 【关键指令】：每次执行缩放操作（zoom）后，你必须在下一回合再次调用 `get_viewport_screenshot` 获取放大后的新截图！
4. 反复观察和放大，直到你 100% 看清细节，再用自然语言向用户输出最终结论。"""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    
    MAX_TURNS = 8  # 熔断机制：最多允许思考 8 个回合，防止无限套娃
    
    for turn in range(MAX_TURNS):
        print(f"\n--- [第 {turn + 1} 回合] 思考中 ---")
        
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            tools=tools
        )
        
        assistant_message = response.choices[0].message
        messages.append(assistant_message) # 记录思维路径
        
        # 1. 检查是否需要执行物理动作
        if assistant_message.tool_calls:
            for tool_call in assistant_message.tool_calls:
                tool_name = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                
                print(f"[动作决定] ⚡ 调用工具: {tool_name}")
                if args: print(f"          参数: {args}")
                
                # 执行 CAD 通信
                cad_res = call_cad_mcp_server(tool_name, args)
                time.sleep(0.5) # 给予 CAD 渲染刷新窗口的时间缓冲
                
                # 处理执行结果，并准备发回给 LLM 的观测报告
                tool_result_content = []
                
                if cad_res and "result" in cad_res:
                    for item in cad_res["result"]["content"]:
                        if item["type"] == "text":
                            tool_result_content.append({"type": "text", "text": item["text"]})
                        elif item["type"] == "image":
                            base64_img = item["data"]
                            mime = item.get("mimeType", "image/png")
                            tool_result_content.append({
                                "type": "image_url",
                                "image_url": {"url": f"data:{mime};base64,{base64_img}"}
                            })
                            print("[观测反馈] 📸 已将最新 CAD 屏幕画面传回视觉中枢。")
                else:
                    tool_result_content.append({"type": "text", "text": "CAD 工具执行失败或无响应。"})
                
                # 必须向大模型提交 Tool 返回结果
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result_content
                })
                
        # 2. 如果不调工具，说明任务已完成，输出最终自然语言结论
        else:
            print("\n==================================")
            print(f"🎯 [最终结论]:\n{assistant_message.content}")
            print("==================================")
            break
            
    else:
        print("\n[警告] 达到最大思考回合数限制，Agent 强行中止。")

if __name__ == "__main__":
    run_agent("请仔细检查当前图纸，告诉我图纸中心那些小圆的内部，是否还包含了更小的同心圆？如果看不清，请务必放大确认。")