import { WebSocketServer } from "ws";
import type { WebSocket } from "ws";
import type { Server } from "http";
import type { IncomingMessage } from "http";
import { handleUpgrade } from "./auth.js";
import { handleMessage, type AuthenticatedUser } from "./router.js";
import { removeConnection } from "./connections.js";
import { handleLobbyLeave } from "./handlers/lobbyHandlers.js";

export const setupWebSocket = (server: Server): WebSocketServer => {
  const wss = new WebSocketServer({ noServer: true });

  server.on("upgrade", (request, socket, head) => {
    if (request.url !== "/ws") {
      socket.destroy();
      return;
    }
    void handleUpgrade(request, socket, head, wss).catch((err) => {
      console.error("WebSocket upgrade error:", err);
      socket.destroy();
    });
  });

  wss.on(
    "connection",
    (ws: WebSocket, _request: IncomingMessage, auth: AuthenticatedUser) => {
      ws.on("message", (rawData) => {
        void handleMessage(ws, rawData, auth).catch((err) => {
          console.error(
            `WebSocket message error for user ${auth.user.id}:`,
            err,
          );
        });
      });

      ws.on("close", () => {
        void handleDisconnect(ws, auth).catch((err) => {
          console.error(
            `WebSocket disconnect error for user ${auth.user.id}:`,
            err,
          );
        });
      });

      ws.on("error", (err) => {
        console.error(`WebSocket error for user ${auth.user.id}:`, err);
      });
    },
  );

  return wss;
};

const handleDisconnect = async (
  ws: WebSocket,
  auth: AuthenticatedUser,
): Promise<void> => {
  if (!ws.lobbyId) return; // user connected but never joined a lobby
  removeConnection(ws.lobbyId, auth.user.id);
  await handleLobbyLeave(
    ws,
    { type: "lobby:leave", lobbyId: ws.lobbyId },
    auth.user,
  );
};