gomoku/src/web-socket-handler.ts

import { ElysiaWS } from 'elysia/dist/ws';
import { GameInstance } from './game/game-instance';
import { renderGameBoardHtml } from './view/board-renderer';

interface MakeMoveMessage {
  gameId: string;
  playerId: string;
  row: number;
  col: number;
}

type WS = ElysiaWS<{ query: { playerId: string; gameId: string } }>;
export class WebSocketHandler {
  private connections: Map<string, Array<WS>>;
  private games: Map<string, GameInstance>;

  constructor() {
    this.connections = new Map();
    this.games = new Map();
  }

  public handleConnection(ws: WS): void {
    const { gameId, playerId } = ws.data.query;

    if (!this.connections.has(gameId)) {
      this.connections.set(gameId, []);
    }
    this.connections.get(gameId)?.push(ws);

    const game = this.getGame(gameId);
    if (game) {
      this.broadcastGameState(game.id);
    } else {
      ws.send('Error: game not found');
      ws.close();
    }

    console.log(
      `WebSocket connected, registered for Game ${gameId} as Player ${playerId}`,
    );
  }

  public handleMessage(ws: WS, message: any): void {
    const type: string = message.type;
    // Someday we might have other message types
    if (type === 'make_move') {
      this.handleMakeMove(ws, message as MakeMoveMessage);
    }
  }

  private handleMakeMove(ws: WS, message: MakeMoveMessage): void {
    const { row, col } = message;
    const { gameId, playerId } = ws.data.query;
    console.log(
      `Handling make_move message in game ${gameId} from player ${playerId}: ${{ message }}`,
    );

    if (!gameId || !playerId || row === undefined || col === undefined) {
      this.sendMessage(ws, 'Error: missing gameId, playerId, row, or col');
      return;
    }

    const game = this.games.get(gameId);
    if (!game) {
      this.sendMessage(ws, 'Error: game not found');
      return;
    }

    const playerColor = Object.entries(game.players).find(
      ([_, id]) => id === playerId,
    )?.[0] as ('black' | 'white') | undefined;
    if (!playerColor) {
      this.sendMessage(ws, 'Error: you are not a player in this game');
      return;
    }

    if (game.currentPlayerColor !== playerColor) {
      this.sendMessage(ws, "Error: It's not your turn");
      return;
    }

    try {
      const result = game.makeMove(playerId, row, col);
      if (result.success) {
        this.broadcastGameState(game.id);
        console.log(
          `Move made in game ${game.id} by ${playerId}: (${row}, ${col})`,
        );
      } else {
        this.sendMessage(ws, result.error || 'Error: invalid move');
      }
    } catch (e: any) {
      this.sendMessage(ws, 'Error: ' + e.message);
    }
  }

  public handleDisconnect(ws: WS): void {
    const { gameId, playerId } = ws.data.query;

    const connectionsInGame = this.connections.get(gameId);
    if (!connectionsInGame) {
      console.error(
        `Disconnecting WebSocket for player ${playerId} from game ${gameId}, but that game has no connections!`,
      );
      return;
    }
    this.connections.set(
      gameId,
      connectionsInGame.filter(
        (conn) => conn.data.query.playerId !== ws.data.query.playerId,
      ),
    );
    if (this.connections.get(gameId)?.length === 0) {
      this.connections.delete(gameId);
    }

    if (this.connections.has(gameId)) {
      // Notify remaining players about disconnect
      this.sendMessageToGame(gameId, `${playerId} disconnected.`);
    }

    this.sendTitleBoxesForGame(gameId);

    console.log(`${playerId} disconnected from game ${gameId}`);
  }

