import asyncio
import base64
from io import BytesIO
import mss
import mss.tools
import tkinter as tk
from tkinter import messagebox
import websockets
import threading
import time  # 导入 time 模块
import tempfile  # 导入 tempfile 模块

# 全局变量，用于控制WebSocket服务器的状态
server = None
loop = None

async def handle_screenshot(websocket, path=None):
    print(f"Handling connection at path: {path or '/'}")  # 调试信息
    async for message in websocket:
        try:
            # 解析收到的消息
            x1, y1, x2, y2 = map(int, message.split(','))

            # 记录开始时间
            # start_time = time.time()

            # 使用 mss 进行截图
            with mss.mss() as sct:
                bbox = {'left': x1, 'top': y1, 'width': x2 - x1, 'height': y2 - y1}
                img = sct.grab(bbox)  # 这种方法截图大概在6到12毫秒，非常快

            # 记录结束时间
            # end_time = time.time()

            # 计算并打印花费的时间（毫秒）
            # elapsed_time_ms = (end_time - start_time) * 1000
            # print(f"Screenshot took {elapsed_time_ms:.2f} ms")

            # 使用临时文件保存截图
            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
                mss.tools.to_png(img.rgb, img.size, output=temp_file.name)
                # 读取临时文件的内容
                with open(temp_file.name, 'rb') as f:
                    img_data = f.read()

            # 将截图转换为Base64字符串
            img_str = base64.b64encode(img_data).decode()
            # 发送Base64字符串给客户端
            await websocket.send(img_str)
        except Exception as e:
            print(f"Error handling screenshot request: {e}")

async def start_server():
    global server
    server = await websockets.serve(handle_screenshot, '127.0.0.1', 12345)
    print("Server started")

def start_or_stop_server():
    global server, loop
    if server is None:
        # 提交任务到事件循环
        asyncio.run_coroutine_threadsafe(start_server(), loop)
        button.config(text="停止服务")
        messagebox.showinfo("信息", "WebSocket服务已启动")
    else:
        server.close()
        server = None
        button.config(text="启动服务")
        messagebox.showinfo("信息", "WebSocket服务已停止")

def run_event_loop():
    global loop
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_forever()

# 创建一个线程来运行事件循环
thread = threading.Thread(target=run_event_loop, daemon=True)
thread.start()

# 创建Tkinter窗口
root = tk.Tk()
root.title("WebSocket截图服务")
root.geometry("350x230")  # 设置窗口大小

# 添加文本
label = tk.Label(root, text="注意事项:\n1. 请求地址'ws://127.0.0.1:12345'。\n2. 传递四个参数x1, y1, x2, y2，\n分别是左上角坐标和右下角坐标。\njs语言示例：ws.send(`${x1},${y1},${x2},${y2}`);\n3. 返回base64字符串。")
label.pack(pady=20)

# 添加按钮
button = tk.Button(root, text="启动服务", command=start_or_stop_server)
button.pack(pady=20)

# 运行Tkinter主循环
root.mainloop()