#!/usr/bin/env python3.12

import json
import time

from fastapi import WebSocket
from cah.constructors import CAHClient

class ConnectionManager:
    def __init__(self):
        self.active_connections: dict = {}

    def get_connection_by_ws(self, websocket: WebSocket) -> WebSocket:
        return self.active_connections.get(websocket)
    
    def get_connection_by_csid(self, csid: str) -> WebSocket:
        for connection in self.active_connections:
            if connection.get('csid') == csid:
                return connection
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[websocket] = {
            'client': None,
            'websocket': websocket,
        }
 

    async def handshake_complete(self,
                           websocket: WebSocket,
                           csid: str,
                           handshakedClient: CAHClient):
        self.active_connections[websocket] = {
            'websocket': websocket,
            'csid': csid,
            'client': handshakedClient,
        }
        
        await self.broadcast({
            "event": "client_connected",
            "ts": int(time.time()),
            "data": {
                "connected_resource": handshakedClient.resource, 
            }
        })        

    def disconnect(self, websocket: WebSocket, csid: str = None):
        self.active_connections.pop(websocket)

    async def send(self, message: str, websocket: WebSocket):
        await websocket.send_json(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_json(message)