# %%
import time
import cv2
import numpy as np
from mss import mss
from collections import deque
import pyautogui

# %%
from typing import Callable, Optional

# %%
import json
import threading
import websocket
from queue import Queue


class WebSocketClient:
    def __init__(self, server_url):
        self.server_url = server_url
        self.ws = None
        self.connected = False
        self.message_queue = Queue()
        self.lock = threading.Lock()
        self.reconnect_delay = 3  # 重连延迟(秒)
        self.running = False

    def on_open(self, ws):
        """连接成功回调"""
        print("WebSocket 连接成功")
        self.connected = True

    def on_message(self, ws, message):
        """接收消息回调"""
        print(f"收到前端消息: {message}")
        data = json.loads(message)
        if data.get("type") == "click":
            x, y = data["position"].values()
            pyautogui.click(x, y)

    def on_error(self, ws, error):
        """错误回调"""
        print(f"WebSocket 错误: {error}")
        self.connected = False

    def on_close(self, ws, close_status_code, close_msg):
        """关闭回调"""
        print(f"WebSocket 关闭: {close_status_code} - {close_msg}")
        self.connected = False
        if self.running:
            self.reconnect()

    def reconnect(self):
        """断线重连"""
        print(f"{self.reconnect_delay}秒后尝试重连...")
        time.sleep(self.reconnect_delay)
        self.connect()

    def connect(self):
        """连接WebSocket服务器"""
        self.running = True
        self.ws = websocket.WebSocketApp(
            self.server_url,
            on_open=self.on_open,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close,
        )
        # 在独立线程中运行WebSocket
        threading.Thread(target=self.ws.run_forever, daemon=True).start()

    def send(self, data):
        """发送消息(线程安全)"""
        with self.lock:
            if self.connected and self.ws:
                try:
                    self.ws.send("xtleto|" + json.dumps(data))
                    return True
                except Exception as e:
                    print(f"消息发送失败: {e}")
                    return False
            return False

    def stop(self):
        """停止WebSocket连接"""
        self.running = False
        if self.ws:
            self.ws.close()


# %%
class DynamicAdjuster:
    def __init__(self):
        # 参数跟踪窗口（保留最近20帧数据）
        fps = 30
        self.diff_history = deque(maxlen=fps)
        self.area_history = deque(maxlen=fps)

        # 初始参数值
        self.threshold = 10
        self.min_area = 100

        # 调整速率系数
        self.THRESH_STEP = 2  # 阈值调整步长
        self.AREA_STEP = 5  # 面积调整步长
        self.TARGET_FPS = fps  # 目标处理速度（用于CPU优化）

    def adjust_by_fps(self, actual_fps):
        """根据实际帧率动态调整参数"""
        if actual_fps < self.TARGET_FPS * 0.8:
            # 当帧率过低时激进调节参数
            self.THRESH_STEP = max(3, self.THRESH_STEP + 1)
            self.min_area = min(300, self.min_area + 10)
        elif actual_fps > self.TARGET_FPS * 1.2:
            # 当帧率过高时放松限制
            self.THRESH_STEP = max(1, self.THRESH_STEP - 1)

    def update(self, diff_frame, detected_areas):
        """根据当前帧数据动态调整参数"""
        # 计算当前帧的运动强度
        curr_diff = cv2.mean(diff_frame)[0]
        self.diff_history.append(curr_diff)

        # 计算当前检测区域面积
        curr_area = sum([(x2 - x1) * (y2 - y1) for (x1, y1, x2, y2) in detected_areas])
        self.area_history.append(curr_area)

        # 计算平均运动强度（最近N帧的指数加权平均值）
        avg_diff = (
            sum(self.diff_history) / len(self.diff_history) if self.diff_history else 0
        )

        # 自动调整阈值（运动强度越低，阈值越高）
        if avg_diff < 10:  # 低运动量阶段
            self.threshold = min(20, self.threshold + self.THRESH_STEP)
        elif avg_diff > 30:  # 高运动量阶段
            self.threshold = max(15, self.threshold - self.THRESH_STEP * 2)
        else:  # 正常调整
            if len(detected_areas) > 5:
                self.threshold = max(10, self.threshold - self.THRESH_STEP)
            elif not detected_areas:
                self.threshold = min(50, self.threshold + self.THRESH_STEP)

        # 根据检测区域面积自动调整最小面积
        avg_area = (
            sum(self.area_history) / len(self.area_history) if self.area_history else 0
        )
        if avg_area < 500:  # 小范围变化
            self.min_area = min(200, self.min_area + self.AREA_STEP)
        else:  # 大范围变化
            self.min_area = max(50, self.min_area - self.AREA_STEP)


