# -*- coding: utf-8 -*-
"""
@File    : ws_connection_manager.py
@Time    : 2023/9/12 17:09
@Author  : geekbing
@LastEditTime : -
@LastEditors : -
@Description : -
"""
import asyncio
import logging
import json
from typing import TypeVar

from channels.generic.websocket import AsyncWebsocketConsumer

logger = logging.getLogger(__name__)

MsgType = TypeVar("MsgType", str, dict, bytes)


class ConnectionManager(AsyncWebsocketConsumer):
    active_connections = {}

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.user_id = None  # 初始化 user_id 属性

    async def connect(self):
        self.user_id = int(self.scope["url_route"]["kwargs"]["user_id"])
        exist = self.active_connections.get(self.user_id)
        if exist:
            await self.close()
        else:
            self.active_connections[self.user_id] = self
            logger.debug(f"websocket: 用户[{self.user_id}]建立连接成功！")
        await self.accept()
        # asyncio.create_task(self.send_heartbeat())

    async def disconnect(self, close_code):
        if self.user_id is not None:
            del self.active_connections[self.user_id]
            logger.debug(f"websocket: 用户[{self.user_id}] 已安全断开！")

    async def receive(self, text_data=None, bytes_data=None):
        if text_data:
            data = json.loads(text_data)
            message_type = data.get("type")
            if message_type in self.questions_and_answers_map:
                response = self.questions_and_answers_map[message_type].format(
                    user_id=self.user_id
                )  # 格式化字符串以包含 user_id
                await self.send_personal_message(
                    self.user_id, response
                )  # 使用 self.user_id
            else:
                # 添加其他消息类型的处理逻辑
                pass
        elif bytes_data:
            # 在这里添加你的二进制数据处理逻辑
            pass

    @staticmethod
    async def pusher(connection: AsyncWebsocketConsumer, message: MsgType) -> None:
        if isinstance(message, str):
            await connection.send(text_data=message)
        elif isinstance(message, dict):
            await connection.send(text_data=json.dumps(message))
        elif isinstance(message, bytes):
            await connection.send(bytes_data=message)
        else:
            raise TypeError(f"websocket不能发送{type(message)}的内容！")

    @classmethod
    async def send_personal_message(cls, user_id: int, message: MsgType) -> None:
        """
        发送个人信息
        """
        connection = cls.active_connections.get(user_id)
        if connection:
            await cls.pusher(connection, message)

    @classmethod
    async def send_data(cls, user_id, msg_type, record_msg):
        msg = dict(type=msg_type, record_msg=record_msg)
        await cls.send_personal_message(user_id, msg)

    questions_and_answers_map = {
        "HELLO SERVER": "Hello {user_id}",
        "HEARTBEAT": "{user_id}",
    }

    async def send_heartbeat(self):
        while True:
            await self.send_personal_message(self.user_id, {"type": 3})
            await asyncio.sleep(60)


# ws_manage = ConnectionManager()