const { WebSocketServer } = require('ws');
const log = require('./logger');

let wss = null;
// Map jobId → Set of connected clients
const subscribers = new Map();

function attach(server) {
  wss = new WebSocketServer({ server, path: '/ws' });

  wss.on('connection', (ws) => {
    ws._jobId = null;

    ws.on('message', (data) => {
      try {
        const msg = JSON.parse(data);
        if (msg.type === 'subscribe' && msg.jobId) {
          // Unsubscribe from previous
          if (ws._jobId) {
            const prev = subscribers.get(ws._jobId);
            if (prev) { prev.delete(ws); if (prev.size === 0) subscribers.delete(ws._jobId); }
          }
          // Subscribe to new job
          ws._jobId = msg.jobId;
          if (!subscribers.has(msg.jobId)) subscribers.set(msg.jobId, new Set());
          subscribers.get(msg.jobId).add(ws);
        }
      } catch {}
    });

    ws.on('close', () => {
      if (ws._jobId) {
        const subs = subscribers.get(ws._jobId);
        if (subs) { subs.delete(ws); if (subs.size === 0) subscribers.delete(ws._jobId); }
      }
    });
  });

  log.info('WebSocket server attached at /ws');
}

function notify(jobId, data) {
  const subs = subscribers.get(jobId);
  if (!subs || subs.size === 0) return;
  const msg = JSON.stringify({ jobId, ...data });
  for (const ws of subs) {
    if (ws.readyState === 1) { // OPEN
      ws.send(msg);
    }
  }
}

function getConnectionCount() {
  return wss ? wss.clients.size : 0;
}

module.exports = { attach, notify, getConnectionCount };