  public broadcastGameState(gameId: string): void {
    const game = this.games.get(gameId);
    if (!game) {
      console.warn(
        'Attempted to broadcast state of game ${gameId}, which is not loaded.',
      );
      return;
    }

    const connectionsToUpdate = this.connections.get(gameId);
    if (connectionsToUpdate) {
      connectionsToUpdate.forEach((ws) => {
        const { gameId, playerId } = ws.data.query;

        const updatedBoardHtml = renderGameBoardHtml(game, playerId);
        ws.send(updatedBoardHtml);

        if (game.status === 'finished') {
          if (game.winnerColor === 'draw') {
            this.sendMessageToGame(gameId, 'Game ended in draw.');
          } else if (game.winnerColor) {
            this.sendMessageToGame(
              gameId,
              `${game.winnerColor.toUpperCase()} wins!`,
            );
          }
        } else if (game.status === 'playing') {
          const clientPlayerColor = Object.entries(game.players).find(
            ([_, id]) => id === playerId,
          )?.[0] as ('black' | 'white') | undefined;
          if (
            game.currentPlayerColor &&
            clientPlayerColor === game.currentPlayerColor
          ) {
            this.sendMessage(ws, "It's your turn!");
          } else if (game.currentPlayerColor) {
            this.sendMessage(
              ws,
              `Waiting for ${game.currentPlayerColor}'s move.`,
            );
          }
        } else if (game.status === 'waiting') {
          this.sendMessage(ws, 'Waiting for another player...');
        }
      });
    } else {
      console.log(`No connections to update for game ${gameId}.`);
    }

    this.sendTitleBoxesForGame(gameId);
  }

  public sendMessageToGame(gameId: string, message: string): void {
    const connections = this.connections.get(gameId);
    if (connections) {
      connections.forEach((ws) => {
        this.sendMessage(ws, message);
      });
    }
  }

  public sendMessage(targetWs: WS, message: string): void {
    targetWs.send('<div id="messages">' + message + '</div>');
  }

  private sendTitleBox(targetWs: WS, message: string): void {
    targetWs.send('<div id="title-box">' + message + '</div>');
  }

  private sendTitleBoxesForGame(gameId: string): void {
    const game = this.games.get(gameId);
    if (!game) {
      console.error(
        `Tried to send title boxes for game ${gameId}, but it doesn't exist!`,
      );
      return;
    }
    const connections = this.connections.get(gameId);
    if (!connections) {
      console.log(
        `Attempted to send title boxes for game ${gameId}, but no players are connected.`,
      );
      return;
    }

    var message = '';
    switch (game.status) {
      case 'waiting': {
        message = 'Waiting for players...';
        break;
      }
      case 'playing': {
        const blackTag = game.players.black
          ? this.playerTag(gameId, game.players.black)
          : 'Unknown';
        const whiteTag = game.players.white
          ? this.playerTag(gameId, game.players.white)
          : 'Unknown';
        message = `${blackTag} vs ${whiteTag}`;
        break;
      }
      case 'finished': {
        switch (game.winnerColor) {
          case 'draw': {
            message = 'Game ended in draw.';
            break;
          }
          case 'black': {
            message = `${game.players.black} wins!`;
            break;
          }
          case 'white': {
            message = `${game.players.white} wins!`;
            break;
          }
        }
        break;
      }
    }

    connections.forEach((connection) => {
      this.sendTitleBox(connection, message);
    });
  }

  private playerTag(gameId: string, playerId: string) {
    // Determine whether the player is disconnected
    var connectionIcon = `<img src="/icons/disconnected.svg" alt="Disconnected" class="icon" />`;
    const connections = this.connections.get(gameId);
    if (connections) {
      connections.forEach((ws) => {
        if (ws.data.query.playerId === playerId) {
          console.log(`Connection exists for player ${playerId}`);
          connectionIcon = '';
        }
      });
    }

    // Set the correct name color for the player
    var colorClass = '';
    var turnClass = '';
    const game = this.games.get(gameId);
    if (game) {
      if (game.players.white === playerId) {
        colorClass = 'player-white';
      } else if (game.players.black === playerId) {
        colorClass = 'player-black';
      }
      if (game.getCurrentPlayerId() === playerId) {
        turnClass = 'player-to-play';
      }
    }
    const classes = `player-name ${colorClass} ${turnClass}`.trim();

    return `<span class="${classes}">${playerId}${connectionIcon}</span>`;
  }

  public getGame(gameId: string): GameInstance | undefined {
    return this.games.get(gameId);
  }

  createGame(gameId?: string): GameInstance {
    const game = new GameInstance(gameId);
    this.games.set(game.id, game);
    return game;
  }
}