import functools
import json
import logging
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.db import connections
from django.utils import timezone
from rest_framework.utils.encoders import JSONEncoder as DRFJSONEncoder
from .models import (
    ChatRoom,
    Message,
    ChatParticipantTeacher,
    ChatParticipantParent,
    ChatParticipantStudent,
    MessageReadReceipt,
    MessageDelivery,
    ChatMention,
)

logger = logging.getLogger(__name__)


def school_scoped(func):
    """
    Like @database_sync_to_async, but first re-points the shared "school"
    DB connection at *this connection's* tenant.

    Channels runs every @database_sync_to_async call through a single
    thread-sensitive executor shared by every concurrently-open websocket
    connection in the process. Two different schools' connections can each
    get scheduled between the other's awaits, so the "school" alias must be
    reasserted at the start of every individual DB call — not just once at
    connect() — or one school's queries can silently run against another
    school's database.
    """

    @database_sync_to_async
    def wrapper(self, *args, **kwargs):
        db = self.school_db
        connections["school"].settings_dict.update(
            {
                "NAME": db["name"],
                "USER": db["user"],
                "PASSWORD": db["password"],
                "HOST": db["host"],
                "PORT": db["port"],
            }
        )
        connections["school"].close()
        return func(self, *args, **kwargs)

    return functools.wraps(func)(wrapper)


