CameraMachine/python/sockets/socket_server.py

import json
import asyncio
from models import *
from .connect_manager import ConnectionManager
from .message_handler import *
from mcu.DeviceControl import DeviceControl, checkMcuConnection
from mcu.BlueToothMode import BlueToothMode
from mcu.capture.smart_shooter_class import SmartShooter
import time
from .socket_client import socket_manager
from sqlalchemy.exc import NoResultFound
import os, datetime
import traceback
import logging
from utils import common
from utils.common import message_queue

logger = logging.getLogger(__name__)
conn_manager = ConnectionManager()
active_connections = set()
device_ctrl = DeviceControl(websocket_manager=conn_manager)
blue_tooth = BlueToothMode(websocket_manager=conn_manager)
smart_shooter = SmartShooter(websocket_manager=conn_manager)


async def updateDataRecord(PhotoFilename, id):
    await asyncio.sleep(0.01)
    create_time = datetime.datetime.fromtimestamp(os.path.getctime(PhotoFilename))
    data = {"id": id, "image_path": PhotoFilename, "photo_create_time": create_time}
    # record_model = PhotoRecord(**data)
    session = SqlQuery()
    record_model = CRUD(PhotoRecord)
    model = record_model.read(session, conditions={"id": id, "delete_time": None})
    if model == None:
        print(f"smart shooter 拍照记录更新失败,记录id:{id},不存在")
        session.close()
        return None
    else:
        # 走编辑逻辑
        settings.syncPhotoRecord(data, action_type=3)
        record_data = record_model.updateConditions(session, conditions={"id": id}, **data)
        session.close()
        print(f"smart shooter 拍照记录更新成功,记录id:{id}")
        return record_data


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # await websocket.accept()
    main_loop = asyncio.get_running_loop()
    smart_shooter.main_loop = main_loop  # <--- 添加这一行
    await conn_manager.connect(websocket)
    active_connections.add(websocket)
    smart_shooter.websocket = websocket
    device_ctrl.websocket = websocket
    blue_tooth.websocket = websocket
    common.websocket_manager = conn_manager
    common.websocket = websocket
    # 启动 smart_shooter.connect_listen 服务
    listen_task = None
    tasks = set()
    send_task = None  # <--- 新增
    try:
        # 初始化回调函数
        smart_shooter.callback_listen = MsgCallback
        # 创建任务来并发处理不同类型的消息
        handler_task = asyncio.create_task(handler_messages(websocket))
        # send_task = asyncio.create_task(send_message(websocket))
        send_task = asyncio.create_task(send_message(websocket))  # <--- 启动消费者
        loop = asyncio.get_event_loop()
        listen_task = loop.run_in_executor(None, smart_shooter.connect_listen)
        # send_task = loop.run_in_executor(None, send_message(websocket))
        # 创建任务来启动 connect_listen
        # listen_task = asyncio.create_task(restart_smart_shooter_listener())
        # 等待所有任务完成
        await asyncio.gather(handler_task, listen_task, return_exceptions=True)

    except WebSocketDisconnect:
        print("Client disconnected")
    except asyncio.CancelledError:
        print("Connection cancelled")
    finally:
        # 确保任务被正确取消和清理
        tasks_to_cancel = []

        if handler_task and not handler_task.done():
            tasks_to_cancel.append(handler_task)

        if send_task and not send_task.done():
            tasks_to_cancel.append(send_task)

        # 取消所有待处理的任务
        for task in tasks_to_cancel:
            task.cancel()

        # 等待任务取消完成
        if tasks_to_cancel:
            await asyncio.gather(*tasks_to_cancel, return_exceptions=True)

        # 清理连接
        active_connections.discard(websocket)
        print("WebSocket connection cleaned up")


async def start_smart_shooter_listen():
    """启动 smart_shooter 监听服务"""
    loop = asyncio.get_event_loop()
    # 在执行器中运行 connect_listen 方法
    try:
        loop.create_task(None, smart_shooter.connect_listen())
    except Exception as e:
        print(f"Smart shooter listen error: {e}")


async def handler_messages(websocket):
    while True:
        try:
            byteDats = await websocket.receive()
            socket_type = byteDats.get("type")
            if socket_type == "websocket.disconnect":
                print("socket_type===>", byteDats)
                if byteDats.get("code") == 1006:
                    continue
                logger.info("socket强制断开")
                diviceList = blue_tooth.devices
                if len(diviceList) == 0:
                    blue_tooth.bluetooth_exit = True
                    blue_tooth.clearMyInstance()
                    break
                diviceAddress = (
                    ""
                    if len(list(diviceList.keys())) == 0
                    else list(diviceList.keys())[0]
                )
                if diviceAddress != "":
                    print(diviceList.get(diviceAddress))
                    diviceName = diviceList[diviceAddress]["name"]
                    blue_tooth.disconnect_device(diviceAddress, diviceName)
                blue_tooth.bluetooth_exit = True
                blue_tooth.clearMyInstance()
                print("所有设备已断开连接")
                break
            print("byteDats", byteDats)
            # 使用create_task来避免阻塞消息处理循环
            asyncio.create_task(
                handlerSend(
                    conn_manager, json.dumps(byteDats), websocket, smart_shooter
                )
            )
        except Exception as e:
            print("socket error", e)
            break


