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

logger = logging.getLogger(__name__)

class NotificationConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        """建立WebSocket连接"""
        try:
            # 获取token
            headers = dict(self.scope['headers'])
            auth_header = headers.get(b'authorization', b'').decode()
            
            if not auth_header.startswith('Token '):
                await self.close()
                return
                
            token_key = auth_header.split(' ')[1]
            
            # 验证token
            self.user = await self.get_user_from_token(token_key)
            if not self.user:
                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()
            
        except Exception as 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)}")