video/ws/ws.py

import asyncio
import datetime
import json
import os
from contextlib import asynccontextmanager
from typing import Dict, Optional, AsyncGenerator
from service.device_service import update_online_status_by_ip, increment_alarm_count_by_ip
from service.device_action_service import add_device_action
from schema.device_action_schema import DeviceActionCreate
from core.all import detect

import cv2
import numpy as np
from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
from core.all import load_model

# 配置常量
HEARTBEAT_INTERVAL = 30  # 心跳检查间隔（秒）
HEARTBEAT_TIMEOUT = 600  # 客户端超时阈值（秒）
WS_ENDPOINT = "/ws"  # WebSocket端点路径
FRAME_QUEUE_SIZE = 1  # 帧队列大小限制


# 工具函数: 获取格式化时间字符串（统一时间戳格式）
def get_current_time_str() -> str:
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def get_current_time_file_str() -> str:
    return datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")


# 客户端连接封装
class ClientConnection:
    def __init__(self, websocket: WebSocket, client_ip: str):
        self.websocket = websocket
        self.client_ip = client_ip
        self.last_heartbeat = datetime.datetime.now()
        self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE)
        self.consumer_task: Optional[asyncio.Task] = None

    def update_heartbeat(self):
        """更新心跳时间（客户端发送心跳时调用）"""
        self.last_heartbeat = datetime.datetime.now()

    def is_alive(self) -> bool:
        """判断客户端是否存活（心跳超时检查）"""
        timeout = (datetime.datetime.now() - self.last_heartbeat).total_seconds()
        return timeout < HEARTBEAT_TIMEOUT

    def start_consumer(self):
        """启动帧消费任务"""
        self.consumer_task = asyncio.create_task(self.consume_frames())
        return self.consumer_task

    async def send_frame_permit(self):
        """
        发送「帧发送许可信号」
        通知客户端可发送下一帧图像
        """
        try:
            frame_permit_msg = {
                "type": "frame",
                "timestamp": get_current_time_str(),
                "client_ip": self.client_ip
            }
            await self.websocket.send_json(frame_permit_msg)
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 已发送帧发送许可信号（取帧后立即通知）")
        except Exception as e:
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧许可信号发送失败 - {str(e)}")

    async def consume_frames(self) -> None:
        """消费队列中的帧并处理（核心调整: 取帧后立即发许可、再处理帧）"""
        try:
            while True:
                # 1. 从队列取出帧（阻塞直到有帧可用）
                frame_data = await self.frame_queue.get()

                # -------------------------- 核心修改: 取出帧后立即发送下一帧许可 --------------------------
                await self.send_frame_permit()  # 取帧即通知客户端发下一帧、无需等处理完成
                # -----------------------------------------------------------------------------------------

                try:
                    # 2. 处理取出的帧（即使处理慢、客户端也已收到许可、可提前准备下一帧）
                    await self.process_frame(frame_data)
                finally:
                    # 3. 标记帧任务完成（无论处理成功/失败、都需清理队列）
                    self.frame_queue.task_done()

        except asyncio.CancelledError:
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费任务已取消")
        except Exception as e:
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费逻辑错误 - {str(e)}")

    async def process_frame(self, frame_data: bytes) -> None:
        """处理单帧图像数据（检测违规后发送危险通知 + 调用违规次数加一方法）"""
        # 二进制数据转OpenCV图像
        nparr = np.frombuffer(frame_data, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img is None:
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 无法解析图像数据")
            return

        # 确保图像保存目录存在
        os.makedirs('images', exist_ok=True)

        # 保存图像（按IP+时间戳命名、避免冲突）
        filename = f"images/{self.client_ip.replace('.', '_')}_{get_current_time_file_str()}.jpg"
        try:
            cv2.imwrite(filename, img)
            print(f"[{get_current_time_str()}] 图像已保存至: {filename}")
            has_violation, data, type = detect(img)
            print(has_violation)
            print(type)
            print(data)
            if has_violation:
                print(
                    f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测到违规 - 类型: {type}, 详情: {data}")

                # 调用违规次数加一方法
                try:
                    await asyncio.to_thread(increment_alarm_count_by_ip, self.client_ip)
                    print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数已+1")
                except Exception as e:
                    print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数更新失败 - {str(e)}")

                # 发送「危险通知」
                danger_msg = {
                    "type": "danger",
                    "timestamp": get_current_time_str(),
                    "client_ip": self.client_ip
                }
                await self.websocket.send_json(danger_msg)
            else:
                print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 未检测到违规")
        except Exception as e:
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 图像处理错误 - {str(e)}")


# 全局状态管理
connected_clients: Dict[str, ClientConnection] = {}
heartbeat_task: Optional[asyncio.Task] = None


# 心跳检查（定时清理超时客户端 + 调用离线状态更新方法）
async def heartbeat_checker():
    while True:
        current_time = get_current_time_str()
        timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]

        if timeout_ips:
            print(f"[{current_time}] 心跳检查: {len(timeout_ips)}个客户端超时（IP: {timeout_ips}）")
            for ip in timeout_ips:
                try:
                    conn = connected_clients[ip]
                    if conn.consumer_task and not conn.consumer_task.done():
                        conn.consumer_task.cancel()
                    await conn.websocket.close(code=1008, reason="心跳超时")

                    # 超时设为离线并记录
                    try:
                        await asyncio.to_thread(update_online_status_by_ip, ip, 0)
                        action_data = DeviceActionCreate(client_ip=ip, action=0)
                        await asyncio.to_thread(add_device_action, action_data)
                        print(f"[{current_time}] 客户端{ip}: 已标记为离线并记录操作")
                    except Exception as e:
                        print(f"[{current_time}] 客户端{ip}: 离线状态更新失败 - {str(e)}")
                finally:
                    connected_clients.pop(ip, None)
        else:
            print(f"[{current_time}] 心跳检查: {len(connected_clients)}个客户端在线")

        await asyncio.sleep(HEARTBEAT_INTERVAL)


# 应用生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    global heartbeat_task
    heartbeat_task = asyncio.create_task(heartbeat_checker())
    print(f"[{get_current_time_str()}] 全局心跳检查任务启动（任务ID: {id(heartbeat_task)}）")
    yield
    if heartbeat_task and not heartbeat_task.done():
        heartbeat_task.cancel()
        try:
            await heartbeat_task
            print(f"[{get_current_time_str()}] 全局心跳检查任务已取消")
        except asyncio.CancelledError:
            pass


# 消息处理工具函数
async def send_heartbeat_ack(conn: ClientConnection):
    try:
        heartbeat_ack_msg = {
            "type": "heart",
            "timestamp": get_current_time_str(),
            "client_ip": conn.client_ip
        }
        await conn.websocket.send_json(heartbeat_ack_msg)
        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 已发送心跳确认")
        return True
    except Exception as e:
        connected_clients.pop(conn.client_ip, None)
        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 心跳确认发送失败 - {str(e)}")
        return False


async def handle_text_msg(conn: ClientConnection, text: str):
    try:
        msg = json.loads(text)
        if msg.get("type") == "heart":
            conn.update_heartbeat()
            await send_heartbeat_ack(conn)
        else:
            print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 未知文本消息类型（{msg.get('type')}）")
    except json.JSONDecodeError:
        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 无效JSON文本消息")


async def handle_binary_msg(conn: ClientConnection, data: bytes):
    try:
        conn.frame_queue.put_nowait(data)
        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 图像数据（{len(data)}字节）已加入队列")
    except asyncio.QueueFull:
        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 帧队列已满、丢弃当前图像数据")


# WebSocket路由配置
ws_router = APIRouter()


@ws_router.websocket(WS_ENDPOINT)
async def websocket_endpoint(websocket: WebSocket):
    # 加载模型
    load_model()
    await websocket.accept()
    client_ip = websocket.client.host if websocket.client else "unknown_ip"
    current_time = get_current_time_str()
    print(f"[{current_time}] 客户端{client_ip}: WebSocket连接已建立")

    is_online_updated = False

    try:
        # 处理重复连接
        if client_ip in connected_clients:
            old_conn = connected_clients[client_ip]
            if old_conn.consumer_task and not old_conn.consumer_task.done():
                old_conn.consumer_task.cancel()
            await old_conn.websocket.close(code=1008, reason="同一IP新连接建立")
            connected_clients.pop(client_ip)
            print(f"[{current_time}] 客户端{client_ip}: 已关闭旧连接")

        # 注册新连接
        new_conn = ClientConnection(websocket, client_ip)
        connected_clients[client_ip] = new_conn
        new_conn.start_consumer()
        # 初始许可: 连接建立后立即发一次、让客户端知道可发第一帧（后续靠取帧后自动发）
        await new_conn.send_frame_permit()

        # 标记上线并记录
        try:
            await asyncio.to_thread(update_online_status_by_ip, client_ip, 1)
            action_data = DeviceActionCreate(client_ip=client_ip, action=1)
            await asyncio.to_thread(add_device_action, action_data)
            print(f"[{current_time}] 客户端{client_ip}: 已标记为在线并记录操作")
            is_online_updated = True
        except Exception as e:
            print(f"[{current_time}] 客户端{client_ip}: 上线状态更新失败 - {str(e)}")

        print(f"[{current_time}] 客户端{client_ip}: 新连接注册成功、在线数: {len(connected_clients)}")

        # 消息循环
        while True:
            data = await websocket.receive()
            if "text" in data:
                await handle_text_msg(new_conn, data["text"])
            elif "bytes" in data:
                await handle_binary_msg(new_conn, data["bytes"])

    except WebSocketDisconnect as e:
        print(f"[{get_current_time_str()}] 客户端{client_ip}: 主动断开连接（代码: {e.code}）")
    except Exception as e:
        print(f"[{get_current_time_str()}] 客户端{client_ip}: 连接异常 - {str(e)[:50]}")
    finally:
        # 清理资源并标记离线
        if client_ip in connected_clients:
            conn = connected_clients[client_ip]
            if conn.consumer_task and not conn.consumer_task.done():
                conn.consumer_task.cancel()

            # 主动/异常断开时标记离线
            if is_online_updated:
                try:
                    await asyncio.to_thread(update_online_status_by_ip, client_ip, 0)
                    action_data = DeviceActionCreate(client_ip=client_ip, action=0)
                    await asyncio.to_thread(add_device_action, action_data)
                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 断开后已标记为离线")
                except Exception as e:
                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 断开后离线更新失败 - {str(e)}")

            connected_clients.pop(client_ip, None)
        print(f"[{get_current_time_str()}] 客户端{client_ip}: 资源已清理、在线数: {len(connected_clients)}")