from typing import Dict, List
from fastapi import WebSocket

class JobWSManager:
    def __init__(self):
        self.connections: Dict[str, List[WebSocket]] = {}

    async def connect(self, job_id: str, ws: WebSocket):
        await ws.accept()
        self.connections.setdefault(job_id, []).append(ws)

    def disconnect(self, job_id: str, ws: WebSocket):
        if job_id in self.connections:
            self.connections[job_id].remove(ws)
            if not self.connections[job_id]:
                del self.connections[job_id]

    async def send(self, job_id: str, data: dict):
        for ws in self.connections.get(job_id, []):
            await ws.send_json(data)


manager = JobWSManager()