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() user_info = db.execute_query("SELECT first_name, email FROM users WHERE id = %s", (user_id,)) first_name = user_info[0]['first_name'] if user_info else None email = user_info[0]['email'] if user_info else None message_payload = { "id": msg_id, "is_from_admin": is_admin, "message": json_text, "created_at": now, "type": "message", "first_name": first_name, "email": email } # 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()