itpdp/api/src/routes/party-socket.ts

import { Elysia } from "elysia";

import { betterAuthElysia } from "../auth";

import { db } from "../db";
import { getMemberRecord, getPartyStatus } from "../party-data";
import type { PartySocketEvent, QuizState } from "../party-types";

export function userTopic(userId: string) {
	return `user:${userId}`;
}

export function partyTopic(partyId: string) {
	return `party:${partyId}`;
}

export const socketPartyId = new WeakMap<object, string>();

export const pubsub = {
	_server: null as ReturnType<typeof Bun.serve> | null,
	setServer(server: ReturnType<typeof Bun.serve> | null) {
		this._server = server;
	},
	publish(topic: string, data: string) {
		this._server?.publish(topic, data);
	},
	publishPartyData(partyId: string, data: PartySocketEvent) {
		pubsub.publish(`party:${partyId}`, JSON.stringify(data));
	},
};

export async function broadcastQuizState(
	ws: { publish: (topic: string, message: string) => void },
	partyId: string,
) {
	const partyRecord = await db.query.party.findFirst({
		where: { id: partyId },
	});

	if (!partyRecord) return;

	const quizData = ((partyRecord.data ?? {}) as Record<string, unknown>).quiz as
		| QuizState
		| undefined;
	if (!quizData) return;
	if (!quizData) return;

	ws.publish(
		partyTopic(partyId),
		JSON.stringify({
			type: "quiz_state",
			quiz: quizData,
		}),
	);
}

export const topic = {
	user: userTopic,
	party: partyTopic,
};

export const partySocketApp = new Elysia()
	.use(betterAuthElysia)
	.group("/party-socket", (app) =>
		app
			.get("/test", () => ({ ok: 1 }))
			.ws("/ws", {
				auth: true,
				publishToSelf: true,
				open: async (ws) => {
					const user = ws.data.user;
					if (!user) return;

					ws.subscribe(userTopic(user.id));

					const membership = await getMemberRecord(db, user.id);
					if (!membership) {
						ws.send(
							JSON.stringify({
								type: "party_status",
								party: null,
								members: [],
							} as PartySocketEvent),
						);
						return;
					}

					socketPartyId.set(ws, membership.partyId);
					ws.subscribe(partyTopic(membership.partyId));

					const snapshot = await getPartyStatus(membership.partyId);
					if (snapshot) {
						ws.send(
							JSON.stringify({
								type: "party_status",
								party: snapshot.party,
								members: snapshot.members,
							} as PartySocketEvent),
						);

						await broadcastQuizState(ws, membership.partyId);
					}
				},
				message: async (ws, message) => {
					const data = ws.data;
					const user = data.user;
					if (!user) return;

					if (typeof message !== "string") return;

					let parsed: { type: string; payload?: unknown };
					try {
						parsed = JSON.parse(message);
					} catch {
						ws.send(JSON.stringify({ type: "error", message: "Invalid JSON" }));
						return;
					}

					if (parsed.type === "ping") {
						ws.send(JSON.stringify({ type: "pong" }));
						return;
					}

					if (parsed.type !== "member_payload") return;

					const MAX_MEMBER_PAYLOAD_SIZE = 8_000;
					const payloadString = JSON.stringify(parsed.payload);
					if (payloadString.length > MAX_MEMBER_PAYLOAD_SIZE) {
						ws.send(
							JSON.stringify({ type: "error", message: "Payload too large." }),
						);
						return;
					}

					const membership = await getMemberRecord(db, user.id);
					if (!membership) return;

					const currentParty = await db.query.party.findFirst({
						where: { id: membership.partyId },
					});
					if (!currentParty) return;

					ws.publish(
						partyTopic(membership.partyId),
						JSON.stringify({
							type: "member_payload",
							fromUserId: user.id,
							payload: parsed.payload,
						} as PartySocketEvent),
					);
				},
				close: async (ws) => {
					const user = ws.data.user;
					if (!user) return;

					ws.unsubscribe(userTopic(user.id));

					const partyId = socketPartyId.get(ws);
					if (!partyId) return;

					ws.unsubscribe(partyTopic(partyId));
				},
			}),
	);