async def message_generator():
    """异步生成器，用于从队列中获取消息"""
    while True:
        try:
            # 使用asyncio.wait_for设置合理的超时时间
            message = await message_queue.get()
            yield message
        except asyncio.TimeoutError:
            # 超时继续，允许其他协程运行
            await asyncio.sleep(0.01)
            # print("超时继续，允许其他协程运行")
            continue
        except Exception as e:
            print("消息生成器错误", e)
            break


async def send_message(websocket):
    """使用异步生成器发送消息"""
    print("构建消息监听   send_message")
    while True:
        try:
            # 1. 异步等待消息，不会阻塞其他协程
            data = await message_queue.get()

            # 2. 发送消息
            if common.websocket_manager:
                # 假设 broadcast 或 send_personal_message 是异步方法
                await common.websocket_manager.send_personal_message(data, common.websocket)
                # 或者根据你的 ConnectionManager 实现调用具体发送方法

            message_queue.task_done()
        except Exception as e:
            print(f"消息消费错误: {e}")
            await asyncio.sleep(1)  # 防止死循环报错


async def getActionInfo(record_info):
    await asyncio.sleep(0.01)
    if not record_info:
        return None
    action_id = record_info.action_id
    if not action_id:
        return None
    session = SqlQuery()
    device_model = CRUD(DeviceConfig)
    model = device_model.read(session, conditions={"id": action_id})
    if not model:
        session.close()
        return None
    return model


async def MsgCallback(msg):
    msg_id = msg.get("msg_id")
    match msg_id:
        case "PhotoUpdated":
            PhotoFilename = msg.get("PhotoFilename")
            PhotoLocation = msg.get("PhotoLocation")
            PhotoOrigin = msg.get("PhotoOrigin")
            if (PhotoFilename != "" and PhotoFilename != None) and (
                    PhotoLocation == "Local Disk"
            ):
                # temp_photo_name = PhotoFilename
                # 更新拍照记录
                logger.info(f"PhotoUpdated,{PhotoFilename}--{PhotoOrigin}")
                goods_art_no = None
                id = None
                try:
                    if PhotoOrigin != "" and PhotoOrigin not in ["external", "ui"]:
                        goods_art_no, id = PhotoOrigin.split(",")
                    # 创建任务来处理数据库更新，避免阻塞回调
                    recordResult = await updateDataRecord(PhotoFilename, id)
                except Exception as e:
                    print("拍照更新异常", e)
                    recordResult = None
                actionModel = None
                if recordResult:
                    actionModel = await getActionInfo(recordResult)
                data = conn_manager.jsonMessage(
                    code=0,
                    msg=f"照片获取成功",
                    data={
                        "photo_file_name": PhotoFilename,
                        "goods_art_no": goods_art_no,
                        "action_name": actionModel.action_name if actionModel else None,
                    },
                    msg_type="smart_shooter_photo_take",
                )
                await conn_manager.send_personal_message(data, smart_shooter.websocket)
        case "LiveviewUpdated":
            CameraLiveviewImage = msg.get("CameraLiveviewImage", None)
            # base64_to_image(CameraLiveviewImage, "liveview.jpg")
            # print("收到直播画面：CameraLiveviewImage")
            data = conn_manager.jsonMessage(
                code=1,
                msg=f"预览数据发送",
                data={"smart_shooter_preview": CameraLiveviewImage},
                msg_type="smart_shooter_enable_preview",
            )
            await conn_manager.send_personal_message(data, smart_shooter.websocket)
        # case _:
        #     print("收到未知数据：{}".format(msg))


# @app.on_event("startup")
# async def startup_event():
#     loop = asyncio.get_event_loop()
#     loop.run_in_executor(None, await smart_shooter.connect_listen)
#     print("监听服务已启动")


@app.on_event("shutdown")
async def shutdown_event():
    print("Shutting down...")
    # socket_manager.close()
    # 清理操作
    for connection in list(active_connections):
        try:
            await connection.close()
        except Exception as e:
            print(f"Error closing connection: {e}")
    smart_shooter.stop_listen = True
    smart_shooter.is_init_while = False
    device_ctrl.close_connect()
    device_ctrl.close_lineConnect()
    device_ctrl.mcu_exit = True
    device_ctrl.p_list = []
    device_ctrl.temp_ports_dict = {}
    device_ctrl.clearMyInstance()
    diviceList = blue_tooth.devices
    if len(diviceList) == 0:
        blue_tooth.bluetooth_exit = True
        blue_tooth.clearMyInstance()
    diviceAddress = "" if len(list(diviceList.keys())) == 0 else list(diviceList.keys())[0]
    if diviceAddress != "":
        print(diviceList.get(diviceAddress))
        diviceName = diviceList[diviceAddress]["name"]
        blue_tooth.disconnect_device(diviceAddress, diviceName)
    blue_tooth.bluetooth_exit = True
    blue_tooth.clearMyInstance()
    print("所有设备已断开连接")