# %%
class FrameDiffProcessor:
    def __init__(self, ws_client: Optional[WebSocketClient] = None):
        self.prev_frame = None
        self.adjuster = DynamicAdjuster()
        self.kernel_size = (21, 21)
        self.last_diff = None
        self.ws_client = ws_client

    def _preprocess(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        return cv2.GaussianBlur(gray, self.kernel_size, 0)

    def process(self, frame):
        processed = self._preprocess(frame)
        if self.prev_frame is None:
            self.prev_frame = processed
            return [], self.adjuster.threshold, self.adjuster.min_area, 0

        # 计算原始差异
        raw_diff = cv2.absdiff(self.prev_frame, processed)
        self.last_diff = raw_diff.copy()

        # 应用当前阈值
        _, thresh = cv2.threshold(
            raw_diff, self.adjuster.threshold, 255, cv2.THRESH_BINARY
        )

        # 形态学优化
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

        # 查找轮廓并过滤
        contours, _ = cv2.findContours(
            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        rects = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > self.adjuster.min_area:
                x, y, w, h = cv2.boundingRect(cnt)
                rects.append((x, y, x + w, y + h))

        if rects:
            payload = {
                "type": "motion_detection",
                "timestamp": int(time.time() * 1000),
                "regions": [],
                "stats": {
                    "threshold": self.adjuster.threshold,
                    "min_area": self.adjuster.min_area,
                    "fps": self.adjuster.TARGET_FPS,
                },
            }
            for x1, y1, x2, y2 in rects:
                roi = frame[y1:y2, x1:x2]
                # 将图像转换为base64编码
                _, buffer = cv2.imencode(".webp", roi, [cv2.IMWRITE_WEBP_QUALITY, 80])
                payload["regions"].append(
                    {
                        "x": x1,
                        "y": y1,
                        "width": x2 - x1,
                        "height": y2 - y1,
                        "image": buffer.tobytes().hex(),  # 转换为十六进制字符串
                    }
                )

            # 通过WebSocket发送
            self.ws_client.send(payload)
        # 动态参数调整
        self.adjuster.update(raw_diff, rects)

        self.prev_frame = processed
        return (
            rects,
            self.adjuster.threshold,
            self.adjuster.min_area,
            self.adjuster.TARGET_FPS,
        )


# %%
# 初始化捕获
sct = mss()
monitor = sct.monitors[1]
ws_client = WebSocketClient("ws://43.248.184.71:48905?mac=server")
ws_client.connect()
processor = FrameDiffProcessor(ws_client)

cv2.namedWindow("Adaptive Detection")
prev_time = time.time()
while True:
    sct_img = sct.grab(monitor)
    frame = np.array(sct_img)

    rects, curr_thresh, curr_area, fps = processor.process(frame)

    # 计算实际FPS
    curr_time = time.time()
    actual_fps = 1 / (curr_time - prev_time)
    prev_time = curr_time

    # 调节参数
    processor.adjuster.adjust_by_fps(actual_fps)
    # 可视化
    display = frame.copy()
    for x1, y1, x2, y2 in rects:
        cv2.rectangle(display, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 显示调整信息
    status = f"Thresh: {curr_thresh} | MinArea: {curr_area}| FPS:{ fps}"
    cv2.putText(
        display, status, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2
    )

    # 显示差异图（可选）
    # if processor.last_diff is not None:
    #     diff_display = cv2.normalize(processor.last_diff, None, 0, 255, cv2.NORM_MINMAX)
    #     cv2.imshow("Difference", diff_display.astype(np.uint8))

    cv2.imshow("Adaptive Detection", display)

    if cv2.waitKey(25) & 0xFF == 27:
        break

cv2.destroyAllWindows()
ws_client.stop()