Toju/server/src/websocket/index.ts

import {
  IncomingMessage,
  Server,
  ServerResponse
} from 'http';
import { WebSocketServer, WebSocket } from 'ws';
import { v4 as uuidv4 } from 'uuid';
import { connectedUsers } from './state';
import {
  broadcastToServer,
  getServerIdsForOderId,
  isOderIdConnectedToServer
} from './broadcast';
import { handleWebSocketMessage } from './handler';

type IncomingWebSocketMessage = Parameters<typeof handleWebSocketMessage>[1];

/** How often to ping all connected clients (ms). */
const PING_INTERVAL_MS = 30_000;
/** Maximum time a client can go without a pong before we consider it dead (ms). */
const PONG_TIMEOUT_MS = 45_000;

function removeDeadConnection(connectionId: string): void {
  const user = connectedUsers.get(connectionId);

  if (user) {
    console.log(`Removing dead connection: ${user.displayName ?? 'Unknown'} (${user.oderId})`);

    const remainingServerIds = getServerIdsForOderId(user.oderId, connectionId);

    user.serverIds.forEach((sid) => {
      if (isOderIdConnectedToServer(user.oderId, sid, connectionId)) {
        return;
      }

      broadcastToServer(sid, {
        type: 'user_left',
        oderId: user.oderId,
        displayName: user.displayName,
        serverId: sid,
        serverIds: remainingServerIds
      }, user.oderId);
    });

    try {
      user.ws.terminate();
    } catch {
      console.warn(`Failed to terminate WebSocket for ${user.displayName ?? 'Unknown'} (${user.oderId})`);
    }
  }

  connectedUsers.delete(connectionId);
}

export function setupWebSocket(server: Server<typeof IncomingMessage, typeof ServerResponse>): void {
  const wss = new WebSocketServer({ server });
  // Periodically ping all clients and reap dead connections
  const pingInterval = setInterval(() => {
    const now = Date.now();

    connectedUsers.forEach((user, connectionId) => {
      if (now - user.lastPong > PONG_TIMEOUT_MS) {
        removeDeadConnection(connectionId);
        return;
      }

      if (user.ws.readyState === WebSocket.OPEN) {
        try {
          user.ws.ping();
        } catch {
          console.warn(`Failed to ping client ${user.displayName ?? 'Unknown'} (${user.oderId})`);
        }
      }
    });
  }, PING_INTERVAL_MS);

  wss.on('close', () => clearInterval(pingInterval));

  wss.on('connection', (ws: WebSocket) => {
    const connectionId = uuidv4();
    const now = Date.now();

    connectedUsers.set(connectionId, { oderId: connectionId, ws, serverIds: new Set(), lastPong: now });

    ws.on('pong', () => {
      const user = connectedUsers.get(connectionId);

      if (user) {
        user.lastPong = Date.now();
      }
    });

    ws.on('message', async (data) => {
      let message: IncomingWebSocketMessage;

      try {
        message = JSON.parse(data.toString()) as IncomingWebSocketMessage;
      } catch (err) {
        console.error('Invalid WebSocket message:', err);

        return;
      }

      try {
        await handleWebSocketMessage(connectionId, message);
      } catch (err) {
        console.error('WebSocket message handler failed:', err);
      }
    });

    ws.on('close', () => {
      removeDeadConnection(connectionId);
    });

    ws.send(JSON.stringify({ type: 'connected', connectionId, serverTime: Date.now() }));
  });
}