infoalloggi-monorepo/apps/infoalloggi/src/server/api/routers/messages_sse.ts

import { TRPCError, tracked } from "@trpc/server";
import { z } from "zod/v4";
import {
	adminProcedure,
	createTRPCRouter,
	protectedProcedure,
} from "~/server/api/trpc";
import {
	getCurrentlyTyping,
	removeStaleTyping,
	removeTyping,
} from "~/server/controllers/chat.controller";
import { db } from "~/server/db";
import {
	type ChatMessage,
	editMessage,
	setReadMessage,
} from "~/server/services/messages.service";
import { ee, type MessageUpdate } from "~/server/sse";
import { zAsyncIterable } from "~/server/utils/sse_zod";
import { zChatId, zMessageId, zUserId } from "~/server/utils/zod_types";

export const messagesSSERouter = createTRPCRouter({
	add: protectedProcedure
		.input(
			z.object({
				chatId: zChatId,
				message: z.string(),
				sender: zUserId,
			}),
		)
		.mutation(async ({ input }) => {
			const newMessage = await db
				.insertInto("messages")
				.values({
					chatid: input.chatId,
					message: input.message,
					sender: input.sender,
					time: new Date(),
				})
				.returningAll()
				.executeTakeFirstOrThrow();

			const userInfos = await db
				.selectFrom("users")
				.select([
					"users.nome as sender_nome",
					"users.isAdmin as sender_isAdmin",
				])
				.where("id", "=", input.sender)
				.executeTakeFirst();

			if (!userInfos) {
				throw new TRPCError({
					code: "NOT_FOUND",
					message: "User not found or no name",
				});
			}

			await removeTyping({
				chatId: input.chatId,
				userId: input.sender,
				username: userInfos.sender_nome,
			});
			await removeStaleTyping(input.chatId);

			ee.emit("chatUpdates", input.chatId, {
				data: await getCurrentlyTyping(input.chatId),
				eventType: "typing",
			});

			ee.emit("add", input.chatId, {
				...newMessage,
				...userInfos,
			});

			return { ...newMessage, ...userInfos };
		}),
	delete: adminProcedure
		.input(z.object({ chatId: zChatId, messageId: zMessageId }))
		.mutation(async ({ input }) => {
			await db
				.deleteFrom("messages")
				.where("messageid", "=", input.messageId)
				.execute();

			ee.emit("messageUpdates", input.chatId, {
				data: input.messageId,
				eventType: "delete",
			});
		}),

	edit: adminProcedure
		.input(
			z.object({ chatId: zChatId, message: z.string(), messageId: zMessageId }),
		)
		.mutation(async ({ input }) => {
			await editMessage({
				db,
				message: input.message,
				messageId: input.messageId,
			});

			ee.emit("messageUpdates", input.chatId, {
				data: {
					message: input.message,
					messageId: input.messageId,
				},
				eventType: "edit",
			});
		}),
	infinite: protectedProcedure
		.input(
			z.object({
				chatId: zChatId,
				cursor: z.date().nullish(),
				take: z.number().min(1).max(50).nullish(),
			}),
		)
		.query(async ({ input, ctx }) => {
			try {
				await setReadMessage({
					chatId: input.chatId,
					db,
					reader: ctx.session.id,
				});

				const take = input.take ?? 20;
				const cursor = input.cursor;

				let qry = db
					.selectFrom("messages")
					.selectAll()
					.innerJoin("users", "users.id", "messages.sender")
					.select([
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
					])
					.where("users.nome", "is not", null)
					.where("messages.chatid", "=", input.chatId)
					.orderBy("messages.time", "desc")
					.limit(take + 1);

				if (cursor) {
					qry = qry.where("messages.time", "<=", cursor);
				}

				const page = await qry.execute();

				const items = page.reverse();
				let nextCursor: typeof cursor | null = null;
				if (items.length > take) {
					const prev = items.shift();
					// biome-ignore lint/style/noNonNullAssertion: <ok>
					nextCursor = prev!.time;
				}
				return {
					items,
					nextCursor,
				};
			} catch (e) {
				throw new TRPCError({
					code: "INTERNAL_SERVER_ERROR",
					message: `Error while fetching messages: ${(e as Error).message}`,
				});
			}
		}),

	onAdd: protectedProcedure
		.input(
			z.object({
				chatId: zChatId,
				// lastEventId is the last event id that the client has received
				// On the first call, it will be whatever was passed in the initial setup
				// If the client reconnects, it will be the last event id that the client received
				lastEventId: zMessageId.nullish(),
			}),
		)
		.subscription(async function* ({ input, signal, ctx }) {
			try {
				// We start by subscribing to the event emitter so that we don't miss any new events while fetching
				const iterable = ee.toIterable("add", {
					signal,
				});

				// Fetch the last message createdAt based on the last event id
				let lastMessageCreatedAt = await (async () => {
					const lastEventId = input.lastEventId;
					console.log("lastEventId", lastEventId);
					if (!lastEventId) return null;

					const messageById = await db
						.selectFrom("messages")
						.selectAll()
						.where("messageid", "=", lastEventId)
						.executeTakeFirstOrThrow();

					return messageById?.time ?? null;
				})();

				let qry = db
					.selectFrom("messages")
					.selectAll()
					.innerJoin("users", "users.id", "messages.sender")
					.select([
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
					])
					.where("messages.chatid", "=", input.chatId)
					.orderBy("messages.time", "asc");

				if (lastMessageCreatedAt) {
					qry = qry.where("messages.time", ">", lastMessageCreatedAt);
				}

				const newMessagesSinceLastMessage = await qry.execute();

				await setReadMessage({
					chatId: input.chatId,
					db,
					reader: ctx.session.id,
				});

				function* maybeYield(msg: ChatMessage) {
					if (msg.chatid !== input.chatId) {
						// ignore messages from other channels - the event emitter can emit from other channels
						return;
					}
					if (lastMessageCreatedAt && msg.time <= lastMessageCreatedAt) {
						// ignore messages that we've already sent - happens if there is a race condition between the query and the event emitter
						return;
					}

					yield tracked(msg.messageid, msg);

					// update the cursor so that we don't send this post again
					lastMessageCreatedAt = msg.time;
				}

				// yield the messages we fetched from the db
				for (const post of newMessagesSinceLastMessage) {
					yield* maybeYield(post);
				}

				// yield any new messages from the event emitter
				for await (const [, post] of iterable) {
					yield* maybeYield(post);
				}
			} catch (e) {
				if (e instanceof Error && e.name === "AbortError") return;
				console.error("Error in onAdd subscription:", e);
			}
		}),

	onMessageEvent: protectedProcedure
		.input(z.object({ chatId: zChatId }))
		.output(
			zAsyncIterable({
				tracked: true,
				yield: z.custom<MessageUpdate>(),
			}),
		)
		.subscription(async function* ({ input, signal }) {
			try {
				const iterable = ee.toIterable("messageUpdates", {
					signal,
				});

				for await (const [chatId, data] of iterable) {
					if (chatId === input.chatId) {
						yield tracked(chatId, data);
					}
				}
			} catch (e) {
				if (e instanceof Error && e.name === "AbortError") return;
				console.error("Error in onMessageEvent subscription:", e);
			}
		}),
	read: protectedProcedure
		.input(z.object({ chatId: zChatId }))
		.mutation(async ({ input, ctx }) => {
			await setReadMessage({
				chatId: input.chatId,
				db,
				reader: ctx.session.id,
			});

			return true;
		}),
});