radionica3d/backend/services/chat_manager.py

import json
import datetime
import db
import auth_utils
from typing import Dict, List, Any
from fastapi import WebSocket

class ChatConnectionManager:
    def __init__(self):
        # order_id -> list of active websockets
        # order_id -> list of dicts: {"ws": websocket, "role": str}
        self.active_connections: Dict[int, List[Dict[str, Any]]] = {}

    async def connect(self, websocket: WebSocket, order_id: int, role: str):
        await websocket.accept()
        if order_id not in self.active_connections:
            self.active_connections[order_id] = []
        self.active_connections[order_id].append({"ws": websocket, "role": role})
        print(f"DEBUG: Chat WS connected for Order #{order_id} (Role: {role}). Total connections: {len(self.active_connections[order_id])}")
        await self.broadcast_presence(order_id)

    def disconnect(self, websocket: WebSocket, order_id: int):
        if order_id in self.active_connections:
            self.active_connections[order_id] = [c for c in self.active_connections[order_id] if c["ws"] != websocket]
            print(f"DEBUG: Chat WS disconnected for Order #{order_id}. Remaining: {len(self.active_connections[order_id])}")
            if not self.active_connections[order_id]:
                del self.active_connections[order_id]
            else:
                import asyncio
                try:
                    asyncio.create_task(self.broadcast_presence(order_id))
                except Exception:
                    pass

    async def broadcast_presence(self, order_id: int):
        if order_id in self.active_connections:
            roles = [c["role"] for c in self.active_connections[order_id]]
            payload = json.dumps({"type": "presence", "online_roles": list(set(roles))})
            for connection in self.active_connections[order_id]:
                try:
                    await connection["ws"].send_text(payload)
                except:
                    pass

    async def broadcast_to_order(self, order_id: int, message: Any):
        if order_id in self.active_connections:
            # We must serialize datetime to JSON string
            if isinstance(message, dict) and 'created_at' in message:
                if isinstance(message['created_at'], datetime.datetime):
                    message['created_at'] = message['created_at'].isoformat()
            
            if "type" not in message:
                message["type"] = "message"
            payload = json.dumps(message)
            print(f"DEBUG: Broadcasting to Order #{order_id} to {len(self.active_connections[order_id])} receivers")
            for connection in self.active_connections[order_id]:
                try:
                    await connection["ws"].send_text(payload)
                except Exception as e:
                    print(f"DEBUG: Failed to send WS message to a connection: {e}")
        else:
            print(f"DEBUG: Attempted broadcast to Order #{order_id} but NO active connections found")

    async def handle_message(self, order_id: int, user_id: int, role: str, text: str):
        is_admin = (role == 'admin')
        text = text.strip()
        if not text:
            return

        # Save to DB
        json_text = json.dumps({"text": text})
        query = "INSERT INTO order_messages (order_id, user_id, is_from_admin, message) VALUES (%s, %s, %s, %s)"
        msg_id = db.execute_commit(query, (order_id, user_id, is_admin, json_text))
        
        now = datetime.datetime.utcnow().isoformat()
        message_payload = {
            "id": msg_id,
            "is_from_admin": is_admin,
            "message": json_text,
            "created_at": now,
            "type": "message"
        }
        
        # Broadcast to all in this chat
        await self.broadcast_to_order(order_id, message_payload)
        
        # Global notifications (toasts, unread counts)
        from services.global_manager import global_manager
        if is_admin:
            # Notify the client who owns the order
            order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
            if order:
                await global_manager.notify_user(order[0]['user_id'])
        else:
            # Notify admins
            await global_manager.notify_admins()
            await global_manager.notify_admins_new_message(order_id, text)
            # External hooks (email, etc.)
            try:
                from services import event_hooks
                event_hooks.on_message_received(order_id, user_id, text)
            except Exception as e:
                print(f"ERROR in chat event hooks: {e}")
        
        return msg_id

manager = ChatConnectionManager()