# -*- coding = utf-8 -*-
# @Time : 2025/3/1 下午8:47
# @File : managers.py
# @Software : PyCharm
# @Author : xingc
# @Desc :
import asyncio
import logging

import redis.asyncio as aioredis
import ujson
from fastapi import WebSocket
from fastapi.websockets import WebSocketState

from core.db.engine import get_cache

logger = logging.getLogger(__name__)


class RedisPubSubManager:
    def __init__(self):
        self.redis_connection = None
        self.pubsub = None

    async def _get_redis_connection(self) -> aioredis:
        return await get_cache()

    async def connect(self):
        self.redis_connection = await self._get_redis_connection()
        self.pubsub = self.redis_connection.pubsub()

    async def subscribe(self, chat_guid: str) -> aioredis:
        await self.pubsub.subscribe(chat_guid)
        return self.pubsub

    async def unsubscribe(self, chat_guid: str):
        await self.pubsub.unsubscribe(chat_guid)

    async def publish(self, chat_guid: str, message: str):
        await self.redis_connection.publish(chat_guid, message)

    async def disconnect(self):
        if self.redis_connection is not None:
            await self.redis_connection.close()


class SocketManager:

    def __init__(self):
        self.handlers: dict = {}
        self.chats: dict = {}
        self.pubsub_client = RedisPubSubManager()
        self.user_guid_to_websocket: dict = {}
        self._task_cancel = True

    def handler(self, message_type):
        def decorator(func):
            self.handlers[message_type] = func
            return func

        return decorator

    async def connect_socket(self, websocket: WebSocket):
        await websocket.accept()

    async def add_user_socket_connection(self, user_guid: str, websocket: WebSocket):
        self.user_guid_to_websocket.setdefault(user_guid, set()).add(websocket)

    async def add_user_to_chat(self, chat_guid: str, websocket: WebSocket):
        if chat_guid in self.chats:
            self.chats[chat_guid].add(websocket)
        else:
            self.chats[chat_guid] = {websocket}
            await self.pubsub_client.connect()
            pubsub_subscriber = await self.pubsub_client.subscribe(chat_guid)
            asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))

    async def broadcast_to_chat(self, chat_guid: str, message: str | dict) -> None:
        if isinstance(message, dict):
            message = ujson.dumps(message)
        await self.pubsub_client.publish(chat_guid, message)

    async def remove_user_from_chat(self, chat_guid: str, websocket: WebSocket) -> None:
        self.chats[chat_guid].remove(websocket)
        if len(self.chats[chat_guid]) == 0:
            del self.chats[chat_guid]
            logger.info(f'Removing user from PubSub channel {chat_guid}')
            await self.pubsub_client.unsubscribe(chat_guid)

    async def remove_user_guid_to_websocket(self, user_guid: str, websocket: WebSocket):
        if user_guid in self.user_guid_to_websocket:
            self.user_guid_to_websocket.get(user_guid).remove(websocket)

    # https://github.com/redis/redis-py/issues/2523
    async def _pubsub_data_reader(self, pubsub_subscriber):
        try:
            while True:
                message = await pubsub_subscriber.get_message(ignore_subscribe_messages=True)
                if message is not None:
                    chat_guid = message['channel']
                    sockets = self.chats.get(chat_guid)
                    if sockets:
                        for socket in sockets:
                            if socket.client_state != WebSocketState.CONNECTED:
                                continue
                            data = message['data']
                            await socket.send_text(data)
                else:
                    if self._task_cancel is False:
                        break
        except Exception as exc:
            logger.exception(f'Exception occurred: {exc}')

    async def send_error(self, message: str, websocket: WebSocket):
        if websocket.client_state == WebSocketState.CONNECTED:
            await websocket.send_json({'status': 'error', 'content': message})

    async def quit(self):
        self._task_cancel = False
        await self.pubsub_client.disconnect()