import type { FastifyInstance } from "fastify";
import { WebSocket } from "ws";
import { canAccessProject, bearerToken, userFromToken } from "./auth.js";
import { toUser } from "./mapping.js";

interface Client {
  ws: WebSocket;
  projectId: string;
  userId: string;
}

/**
 * Register the `/api/ws` endpoint. Clients connect with `?projectId=X` and an
 * auth token (bearer header or cookie). The server validates access, then
 * relays `autosave` and `presence` events to the other clients in the same
 * project (never back to the sender).
 */
export function registerWs(app: FastifyInstance): void {
  const rooms = new Map<string, Set<Client>>();

  app.get("/api/ws", { websocket: true }, (socket: WebSocket, req) => {
      const db = app.db;
      const query = req.query as { projectId?: string };
      const projectId = query.projectId;

      // Authenticate (bearer header or cookie).
      const cookieToken = (req.headers.cookie ?? "")
        .split("; ")
        .find((c) => c.startsWith("ps_session="))
        ?.split("=")[1] ?? null;
      const token = bearerToken(req) ?? cookieToken;
      const row = token ? userFromToken(db, token) : null;
      if (!row || !projectId || !canAccessProject(db, toUser(row), projectId)) {
        socket.close(4401);
        return;
      }
      const client: Client = { ws: socket, projectId, userId: row.id };

      const room = rooms.get(projectId) ?? new Set<Client>();
      room.add(client);
      rooms.set(projectId, room);

      socket.on("message", (data: WebSocket.RawData) => {
        let msg: any;
        try {
          msg = JSON.parse(String(data));
        } catch {
          return;
        }
        if (msg?.type !== "autosave" && msg?.type !== "presence") return;
        const payload = JSON.stringify({ ...msg, userId: client.userId, at: new Date().toISOString() });
        for (const other of room) {
          if (other !== client && other.ws.readyState === WebSocket.OPEN) {
            other.ws.send(payload);
          }
        }
      });

      const cleanup = (): void => {
        room.delete(client);
        if (room.size === 0) rooms.delete(projectId);
      };
      socket.on("close", cleanup);
      socket.on("error", cleanup);
    });
}