video/ws/ws.py

import asyncio
import datetime
import json
import os
from contextlib import asynccontextmanager
from typing import Dict, Optional, AsyncGenerator

import cv2
import numpy as np
from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI

from ocr.model_violation_detector import MultiModelViolationDetector

# 配置文件相对路径（根据实际目录结构调整）
YOLO_MODEL_PATH = r"D:\Git\bin\video\ocr\models\best.pt"
FORBIDDEN_WORDS_PATH = r"D:\Git\bin\video\ocr\forbidden_words.txt"
OCR_CONFIG_PATH = r"D:\Git\bin\video\ocr\config\1.yaml"
KNOWN_FACES_DIR = r"D:\Git\bin\video\ocr\known_faces"

# 创建检测器实例
detector = MultiModelViolationDetector(
    forbidden_words_path=FORBIDDEN_WORDS_PATH,
    ocr_config_path=OCR_CONFIG_PATH,
    yolo_model_path=YOLO_MODEL_PATH,
    known_faces_dir=KNOWN_FACES_DIR,
    ocr_confidence_threshold=0.5
)

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

# -------------------------- 核心数据结构与全局变量 --------------------------
ws_router = APIRouter()


# 客户端连接封装（包含帧队列）
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)  # 帧队列，长度为1
        self.consumer_task: Optional[asyncio.Task] = None  # 消费者任务

    # 更新心跳时间
    def update_heartbeat(self):
        self.last_heartbeat = datetime.datetime.now()

    # 检查是否存活（超时返回False）
    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 consume_frames(self) -> None:
        """从队列中获取帧并进行处理"""
        try:
            while True:
                # 从队列获取帧数据
                frame_data = await self.frame_queue.get()
                try:
                    # 处理帧数据
                    await self.process_frame(frame_data)
                finally:
                    # 标记任务完成
                    self.frame_queue.task_done()
        except asyncio.CancelledError:
            print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧消费任务已取消")
        except Exception as e:
            print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧处理错误 - {str(e)}")

    async def process_frame(self, frame_data: bytes) -> None:
        """处理单帧图像数据"""
        # 将二进制数据转换为NumPy数组（uint8类型）
        nparr = np.frombuffer(frame_data, np.uint8)
        # 解码为图像，返回与cv2.imread相同的格式（BGR通道的ndarray）
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        # 确保images文件夹存在
        if not os.path.exists('images'):
            os.makedirs('images')

        # 生成唯一的文件名，包含时间戳和客户端IP，避免文件名冲突
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"images/{self.client_ip.replace('.', '_')}_{timestamp}.jpg"

        try:
            # 保存图像到本地
            cv2.imwrite(filename, img)
            print(f"[{datetime.datetime.now():%H:%M:%S}] 图像已保存至：{filename}")

            # 进行检测
            if img is not None:
                has_violation, violation_type, details = detector.detect_violations(img)
                if has_violation:
                    print(f"检测到违规 - 类型: {violation_type}, 详情: {details}")
                    # 可以在这里添加发送检测结果回客户端的逻辑
                    await self.websocket.send_json({
                        "type": "detection_result",
                        "has_violation": has_violation,
                        "violation_type": violation_type,
                        "details": details,
                        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
                else:
                    print("未检测到任何违规内容")
            else:
                print(f"无法解析图像数据")
        except Exception as e:
            print(f"[{datetime.datetime.now():%H:%M:%S}] 图像处理错误 - {str(e)}")


# 全局连接管理（IP -> 连接实例）
connected_clients: Dict[str, ClientConnection] = {}
# 心跳任务（全局引用，用于关闭时清理）
heartbeat_task: Optional[asyncio.Task] = None


# -------------------------- 心跳检查逻辑 --------------------------
async def heartbeat_checker():
    while True:
        now = datetime.datetime.now()
        # 1. 筛选超时客户端（避免遍历中修改字典）
        timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]

        # 2. 处理超时连接（关闭+移除）
        if timeout_ips:
            print(f"[{now:%H:%M:%S}] 心跳检查：{len(timeout_ips)}个客户端超时（{timeout_ips}）")
            for ip in timeout_ips:
                try:
                    # 取消消费者任务
                    if connected_clients[ip].consumer_task and not connected_clients[ip].consumer_task.done():
                        connected_clients[ip].consumer_task.cancel()
                    await connected_clients[ip].websocket.close(code=1008, reason="心跳超时")
                finally:
                    connected_clients.pop(ip, None)
        else:
            print(f"[{now:%H:%M:%S}] 心跳检查：{len(connected_clients)}个客户端在线，无超时")

        # 3. 等待下一轮检查
        await asyncio.sleep(HEARTBEAT_INTERVAL)


# -------------------------- 应用生命周期 --------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    global heartbeat_task
    # 启动心跳任务
    heartbeat_task = asyncio.create_task(heartbeat_checker())
    print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务启动（ID：{id(heartbeat_task)}）")
    yield
    # 关闭时取消心跳任务
    if heartbeat_task and not heartbeat_task.done():
        heartbeat_task.cancel()
        try:
            await heartbeat_task
            print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务已取消")
        except asyncio.CancelledError:
            pass


# -------------------------- 消息处理 --------------------------
async def send_heartbeat_ack(client_ip: str):
    """回复心跳确认"""
    if client_ip not in connected_clients:
        return False
    try:
        ack = {
            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "type": "heartbeat"
        }
        await connected_clients[client_ip].websocket.send_json(ack)
        return True
    except Exception:
        connected_clients.pop(client_ip, None)
        return False


async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
    """处理文本消息（核心：心跳+JSON解析）"""
    try:
        msg = json.loads(text)
        # 仅处理心跳类型消息
        if msg.get("type") == "heartbeat":
            conn.update_heartbeat()
            await send_heartbeat_ack(client_ip)
        else:
            print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到消息：{msg}")
    except json.JSONDecodeError:
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：无效JSON消息")


async def handle_binary_msg(client_ip: str, data: bytes):
    """处理二进制消息（使用队列控制帧处理）"""
    if client_ip not in connected_clients:
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接不存在，丢弃{len(data)}字节数据")
        return

    conn = connected_clients[client_ip]

    # 检查队列是否已满
    if conn.frame_queue.full():
        # 队列已满，丢弃当前帧
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列已满，丢弃{len(data)}字节数据")
        return

    # 队列未满，添加帧到队列
    try:
        # 非阻塞添加（因为已检查队列未满，所以不会阻塞）
        conn.frame_queue.put_nowait(data)
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：已添加{len(data)}字节数据到队列")
    except asyncio.QueueFull:
        # 理论上不会走到这里，因为上面已检查
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列突然满了，丢弃{len(data)}字节数据")


# -------------------------- WebSocket核心端点 --------------------------
@ws_router.websocket(WS_ENDPOINT)
async def websocket_endpoint(websocket: WebSocket):
    # 接受连接 + 获取客户端IP
    await websocket.accept()
    client_ip = websocket.client.host if websocket.client else "unknown"
    now = datetime.datetime.now()
    print(f"[{now:%H:%M:%S}] 客户端{client_ip}：连接成功")

    consumer_task = None
    try:
        # 处理重复连接（关闭旧连接）
        if client_ip in connected_clients:
            # 取消旧连接的消费者任务
            if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done():
                connected_clients[client_ip].consumer_task.cancel()
            await connected_clients[client_ip].websocket.close(code=1008, reason="同一IP新连接")
            connected_clients.pop(client_ip)
            print(f"[{now:%H:%M:%S}] 客户端{client_ip}：关闭旧连接")

        # 注册新连接
        new_conn = ClientConnection(websocket, client_ip)
        connected_clients[client_ip] = new_conn

        # 启动帧消费任务
        consumer_task = new_conn.start_consumer()
        print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，已启动帧消费任务，当前在线{len(connected_clients)}个")

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

    # 异常处理（断开/错误）
    except WebSocketDisconnect as e:
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：主动断开（代码：{e.code}）")
    except Exception as e:
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接异常（{str(e)[:50]}）")
    finally:
        # 清理连接和任务
        if client_ip in connected_clients:
            # 取消消费者任务
            if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done():
                connected_clients[client_ip].consumer_task.cancel()
            connected_clients.pop(client_ip, None)
        print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个")