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  # 帧队列大小限制（保持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 send_allow_send_frame(self):
        """向客户端发送JSON信号，通知其可发送下一帧二进制数据"""
        try:
            allow_msg = {
                "type": "allow_send_frame",  # 信号类型，与客户端约定
                "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "status": "ready",  # 表示服务器已准备好接收下一帧
                "client_ip": self.client_ip  # 可选：便于客户端确认自身身份
            }
            await self.websocket.send_json(allow_msg)
            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 consume_frames(self) -> None:
        """从队列中获取帧并进行处理，处理完后通知客户端可发送下一帧"""
        try:
            while True:
                # 从队列获取帧数据（队列空时会阻塞，等待客户端发送）
                frame_data = await self.frame_queue.get()
                try:
                    # 处理帧数据
                    await self.process_frame(frame_data)
                finally:
                    # 标记任务完成（队列计数-1，此时队列回到空状态）
                    self.frame_queue.task_done()
                    # ---------- 修改：处理完当前帧后，立即通知客户端可发送下一帧 ----------
                    await self.send_allow_send_frame()
        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(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：未检测到任何违规内容")
            else:
                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)}")


# 全局连接管理（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()
        # ---------- 修改：客户端刚连接时，队列空，立即发送「允许发送帧」信号 ----------
        await new_conn.send_allow_send_frame()
        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)}个")