src/ws_manager.py

from fastapi import WebSocket
from sqlalchemy.orm import Session
from typing import Dict, Optional
import json
from datetime import datetime

from models import Message, Conversation, ConversationMember, GroupMessage, GroupMember, Group, User
from chat import msg_to_dict, group_msg_to_dict

class ConnectionManager:
    def __init__(self):
        self.active: Dict[int, WebSocket] = {}

    async def connect(self, websocket: WebSocket, user_id: int):
        await websocket.accept()
        self.active[user_id] = websocket

    def disconnect(self, user_id: int):
        self.active.pop(user_id, None)

    async def send_to(self, user_id: int, data: dict):
        ws = self.active.get(user_id)
        if ws:
            try:
                await ws.send_text(json.dumps(data))
            except Exception:
                self.active.pop(user_id, None)

    async def handle_dm(self, msg: dict, sender: User, db: Session):
        conversation_id = msg.get("conversation_id")
        target_user_id = msg.get("target_user_id")
        ciphertext_for_recipient = msg.get("ciphertext_for_recipient")
        ciphertext_for_sender = msg.get("ciphertext_for_sender")

        if not ciphertext_for_recipient:
            return

        # Auto-create conversation if target_user_id is provided instead of conversation_id
        if not conversation_id and target_user_id:
            target = db.query(User).filter(User.id == target_user_id).first()
            if not target:
                return
            my_ids = set(m.conversation_id for m in db.query(ConversationMember).filter(ConversationMember.user_id == sender.id).all())
            their_ids = set(m.conversation_id for m in db.query(ConversationMember).filter(ConversationMember.user_id == target_user_id).all())
            common = my_ids & their_ids
            if common:
                conversation_id = common.pop()
            else:
                conv = Conversation()
                db.add(conv)
                db.flush()
                db.add(ConversationMember(conversation_id=conv.id, user_id=sender.id))
                db.add(ConversationMember(conversation_id=conv.id, user_id=target_user_id))
                db.flush()
                conversation_id = conv.id

        if not conversation_id:
            return

        # Verify sender is in conversation
        membership = db.query(ConversationMember).filter(
            ConversationMember.conversation_id == conversation_id,
            ConversationMember.user_id == sender.id
        ).first()
        if not membership:
            return

        # Save message
        message = Message(
            conversation_id=conversation_id,
            sender_id=sender.id,
            ciphertext_for_recipient=ciphertext_for_recipient,
            ciphertext_for_sender=ciphertext_for_sender or ciphertext_for_recipient,
        )
        db.add(message)
        db.commit()
        db.refresh(message)

        # Get all members of this conversation
        members = db.query(ConversationMember).filter(
            ConversationMember.conversation_id == conversation_id
        ).all()

        for member in members:
            payload = {
                "type": "dm",
                "conversation_id": conversation_id,
                "message": msg_to_dict(message, member.user_id)
            }
            await self.send_to(member.user_id, payload)

    async def handle_group(self, msg: dict, sender: User, db: Session):
        group_id = msg.get("group_id")
        ciphertext = msg.get("ciphertext")
        iv = msg.get("iv")
        encrypted_aes_keys = msg.get("encrypted_aes_keys", {})

        if not group_id or not ciphertext or not iv:
            return

        # Verify membership
        membership = db.query(GroupMember).filter(
            GroupMember.group_id == group_id,
            GroupMember.user_id == sender.id
        ).first()
        if not membership:
            return

        import json as _json
        message = GroupMessage(
            group_id=group_id,
            sender_id=sender.id,
            ciphertext=ciphertext,
            iv=iv,
            encrypted_aes_keys=_json.dumps(encrypted_aes_keys),
        )
        db.add(message)
        db.commit()
        db.refresh(message)

        members = db.query(GroupMember).filter(GroupMember.group_id == group_id).all()
        for member in members:
            payload = {
                "type": "group",
                "group_id": group_id,
                "message": group_msg_to_dict(message, member.user_id)
            }
            await self.send_to(member.user_id, payload)

    async def handle_typing(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        group_id = msg.get("group_id")
        is_typing = msg.get("is_typing", False)

        if target_id:
            await self.send_to(target_id, {
                "type": "typing",
                "user_id": sender.id,
                "display_name": sender.display_name,
                "is_typing": is_typing,
                "conversation_id": msg.get("conversation_id")
            })
        elif group_id:
            # Would need to broadcast to all group members
            pass

    async def handle_call_offer(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        await self.send_to(target_id, {
            "type": "call_offer",
            "from": sender.id,
            "from_display_name": sender.display_name,
            "from_avatar_color": sender.avatar_color,
            "call_type": msg.get("call_type", "audio"),
            "sdp": msg.get("sdp")
        })

    async def handle_call_answer(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        await self.send_to(target_id, {
            "type": "call_answer",
            "from": sender.id,
            "sdp": msg.get("sdp")
        })

    async def handle_ice_candidate(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        await self.send_to(target_id, {
            "type": "ice_candidate",
            "from": sender.id,
            "candidate": msg.get("candidate")
        })

    async def handle_call_end(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        await self.send_to(target_id, {
            "type": "call_end",
            "from": sender.id
        })

    async def handle_video_toggle(self, msg: dict, sender: User):
        target_id = msg.get("target_id")
        await self.send_to(target_id, {
            "type": "video_toggle",
            "from": sender.id,
            "enabled": msg.get("enabled", True),
            "display_name": sender.display_name,
            "avatar_color": sender.avatar_color
        })

    async def handle_read(self, msg: dict, sender: User, db: Session):
        conversation_id = msg.get("conversation_id")
        if conversation_id:
            messages = db.query(Message).filter(
                Message.conversation_id == conversation_id,
                Message.sender_id != sender.id,
                Message.is_read == False
            ).all()
            if messages:
                for m in messages:
                    m.is_read = True
                db.commit()
                sender_ids = set(m.sender_id for m in messages)
                for sid in sender_ids:
                    await self.send_to(sid, {
                        "type": "read_receipt",
                        "conversation_id": conversation_id,
                        "read_by": sender.id
                    })

manager = ConnectionManager()