video/ws/ws.py

import asyncio
import datetime
import json
import os
from contextlib import asynccontextmanager
from typing import Dict, Optional, AsyncGenerator
from concurrent.futures import ThreadPoolExecutor  # 新增：显式线程池

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

import cv2
import numpy as np
from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
from queue import Queue  # 线程安全队列，无需额外Lock

from ocr.model_violation_detector import MultiModelViolationDetector

# -------------------------- 配置调整 --------------------------
# 模型路径（建议改为环境变量）
YOLO_MODEL_PATH = r"D:\Git\bin\video\ocr\models\best.pt"
OCR_CONFIG_PATH = r"D:\Git\bin\video\ocr\config\1.yaml"

# 核心优化：模型池大小（决定最大并发任务数，显存占用=大小×单模型显存）
MODEL_POOL_SIZE = 5  # 示例：设为5，支持5个任务并行，显存会明显上升
THREAD_POOL_SIZE = MODEL_POOL_SIZE * 2  # 线程池大小≥模型池，避免线程瓶颈

# 其他配置
HEARTBEAT_INTERVAL = 30  # 心跳间隔（秒）
HEARTBEAT_TIMEOUT = 600  # 客户端超时阈值（秒）
WS_ENDPOINT = "/ws"  # WebSocket端点
FRAME_QUEUE_SIZE = 5  # 增大帧队列，允许缓存更多帧（避免丢帧）

# -------------------------- 工具函数 --------------------------
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")

# -------------------------- 模型池重构（核心修改1） --------------------------
class ModelPool:
    def __init__(self, pool_size: int = MODEL_POOL_SIZE):
        self.pool = Queue(maxsize=pool_size)
        # 移除冗余Lock：Queue.get()/put()本身线程安全
        self._init_models(pool_size)
        print(f"[{get_current_time_str()}] 模型池初始化完成（共{pool_size}个实例，显存已预分配）")

    def _init_models(self, pool_size: int):
        """预加载所有模型实例（初始化时显存会一次性上升）"""
        for i in range(pool_size):
            try:
                detector = MultiModelViolationDetector(
                    ocr_config_path=OCR_CONFIG_PATH,
                    yolo_model_path=YOLO_MODEL_PATH,
                    ocr_confidence_threshold=0.5
                )
                self.pool.put(detector)
                print(f"[{get_current_time_str()}] 模型实例{i+1}/{pool_size}加载完成")
            except Exception as e:
                raise RuntimeError(f"模型实例{i+1}加载失败：{str(e)}")

    def get_model(self) -> MultiModelViolationDetector:
        """获取模型（阻塞直到有空闲实例，确保并发安全）"""
        return self.pool.get()

    def return_model(self, detector: MultiModelViolationDetector):
        """归还模型（立即释放资源供其他任务使用）"""
        self.pool.put(detector)

# -------------------------- 全局资源初始化 --------------------------
model_pool = ModelPool(pool_size=MODEL_POOL_SIZE)  # 初始化模型池（预占显存）
thread_pool = ThreadPoolExecutor(  # 显式创建线程池（核心修改2）
    max_workers=THREAD_POOL_SIZE,
    thread_name_prefix="ModelWorker-"  # 线程命名，便于调试
)

# -------------------------- 客户端连接封装（核心修改3） --------------------------
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
        # 移除“客户端独占模型”：不再持有detector属性

    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:
            await self.websocket.send_json({
                "type": "frame",
                "timestamp": get_current_time_str(),
                "client_ip": 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()
                # 2. 立即发送下一帧许可（让客户端持续发帧，积累并发任务）
                await self.send_frame_permit()
                try:
                    # 3. 并行处理帧（核心：任务级借模型）
                    await self.process_frame(frame_data)
                finally:
                    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:
        """处理单帧（核心修改4：任务级借还模型）"""
        # 1. 临时借用模型（阻塞直到有空闲实例，显存随借用数上升）
        detector = model_pool.get_model()
        try:
            # 2. 二进制转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

            # 3. 保存图像（可选）
            os.makedirs('images', exist_ok=True)
            filename = f"images/{self.client_ip.replace('.', '_')}_{get_current_time_file_str()}.jpg"
            cv2.imwrite(filename, img)

            # 4. 显式线程池执行AI检测（真正并发，无线程瓶颈）
            loop = asyncio.get_running_loop()
            has_violation, violation_type, details = await loop.run_in_executor(
                thread_pool,  # 用自定义线程池，避免默认线程不足
                detector.detect_violations,  # 临时借用的模型
                img  # 输入图像
            )

            # 5. 违规处理（与原逻辑一致）
            if has_violation:
                print(f"[{get_current_time_str()}] 客户端{self.client_ip}：违规 - {violation_type}")
                # 违规次数更新（用线程池避免阻塞事件循环）
                await loop.run_in_executor(thread_pool, increment_alarm_count_by_ip, self.client_ip)
                # 发送危险通知
                await self.websocket.send_json({
                    "type": "danger",
                    "timestamp": get_current_time_str(),
                    "client_ip": self.client_ip,
                    "violation_type": violation_type,
                    "details": details
                })
            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)}")
        finally:
            # 6. 无论成功/失败，强制归还模型（核心：释放资源供其他任务使用）
            model_pool.return_model(detector)
            print(f"[{get_current_time_str()}] 客户端{self.client_ip}：模型已归还（可复用）")

