gomoku/src/game/WebSocketHandler.ts

import { ElysiaWS } from 'elysia/dist/ws';
import { GameInstance } from './GameInstance';
import {
  renderGameBoardHtml,
  renderPlayerInfoHtml,
} 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.currentPlayer !== 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 !== ws),
    );
    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.`);
    }
    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);
        const updatedPlayerInfoHtml = renderPlayerInfoHtml(
          game.id,
          playerId,
        );
        ws.send(updatedPlayerInfoHtml);

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

  public sendMessageToGame(
    gameId: string,
    message: string,
  ): void {
    const connections = this.connections.get(gameId);
    if (connections) {
      connections.forEach((ws) => {
        ws.send('<div id="messages">' + message + '</div>')
      });
    }
  }

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

  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;
  }
}