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
        # 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})
        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]
            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()
            
            message["type"] = "message"
            payload = json.dumps(message)
            for connection in self.active_connections[order_id]:
                try:
                    await connection["ws"].send_text(payload)
                except:
                    pass

manager = ChatConnectionManager()