import json
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from channels.exceptions import StopConsumer
import logging
from rest_framework.authtoken.models import Token
from urllib.parse import parse_qs

logger = logging.getLogger(__name__)

class NotificationConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        """建立WebSocket连接"""
        try:
            # 从URL参数中获取token
            query_string = self.scope.get('query_string', b'').decode()
            query_params = parse_qs(query_string)
            token_key = query_params.get('token', [''])[0]
            
            if not token_key:
                logger.warning("WebSocket连接尝试，但没有提供token")
                await self.close()
                return
                
            # 验证token
            self.user = await self.get_user_from_token(token_key)
            if not self.user:
                logger.warning(f"WebSocket连接尝试，但token无效: {token_key}")
                await self.close()
                return

            # 为用户创建专属房间
            self.room_name = f"notification_user_{self.user.id}"
            await self.channel_layer.group_add(
                self.room_name,
                self.channel_name
            )
            await self.accept()
            logger.info(f"用户 {self.user.username} WebSocket连接成功")
            
        except Exception as e:
            logger.error(f"WebSocket连接错误: {str(e)}")
            await self.close()

    @database_sync_to_async
    def get_user_from_token(self, token_key):
        try:
            token = Token.objects.select_related('user').get(key=token_key)
            return token.user
        except Token.DoesNotExist:
            return None

    async def disconnect(self, close_code):
        """断开WebSocket连接"""
        try:
            if hasattr(self, 'room_name'):
                await self.channel_layer.group_discard(
                    self.room_name,
                    self.channel_name
                )
                logger.info(f"用户 {self.user.username} 已断开连接，关闭代码: {close_code}")
        except Exception as e:
            logger.error(f"断开连接时发生错误: {str(e)}")

    async def notification(self, event):
        """处理并发送通知消息"""
        try:
            await self.send(text_data=json.dumps(event))
            logger.info(f"已发送通知给用户 {self.user.username}")
        except Exception as e:
            logger.error(f"发送通知消息时发生错误: {str(e)}")