class ChatConsumer(AsyncWebsocketConsumer):
    def _dumps(self, payload):
        """
        json.dumps, but UUID/Decimal/date-aware — plain json.dumps chokes on
        the raw UUID objects that DRF's PrimaryKeyRelatedField (e.g.
        MessageSerializer's "room" field) puts straight into serializer
        .data without stringifying. DRF's own JSONEncoder is what the REST
        endpoints rely on for this (via the DRF Response renderer), so use
        the same one here for identical behavior over the socket.
        """
        return json.dumps(payload, cls=DRFJSONEncoder)

    async def connect(self):
        self.user_type = self.scope.get("user_type")
        # Every user_id/sender_id column in the chat models is a CharField
        # (MessageReadReceipt, MessageDelivery, Message.sender_id, ...) —
        # keep this as a string everywhere so equality comparisons against
        # values read back from the DB behave consistently.
        raw_user_id = self.scope.get("user_id")
        self.user_id = str(raw_user_id) if raw_user_id is not None else None
        self.school_db = self.scope.get("school_db")
        self.room_id = self.scope["url_route"]["kwargs"]["room_id"]
        self.room_group_name = f"chat_{self.room_id}"

        if not self.user_type or not self.user_id or not self.school_db:
            # No valid token, or token didn't resolve to a teacher/parent/
            # student — reject before touching any DB.
            await self.close()
            return

        # Check if user is participant
        if await self.is_participant():
            # Join room group
            await self.channel_layer.group_add(self.room_group_name, self.channel_name)
            await self.accept()

            # Update last active
            await self.update_last_active()

            # Send unread count
            await self.send_unread_count()
            
            # Send online status to others
            await self.broadcast_online_status(True)
        else:
            await self.close()

    async def disconnect(self, close_code):
        """Handle WebSocket disconnection"""
        # Broadcast offline status to others
        await self.broadcast_online_status(False)
        
        # Leave room group
        await self.channel_layer.group_discard(self.room_group_name, self.channel_name)

    async def receive(self, text_data):
        """Handle incoming messages"""
        try:
            data = json.loads(text_data)
            message_type = data.get("type", "message")

            if message_type == "message":
                await self.handle_message(data)
            elif message_type == "typing":
                await self.handle_typing(data)
            elif message_type == "stop_typing":
                await self.handle_stop_typing(data)
            elif message_type == "read_receipt":
                await self.handle_read_receipt(data)
            elif message_type == "mark_as_read":
                await self.handle_mark_as_read(data)
            elif message_type == "reaction":
                await self.handle_reaction(data)
            elif message_type == "edit_message":
                await self.handle_edit_message(data)
            elif message_type == "delete_message":
                await self.handle_delete_message(data)
            elif message_type == "get_online_users":
                await self.handle_get_online_users(data)

        except Exception as e:
            logger.error(f"Error receiving message: {e}")
            await self.send_error(str(e))

    # ============================================================
    # MESSAGE HANDLERS
    # ============================================================

    async def handle_message(self, data):
        """Save and broadcast message"""
        content = data.get("content", "")
        message_type = data.get("message_type", "TEXT")
        reply_to_id = data.get("reply_to_id")

        # Save message to database
        message = await self.save_message(content, message_type, reply_to_id)

        # Update room's last message time
        await self.update_room_last_message()

        # Serialize message
        message_data = await self.serialize_message(message)

        # Send to room group
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                "type": "chat_message",
                "message": message_data,
                "sender_type": self.user_type,
                "sender_id": self.user_id,
            },
        )

        # Create delivery records for all participants
        await self.create_delivery_records(message)
        
        # Process mentions
        await self.process_mentions(message, content)

    async def handle_edit_message(self, data):
        """Edit an existing message"""
        message_id = data.get("message_id")
        new_content = data.get("content")
        
        if not message_id or not new_content:
            return
        
        # Check if user can edit this message
        message = await self.get_message(message_id)
        if message and message.sender_id == self.user_id and message.sender_type == self.user_type:
            edited_message = await self.update_message_content(message_id, new_content)
            
            if edited_message:
                # Broadcast edited message to room
                await self.channel_layer.group_send(
                    self.room_group_name,
                    {
                        "type": "message_edited",
                        "message_id": message_id,
                        "new_content": new_content,
                        "edited_at": timezone.now().isoformat()
                    }
                )

    async def handle_delete_message(self, data):
        """Delete a message (soft delete)"""
        message_id = data.get("message_id")
        
        if not message_id:
            return
        
        # Check if user can delete this message
        message = await self.get_message(message_id)
        if message and (message.sender_id == self.user_id or await self.is_admin()):
            await self.soft_delete_message(message_id)
            
            # Broadcast deletion to room
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    "type": "message_deleted",
                    "message_id": message_id,
                    "deleted_by": self.user_type,
                    "deleted_by_id": self.user_id
                }
            )

    # ============================================================
    # TYPING INDICATORS
    # ============================================================

    async def handle_typing(self, data):
        """Handle typing indicator"""
        is_typing = data.get('is_typing', True)
        
        # Broadcast typing status to room
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'typing_indicator',
                'user_type': self.user_type,
                'user_id': self.user_id,
                'user_name': await self.get_user_name(),
                'is_typing': is_typing
            }
        )

    async def handle_stop_typing(self, data):
        """Handle stop typing indicator"""
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'typing_indicator',
                'user_type': self.user_type,
                'user_id': self.user_id,
                'user_name': await self.get_user_name(),
                'is_typing': False
            }
        )

    # ============================================================
    # READ RECEIPTS
    # ============================================================

    async def handle_read_receipt(self, data):
        """Handle read receipt for a specific message"""
        message_id = data.get('message_id')
        
        if message_id:
            await self.mark_message_as_read(message_id)
            
            # Broadcast read receipt to room
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    'type': 'read_receipt_update',
                    'message_id': message_id,
                    'user_type': self.user_type,
                    'user_id': self.user_id,
                    'user_name': await self.get_user_name()
                }
            )

    async def handle_mark_as_read(self, data):
        """Mark all messages in room as read"""
        await self.mark_all_messages_read()
        await self.update_last_read()
        await self.send_unread_count()

    # ============================================================
    # REACTIONS
    # ============================================================

    async def handle_reaction(self, data):
        """Handle message reactions (like, love, etc.)"""
        message_id = data.get('message_id')
        reaction = data.get('reaction')  # 'like', 'love', 'laugh', 'wow', 'sad', 'angry'
        action = data.get('action', 'add')  # 'add' or 'remove'
        
        if message_id and reaction:
            await self.update_message_reaction(message_id, reaction, action)
            
            # Get updated reactions count
            reactions = await self.get_message_reactions(message_id)
            
            # Broadcast reaction update to room
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    'type': 'reaction_update',
                    'message_id': message_id,
                    'reaction': reaction,
                    'user_type': self.user_type,
                    'user_id': self.user_id,
                    'action': action,
                    'reactions': reactions
                }
            )

    # ============================================================
    # ONLINE USERS
    # ============================================================

    async def handle_get_online_users(self, data):
        """Get list of online users in the room"""
        online_users = await self.get_online_users()
        await self.send(text_data=self._dumps({
            'type': 'online_users',
            'users': online_users
        }))

    async def broadcast_online_status(self, is_online):
        """Broadcast user online/offline status to room"""
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'user_status',
                'user_type': self.user_type,
                'user_id': self.user_id,
                'user_name': await self.get_user_name(),
                'is_online': is_online
            }
        )

    # ============================================================
    # GROUP EVENT HANDLERS (Broadcast to all clients)
    # ============================================================

    async def chat_message(self, event):
        """Send message to WebSocket"""
        await self.send(
            text_data=self._dumps({"type": "message", "message": event["message"]})
        )

    async def typing_indicator(self, event):
        """Send typing indicator to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "typing",
                    "user_type": event["user_type"],
                    "user_id": event["user_id"],
                    "user_name": event.get("user_name"),
                    "is_typing": event["is_typing"],
                }
            )
        )

    async def read_receipt_update(self, event):
        """Send read receipt update to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "read_receipt",
                    "message_id": event["message_id"],
                    "user_type": event["user_type"],
                    "user_id": event["user_id"],
                    "user_name": event.get("user_name"),
                }
            )
        )

    async def reaction_update(self, event):
        """Send reaction update to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "reaction",
                    "message_id": event["message_id"],
                    "reaction": event["reaction"],
                    "user_type": event["user_type"],
                    "user_id": event["user_id"],
                    "action": event["action"],
                    "reactions": event.get("reactions", {})
                }
            )
        )

    async def message_edited(self, event):
        """Send message edited update to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "message_edited",
                    "message_id": event["message_id"],
                    "new_content": event["new_content"],
                    "edited_at": event["edited_at"]
                }
            )
        )

    async def message_deleted(self, event):
        """Send message deleted update to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "message_deleted",
                    "message_id": event["message_id"],
                    "deleted_by": event["deleted_by"],
                    "deleted_by_id": event["deleted_by_id"]
                }
            )
        )

    async def user_status(self, event):
        """Send user online/offline status to WebSocket"""
        await self.send(
            text_data=self._dumps(
                {
                    "type": "user_status",
                    "user_type": event["user_type"],
                    "user_id": event["user_id"],
                    "user_name": event["user_name"],
                    "is_online": event["is_online"]
                }
            )
        )

    # ============================================================
    # HELPER METHODS
    # ============================================================

    async def send_error(self, error_message):
        """Send error message to client"""
        await self.send(
            text_data=self._dumps({"type": "error", "message": error_message})
        )

    async def send_unread_count(self):
        """Send unread messages count"""
        count = await self.get_unread_count()
        await self.send(text_data=self._dumps({"type": "unread_count", "count": count}))

    # ============================================================
    # DATABASE OPERATIONS
    # ============================================================

    @school_scoped
    def is_participant(self):
        """Check if user is participant in the room"""
        try:
            if self.user_type == "teacher":
                return ChatParticipantTeacher.objects.filter(
                    room_id=self.room_id, teacher_id=self.user_id, is_active=True
                ).exists()
            elif self.user_type == "parent":
                return ChatParticipantParent.objects.filter(
                    room_id=self.room_id, parent_id=self.user_id, is_active=True
                ).exists()
            elif self.user_type == "student":
                return ChatParticipantStudent.objects.filter(
                    room_id=self.room_id, student_id=self.user_id, is_active=True
                ).exists()
        except Exception:
            return False
        return False

    @school_scoped
    def save_message(self, content, message_type, reply_to_id):
        """Save message to database"""
        message = Message.objects.create(
            room_id=self.room_id,
            sender_type=self.user_type,
            sender_id=self.user_id,
            content=content,
            message_type=message_type,
            reply_to_id=reply_to_id,
        )
        return message

    @school_scoped
    def serialize_message(self, message):
        """Serialize message for JSON response"""
        from .serializers import MessageSerializer
        serializer = MessageSerializer(message)
        return serializer.data

    @school_scoped
    def update_room_last_message(self):
        """Update room's last message timestamp"""
        ChatRoom.objects.filter(id=self.room_id).update(last_message_at=timezone.now())

    @school_scoped
    def update_last_active(self):
        """Update participant's last active timestamp"""
        if self.user_type == "teacher":
            ChatParticipantTeacher.objects.filter(
                room_id=self.room_id, teacher_id=self.user_id
            ).update(last_active_at=timezone.now())
        elif self.user_type == "parent":
            ChatParticipantParent.objects.filter(
                room_id=self.room_id, parent_id=self.user_id
            ).update(last_active_at=timezone.now())
        elif self.user_type == "student":
            ChatParticipantStudent.objects.filter(
                room_id=self.room_id, student_id=self.user_id
            ).update(last_active_at=timezone.now())

    @school_scoped
    def get_unread_count(self):
        """Get count of unread messages for this user in this room"""
        messages = Message.objects.filter(room_id=self.room_id, is_deleted=False)
        read_receipts = MessageReadReceipt.objects.filter(
            user_type=self.user_type, user_id=self.user_id, message__in=messages
        ).values_list("message_id", flat=True)

        unread = (
            messages.exclude(id__in=read_receipts)
            .exclude(sender_type=self.user_type, sender_id=self.user_id)
            .count()
        )
        return unread

    @school_scoped
    def create_delivery_records(self, message):
        """Create delivery records for all participants"""
        if self.user_type == "teacher":
            participants = ChatParticipantTeacher.objects.filter(
                room_id=self.room_id, is_active=True
            ).exclude(teacher_id=self.user_id)

            for participant in participants:
                MessageDelivery.objects.get_or_create(
                    message=message,
                    user_type="teacher",
                    user_id=str(participant.teacher_id),
                    defaults={"status": "SENT"},
                )
        elif self.user_type == "parent":
            participants = ChatParticipantParent.objects.filter(
                room_id=self.room_id, is_active=True
            ).exclude(parent_id=self.user_id)

            for participant in participants:
                MessageDelivery.objects.get_or_create(
                    message=message,
                    user_type="parent",
                    user_id=str(participant.parent_id),
                    defaults={"status": "SENT"},
                )
        elif self.user_type == "student":
            participants = ChatParticipantStudent.objects.filter(
                room_id=self.room_id, is_active=True
            ).exclude(student_id=self.user_id)

            for participant in participants:
                MessageDelivery.objects.get_or_create(
                    message=message,
                    user_type="student",
                    user_id=str(participant.student_id),
                    defaults={"status": "SENT"},
                )

    @school_scoped
    def mark_message_as_read(self, message_id):
        """Mark specific message as read"""
        try:
            message = Message.objects.get(id=message_id)
            receipt, created = MessageReadReceipt.objects.get_or_create(
                message=message,
                user_type=self.user_type,
                user_id=self.user_id
            )
            return True
        except Message.DoesNotExist:
            return False

    @school_scoped
    def mark_all_messages_read(self):
        """Mark all messages in room as read"""
        messages = Message.objects.filter(room_id=self.room_id, is_deleted=False)
        for message in messages:
            MessageReadReceipt.objects.get_or_create(
                message=message,
                user_type=self.user_type,
                user_id=self.user_id
            )

    @school_scoped
    def update_last_read(self):
        """Update participant's last read timestamp"""
        if self.user_type == "teacher":
            ChatParticipantTeacher.objects.filter(
                room_id=self.room_id, teacher_id=self.user_id
            ).update(last_read_at=timezone.now())
        elif self.user_type == "parent":
            ChatParticipantParent.objects.filter(
                room_id=self.room_id, parent_id=self.user_id
            ).update(last_read_at=timezone.now())
        elif self.user_type == "student":
            ChatParticipantStudent.objects.filter(
                room_id=self.room_id, student_id=self.user_id
            ).update(last_read_at=timezone.now())

    @school_scoped
    def update_message_reaction(self, message_id, reaction, action):
        """Update message reactions"""
        try:
            message = Message.objects.get(id=message_id)
            reactions = message.reactions or {}
            
            if action == 'add':
                if reaction not in reactions:
                    reactions[reaction] = []
                if self.user_id not in reactions[reaction]:
                    reactions[reaction].append(self.user_id)
            elif action == 'remove':
                if reaction in reactions and self.user_id in reactions[reaction]:
                    reactions[reaction].remove(self.user_id)
                    if not reactions[reaction]:
                        del reactions[reaction]
            
            message.reactions = reactions
            message.save(update_fields=['reactions'])
            return True
        except Message.DoesNotExist:
            return False

    @school_scoped
    def get_message_reactions(self, message_id):
        """Get all reactions for a message"""
        try:
            message = Message.objects.get(id=message_id)
            return message.reactions or {}
        except Message.DoesNotExist:
            return {}

    @school_scoped
    def update_message_content(self, message_id, new_content):
        """Update message content"""
        try:
            message = Message.objects.get(id=message_id)
            # Save edit history
            edit_history = message.edit_history or []
            edit_history.append({
                'old_content': message.content,
                'edited_at': timezone.now().isoformat()
            })
            
            message.content = new_content
            message.edited_at = timezone.now()
            message.edit_history = edit_history
            message.save(update_fields=['content', 'edited_at', 'edit_history'])
            return message
        except Message.DoesNotExist:
            return None

    @school_scoped
    def soft_delete_message(self, message_id):
        """Soft delete a message"""
        try:
            message = Message.objects.get(id=message_id)
            message.is_deleted = True
            message.deleted_at = timezone.now()
            message.deleted_by_type = self.user_type
            message.deleted_by_id = self.user_id
            message.save(update_fields=['is_deleted', 'deleted_at', 'deleted_by_type', 'deleted_by_id'])
            return True
        except Message.DoesNotExist:
            return False

    @school_scoped
    def get_message(self, message_id):
        """Get message by ID"""
        try:
            return Message.objects.get(id=message_id)
        except Message.DoesNotExist:
            return None

    @school_scoped
    def get_user_name(self):
        """Get user's display name"""
        if self.user_type == "teacher":
            from people.models import Teacher
            try:
                teacher = Teacher.objects.get(id=self.user_id)
                return teacher.full_name
            except Teacher.DoesNotExist:
                return "Teacher"
        elif self.user_type == "parent":
            from people.models import Parent
            try:
                parent = Parent.objects.get(id=self.user_id)
                return parent.full_name
            except Parent.DoesNotExist:
                return "Parent"
        elif self.user_type == "student":
            from people.models import Student
            try:
                student = Student.objects.get(id=self.user_id)
                return student.full_name
            except Student.DoesNotExist:
                return "Student"
        return "User"

    @school_scoped
    def is_admin(self):
        """Check if user is admin of the room"""
        if self.user_type == "teacher":
            participant = ChatParticipantTeacher.objects.filter(
                room_id=self.room_id, teacher_id=self.user_id, is_active=True
            ).first()
            return participant and participant.role == 'ADMIN'
        return False

    @school_scoped
    def get_online_users(self):
        """Get list of online users (users with active WebSocket connections)"""
        # This requires tracking online users - implement with Redis or cache
        # For now, return empty list
        return []

    @school_scoped
    def process_mentions(self, message, content):
        """Process @mentions in message"""
        import re
        mention_pattern = r'@(teacher|parent|student)_(\d+)'
        mentions = re.findall(mention_pattern, content)
        
        for user_type, user_id in mentions:
            ChatMention.objects.create(
                message=message,
                user_type=user_type,
                user_id=user_id,
                is_notified=False
            )