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

class ChatConnectionManager:
    def __init__(self):
        # order_id -> list of active websockets
        self.active_connections: Dict[int, List[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, order_id: int):
        await websocket.accept()
        if order_id not in self.active_connections:
            self.active_connections[order_id] = []
        self.active_connections[order_id].append(websocket)

    def disconnect(self, websocket: WebSocket, order_id: int):
        if order_id in self.active_connections:
            self.active_connections[order_id].remove(websocket)
            if not self.active_connections[order_id]:
                del self.active_connections[order_id]

    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()
            
            payload = json.dumps(message)
            for connection in self.active_connections[order_id]:
                try:
                    await connection.send_text(payload)
                except:
                    # Connection might be dead
                    pass

manager = ChatConnectionManager()