import socket
import json
import base64

def test_screenshot():
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect(("127.0.0.1", 8080))
    
    req = {
        "jsonrpc": "2.0",
        "id": 5,
        "method": "tools/call",
        "params": {
            "name": "get_viewport_screenshot",
            "arguments": {}
        }
    }
    client.sendall(json.dumps(req).encode('utf-8'))
    
    # 图片 base64 报文可能很大，需要循环接收直到拿完
    response_data = b""
    while True:
        chunk = client.recv(8192)
        response_data += chunk
        if len(chunk) < 8192:
            break
            
    client.close()
    
    resp_json = json.loads(response_data.decode('utf-8'))
    
    # 解析并保存图片
    content_list = resp_json["result"]["content"]
    for item in content_list:
        if item["type"] == "image":
            base64_data = item["data"]
            image_bytes = base64.b64decode(base64_data)
            with open("test_screenshot.png", "wb") as f:
                f.write(image_bytes)
            print("截图已保存为 test_screenshot.png，请在当前目录查看！")

if __name__ == "__main__":
    test_screenshot()