# -------------------------- 全局状态与心跳 --------------------------
connected_clients: Dict[str, ClientConnection] = {}
client_lock = asyncio.Lock()  # 保护客户端字典的异步锁
heartbeat_task: Optional[asyncio.Task] = None

async def heartbeat_checker():
    """心跳检查（移除模型归还逻辑，因模型已任务级归还）"""
    while True:
        current_time = get_current_time_str()
        async with client_lock:
            # 筛选超时客户端
            timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]

        for ip in timeout_ips:
            async with client_lock:
                conn = connected_clients.get(ip)
                if not conn:
                    continue
                # 取消消费任务+关闭连接
                if conn.consumer_task and not conn.consumer_task.done():
                    conn.consumer_task.cancel()
                await conn.websocket.close(code=1008, reason="心跳超时")
                # 标记离线（用线程池）
                loop = asyncio.get_running_loop()
                await loop.run_in_executor(thread_pool, update_online_status_by_ip, ip, 0)
                await loop.run_in_executor(
                    thread_pool, add_device_action, DeviceActionCreate(client_ip=ip, action=0)
                )
                connected_clients.pop(ip)
                print(f"[{current_time}] 客户端{ip}：超时离线（资源已清理）")

        # 打印在线状态
        async with client_lock:
            print(f"[{current_time}] 心跳检查：{len(connected_clients)}个客户端在线")
        await asyncio.sleep(HEARTBEAT_INTERVAL)

# -------------------------- 应用生命周期（核心修改5：管理线程池） --------------------------
@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)}）")
    print(f"[{get_current_time_str()}] 线程池启动（最大线程数：{THREAD_POOL_SIZE}）")
    yield  # 应用运行期间
    # 清理资源
    if heartbeat_task and not heartbeat_task.done():
        heartbeat_task.cancel()
        await heartbeat_task
        print(f"[{get_current_time_str()}] 心跳任务已关闭")
    # 关闭线程池（等待所有任务完成）
    thread_pool.shutdown(wait=True)
    print(f"[{get_current_time_str()}] 线程池已关闭")

# -------------------------- WebSocket路由 --------------------------
ws_router = APIRouter()

@ws_router.websocket(WS_ENDPOINT)
async def websocket_endpoint(websocket: WebSocket):
    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}：连接建立")

    new_conn = None
    is_online_updated = False
    try:
        # 处理重复连接（关闭旧连接）
        async with client_lock:
            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="新连接抢占")
                connected_clients.pop(client_ip)
                print(f"[{current_time}] 客户端{client_ip}：旧连接已关闭")

        # 创建新连接+启动消费任务
        new_conn = ClientConnection(websocket, client_ip)
        new_conn.start_consumer()
        # 初始发送帧许可（让客户端立即发帧）
        await new_conn.send_frame_permit()

        # 标记客户端在线
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(thread_pool, update_online_status_by_ip, client_ip, 1)
        await loop.run_in_executor(
            thread_pool, add_device_action, DeviceActionCreate(client_ip=client_ip, action=1)
        )
        is_online_updated = True
        async with client_lock:
            connected_clients[client_ip] = new_conn
        print(f"[{current_time}] 客户端{client_ip}：注册成功（在线数：{len(connected_clients)}）")

        # 消息循环（接收文本/二进制帧）
        while True:
            data = await websocket.receive()
            if "text" in data:
                # 处理文本消息（如心跳）
                try:
                    msg = json.loads(data["text"])
                    if msg.get("type") == "heart":
                        new_conn.update_heartbeat()
                        # 回复心跳确认
                        await websocket.send_json({
                            "type": "heart",
                            "timestamp": get_current_time_str(),
                            "client_ip": client_ip
                        })
                except json.JSONDecodeError:
                    print(f"[{get_current_time_str()}] 客户端{client_ip}：无效JSON")
            elif "bytes" in data:
                # 处理二进制帧（图像）
                try:
                    await new_conn.frame_queue.put(data["bytes"])
                    print(f"[{get_current_time_str()}] 客户端{client_ip}：帧已入队（队列大小：{new_conn.frame_queue.qsize()}）")
                except asyncio.QueueFull:
                    print(f"[{get_current_time_str()}] 客户端{client_ip}：帧队列满（丢弃当前帧）")

    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:
        # 清理资源（无需归还模型，已在process_frame中归还）
        if new_conn and client_ip in connected_clients:
            async with client_lock:
                conn = connected_clients.get(client_ip)
                if conn:
                    if conn.consumer_task and not conn.consumer_task.done():
                        conn.consumer_task.cancel()
                    # 标记离线（仅当在线状态已更新时）
                    if is_online_updated:
                        loop = asyncio.get_running_loop()
                        await loop.run_in_executor(thread_pool, update_online_status_by_ip, client_ip, 0)
                        await loop.run_in_executor(
                            thread_pool, add_device_action, DeviceActionCreate(client_ip=client_ip, action=0)
                        )
                    connected_clients.pop(client_ip)
        async with client_lock:
            print(f"[{get_current_time_str()}] 客户端{client_ip}：资源清理完成（在线数：{len(connected_clients)}）")