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

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

// Minimal response stub so express-session can run during the HTTP upgrade handshake.
const resStub = { getHeader() {}, setHeader() {}, writeHead() {}, on() {}, once() {}, removeListener() {}, emit() {}, end() {} };

function attach(server) {
  // noServer + manual upgrade so we can authenticate before accepting the socket.
  wss = new WebSocketServer({ noServer: true, maxPayload: 8 * 1024 });

  server.on('upgrade', (req, socket, head) => {
    let pathname;
    try { pathname = new URL(req.url, 'http://localhost').pathname; } catch { pathname = req.url; }
    if (pathname !== '/ws') return; // not ours — leave for others / default handling

    sessionMiddleware(req, resStub, () => {
      if (!req.session || !req.session.user) {
        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
        socket.destroy();
        return;
      }
      wss.handleUpgrade(req, socket, head, (ws) => {
        ws._user = req.session.user;
        wss.emit('connection', ws, req);
      });
    });
  });

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

    ws.on('message', (data) => {
      try {
        const msg = JSON.parse(data);
        if (msg.type === 'subscribe' && msg.jobId) {
          // Only allow subscribing to one's OWN job (prevents cross-user IDOR)
          const job = queue.getJob(String(msg.jobId));
          if (!job || String(job.user_id) !== String(ws._user && ws._user.id)) return;
          // 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 = String(msg.jobId);
          if (!subscribers.has(ws._jobId)) subscribers.set(ws._jobId, new Set());
          subscribers.get(ws._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 (authenticated)');
}

function notify(jobId, data) {
  const subs = subscribers.get(String(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 };