From c663cdad31a67400dc88628b17669ddedca2b7ea Mon Sep 17 00:00:00 2001
From: xyj <10908227994@qq.com>
Date: Sat, 6 Jan 2024 15:15:52 +0800
Subject: [PATCH] update

---
 video_task.py | 57 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 43 insertions(+), 14 deletions(-)

diff --git a/video_task.py b/video_task.py
index 133854e..2714fec 100644
--- a/video_task.py
+++ b/video_task.py
@@ -7,6 +7,17 @@ from pydantic import BaseModel
 
 broker = 'mqtt.lihaink.cn'
 port = 1883
+APP = "app"
+SCREEN = "screen"
+WEB = "web"
+SCENE_NAME = "scene"
+user_sched = {}
+user_msg = {}
+user_scene = {
+    SCREEN: False,
+    APP: False,
+    WEB: False
+}
 
 
 class BaseResponse(BaseModel):
@@ -41,9 +52,6 @@ class MQTTClient:
         self.client.connect(self.broker, self.port)
 
 
-user_sched = {}
-
-
 async def close(username, device):
     # 倒计时600秒
     await asyncio.sleep(30)
@@ -115,27 +123,48 @@ async def function_B(client, data):
     await stop(username, device)
 
 
-user_msg = {}
-user_scene = [False, False, False]
-
 # 创建WebSocket连接的处理函数
 async def handler(websocket, path):
+    scene = None
     try:
         # 在循环中等待客户端的消息
         async for message in websocket:
             user_msg[websocket] = message
+            data = json.loads(message)
+            scene = data[SCENE_NAME]
+            # 如果该场景已经开启，那么不在开启了
+            if user_scene[scene] is True:
+                await websocket.send(json.dumps({"code": 200, "msg": scene + "该场景已经开启了，不要再请求我了！"}, ensure_ascii=False))
+                continue
+            if user_scene[APP] is True or user_scene[WEB] is True or user_scene[SCREEN] is True:
+                user_scene[scene] = True
+                await websocket.send(json.dumps({"code": 200, "msg": "已经有其他场景开启了，不需要再推流了！"}, ensure_ascii=False))
+                continue
+            # 如果该场景没有开启，那么进行开启
+            user_scene[scene] = True
             # 当接收到消息时调用函数A
             response = await function_A(websocket, message)
             await websocket.send(response)
-        # 本地情况
-        await function_B(websocket, message)
+        # 没有消息了，则首先关闭该场景
+        if scene is not None:
+            user_scene[scene] = False
+        # 查看所有场景是否在线,只有都不在线就关闭推流
+        if user_scene[APP] is False and user_scene[WEB] is False and user_scene[SCREEN] is False:
+            await websocket.send(json.dumps({"code": 200, "msg": "开始关闭场景" + scene}, ensure_ascii=False))
+            await function_B(websocket, user_msg[websocket])
     except Exception as e:
-        msg = user_msg[websocket]
-        await function_B(websocket, msg)
+        data = json.loads(user_msg[websocket])
+        scene = data[SCENE_NAME]
+        if scene is not None:
+            user_scene[scene] = False
+        # 查看所有场景是否在线,只有都不在线就关闭推流
+        if user_scene[APP] is False and user_scene[WEB] is False and user_scene[SCREEN] is False:
+            await function_B(websocket, user_msg[websocket])
 
 
-# 启动WebSocket服务器
-start_server = websockets.serve(handler, "0.0.0.0", 8765)
+if __name__ == '__main__':
+    # 启动WebSocket服务器
+    start_server = websockets.serve(handler, "0.0.0.0", 8765)
 
-asyncio.get_event_loop().run_until_complete(start_server)
-asyncio.get_event_loop().run_forever()
\ No newline at end of file
+    asyncio.get_event_loop().run_until_complete(start_server)
+    asyncio.get_event_loop().run_forever()