import asyncio
import dataclasses
import logging
from pathlib import Path
from threading import Thread

from fastapi import FastAPI, Response
from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, RobotRtsiMsg, SetSidMsg
from flandre.utils.RfSequence import RfSequence

logger = logging.getLogger(__name__)

import uvicorn

from typing import List

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, APIRouter

app = FastAPI()


# class ConnectionManager:
#     def __init__(self):
#         # 存放激活的ws连接对象
#         self.active_connections: List[WebSocket] = []
#         self.router = APIRouter()
#         self.router.add_api_route("/hello", self.hello, methods=["GET"])
#
#
#     def hello(self):
#         return {"Hello": self.name}

#     async def connect(self, ws: WebSocket):
#         # 等待连接
#         await ws.accept()
#         # 存储ws连接对象
#         self.active_connections.append(ws)
#
#     def disconnect(self, ws: WebSocket):
#         # 关闭时 移除ws对象
#         self.active_connections.remove(ws)
#
#     @staticmethod
#     async def send_personal_message(message: str, ws: WebSocket):
#         # 发送个人消息
#         await ws.send_text(message)
#
#     async def broadcast(self, message: str):
#         # 广播消息
#         for connection in self.active_connections:
#             await connection.send_text(message)
#
#
# manager = ConnectionManager()
#
#
# @app.websocket("/ws/{user}")
# async def websocket_endpoint(websocket: WebSocket, user: str):
#     await manager.connect(websocket)
#
#     await manager.broadcast(f"用户{user}进入聊天室")
#
#     try:
#         while True:
#             data = await websocket.receive_text()
#             await manager.send_personal_message(f"你说了: {data}", websocket)
#             await manager.broadcast(f"用户:{user} 说: {data}")
#
#     except WebSocketDisconnect:
#         manager.disconnect(websocket)
# await manager.broadcast(f"用户-{user}-离开")


class Web(Node):
    topics = [ImageArgMsg, RobotRtsiMsg, SetSidMsg]

    def __init__(self, level=logging.INFO):
        super(Web, self).__init__(level=level)
        self.wss: list[WebSocket] = []
        self.arg = ImageArgMsg(sender='frontend', t_end=123)

    def custom_setup(self):
        self.router = APIRouter()
        self.router.add_api_route("/hello", self.hello, methods=["GET"])
        self.router.add_websocket_route("/ws", self.websocket_endpoint)
        self.router.add_api_route("/p", self.get_image, response_class=Response, responses={
            200: {
                "content": {"image/png": {}}
            }
        })

    def hello(self):
        return {"Hello": 'asd'}

    def get_image(self, i: int):
        rfs = RfSequence()
        f = rfs.frames[i]

        image_bytes: bytes = b''
        return Response(content=image_bytes, media_type="image/png")

    async def websocket_endpoint(self, websocket: WebSocket):
        await websocket.accept()
        self.wss.append(websocket)
        try:
            while True:
                data = await websocket.receive_json()
                if data['type'] == 'ImageArgMsg':
                    t = dataclasses.replace(self.arg)
                    t.sender = data['sender']
                    t.t_start = data['t_start']
                    t.g8 = data['g8']
                    self.send(t)
        except WebSocketDisconnect:
            self.wss.remove(websocket)

    def boardcast(self, data: dict):
        try:
            for ws in self.wss:
                asyncio.run(ws.send_json(data))
        except Exception as e:
            logger.warning(e)

    def wst(self):
        app = FastAPI()
        app.include_router(self.router)
        uvicorn.run(app=app, host="0.0.0.0", port=8010)

    def loop(self):
        Thread(target=self.wst).start()
        while True:
            msg = self.recv()
            if isinstance(msg, ImageArgMsg):
                if msg.sender != 'frontend':
                    self.arg = msg
                    self.boardcast(msg.dict)
            elif isinstance(msg, RobotRtsiMsg):
                self.boardcast(msg.dict)
            elif isinstance(msg, SetSidMsg):
                self.boardcast(msg.dict)


if __name__ == '__main__':
    Web()()
