from fastapi import FastAPI, WebSocket, Depends, WebSocketDisconnect, WebSocketException
from fastapi.responses import HTMLResponse

from extends.redis.conn import redis_connection
from config.settings import get_settings, Settings
from typing import TypedDict, Literal


from contextlib import asynccontextmanager
from asyncio import sleep
import asyncio

import logging
import redis
import orjson

"""
    websocket server
"""



@asynccontextmanager
async def lifespan(app: FastAPI):
    yield
    redis_connection.disconnect_redis_connection()

# 设置日志格式和级别
logging.basicConfig(
    level=logging.INFO,  # 可以根据需要设置为DEBUG、WARNING、ERROR等
    format="[%(name)s]-[%(levelname)s]: %(message)s",
)

# 获取一个logger实例，通常与您的应用或模块名称相关联
logger = logging.getLogger("websocket")

settings: Settings = get_settings()
logger.debug(settings)

app = FastAPI(lifespan=lifespan)

active_connections = set()
stop_flags: dict[WebSocket, asyncio.Event] = {}
task_running_flags: dict[WebSocket, asyncio.Event] = {}

html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat</title>
    </head>
    <body>
        <h1>hhs data send to console</h1>
        <script>
            var ws = new WebSocket("ws://localhost:{}");
            // 接受服务端数据
            ws.onopen = function(event) {{
                ws.send("get_status");
                console.log("Sent request for data");
            }};
            ws.onmessage = function(event) {{
                console.log(JSON.parse(event.data))
            }};
            ws.onclose = function(event) {{
                ws.send("stop_get_status");
                console.log("Sent stop request");
            }};
        </script>
    </body>
</html>
""".format(settings.WEBSOCKET_PORT)


@app.get("/ws")
async def get():
    return HTMLResponse(html)


class InitData(TypedDict):
    current_result: dict
    previous_result: dict
    is_getting_status: bool


async def handle_send_data(websocket: WebSocket, redis_conn: redis.Redis, stop_flag: asyncio.Event, task_running: asyncio.Event):
    """处理发送给客户端的数据"""
    current_result = {}
    previous_result = {}
    while not stop_flag.is_set():
        if task_running.is_set():  
            data = redis_conn.get(settings.DATA_RESULT_KEY)
            if data:
                current_result = orjson.loads(data)
                if current_result != previous_result:
                    await websocket.send_json(current_result)  
                    previous_result = current_result
            await asyncio.sleep(1)
        else:  
            await task_running.wait()


@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket, redis_conn: redis.Redis = Depends(redis_connection.get_redis_connection)):
    active_connections.add(websocket)
    stop_flags[websocket] = asyncio.Event()
    task_running_flags[websocket] = asyncio.Event()  

    await websocket.accept()
    send_data_task = None

    try:
        while True:
            data = await websocket.receive_text()
            if data == "get_status":
                if not task_running_flags[websocket].is_set():  
                    stop_flags[websocket].clear()  
                    task_running_flags[websocket].set()  
                    send_data_task = asyncio.create_task(handle_send_data(websocket, redis_conn, stop_flags[websocket], task_running_flags[websocket]))
            elif data == "stop_get_status":
                stop_flags[websocket].set()  
                task_running_flags[websocket].clear()  
            elif data == "close":
                stop_flags[websocket].set()
                task_running_flags[websocket].clear()
                await websocket.close()
                break
            # 休眠的时间
            await asyncio.sleep(settings.GET_DATA_FREQUENCY)

    except (WebSocketDisconnect, WebSocketException):
        logger.info("websocket disconnect")
    except Exception as e:
        logger.error(e)
    finally:
        if send_data_task and not send_data_task.done():  
            send_data_task.cancel()
        active_connections.remove(websocket)
        del stop_flags[websocket]
        del task_running_flags[websocket]  

    return None


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=settings.WEBSOCKET_HOST, port=settings.WEBSOCKET_PORT)