import { WebSocketServer } from 'ws';
import { URL } from 'node:url';
import { validateToken } from './auth.js';

const studentsById = new Map();
const teachers = new Set();

function send(ws, obj) {
  if (ws.readyState !== ws.OPEN) return;
  try {
    ws.send(JSON.stringify(obj));
  } catch {
    /* ignore */
  }
}

function broadcastToTeachers(obj) {
  for (const tws of teachers) send(tws, obj);
}

function rosterSnapshot() {
  return Array.from(studentsById.values()).map((s) => ({
    userId: s.userId,
    username: s.username,
    state: s.lastState,
  }));
}

export function attachWsServer(server) {
  const wss = new WebSocketServer({ noServer: true });

  server.on('upgrade', async (req, socket, head) => {
    let user = null;
    try {
      const url = new URL(req.url, `http://${req.headers.host}`);
      if (url.pathname !== '/ws') {
        socket.destroy();
        return;
      }
      const token = url.searchParams.get('token');
      user = await validateToken(token);
    } catch {
      user = null;
    }

    if (!user) {
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }

    wss.handleUpgrade(req, socket, head, (ws) => {
      ws.user = user;
      wss.emit('connection', ws, req);
    });
  });

  wss.on('connection', (ws) => {
    const { id: userId, username, role } = ws.user;

    if (role === 'teacher') {
      teachers.add(ws);
      send(ws, { type: 'roster', students: rosterSnapshot() });
    } else {
      const prev = studentsById.get(userId);
      if (prev && prev.ws !== ws) {
        try { prev.ws.close(4000, 'Replaced by new connection'); } catch { /* ignore */ }
      }
      const entry = { userId, username, ws, lastState: prev?.lastState || null };
      studentsById.set(userId, entry);
      broadcastToTeachers({
        type: 'student-update',
        userId,
        username,
        state: entry.lastState,
      });
    }

    ws.on('message', (raw) => {
      let msg;
      try { msg = JSON.parse(raw.toString()); } catch { return; }
      if (!msg || typeof msg !== 'object') return;

      if (role === 'student') {
        if (msg.type === 'workspace') {
          const entry = studentsById.get(userId);
          if (entry) entry.lastState = msg.state || null;
          broadcastToTeachers({
            type: 'student-update',
            userId,
            username,
            state: msg.state || null,
          });
        }
        return;
      }

      if (role === 'teacher') {
        if (msg.type === 'list') {
          send(ws, { type: 'roster', students: rosterSnapshot() });
          return;
        }
        if (msg.type === 'push-code') {
          const target = studentsById.get(Number(msg.targetUserId));
          if (!target) {
            send(ws, { type: 'push-code-result', ok: false, targetUserId: msg.targetUserId, error: 'Student not online' });
            return;
          }
          // Update the cached student state immediately so teacher dashboards
          // reflect pushed code without waiting for the next student edit.
          target.lastState = msg.state || null;
          broadcastToTeachers({
            type: 'student-update',
            userId: target.userId,
            username: target.username,
            state: target.lastState,
          });
          send(target.ws, { type: 'push-code', state: msg.state || null, from: ws.user.username });
          send(ws, { type: 'push-code-result', ok: true, targetUserId: target.userId });
          return;
        }
      }
    });

    ws.on('close', () => {
      if (role === 'teacher') {
        teachers.delete(ws);
      } else {
        const entry = studentsById.get(userId);
        if (entry && entry.ws === ws) {
          studentsById.delete(userId);
          broadcastToTeachers({ type: 'student-leave', userId });
        }
      }
    });
  });

  return wss;
}