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

import { on } from "node:events";
import { TRPCError, tracked } from "@trpc/server";
import { z } from "zod/v4";
import type { MessagesMessageid } from "~/schemas/public/Messages";
import {
	adminProcedure,
	createTRPCRouter,
	protectedProcedure,
	publicProcedure,
} from "~/server/api/trpc";
import { db } from "~/server/db";
import { add } from "~/server/services/chat.service";
import {
	editMessage,
	setReadMessage,
} from "~/server/services/messages.service";
import { zAsyncIterable } from "~/server/utils/zAsyncIterable";
import {
	MessageDataSchema,
	MessageUpdateEventSchema,
	SidebarUpdateSchema,
	TypingDataSchema,
	zChatId,
	zMessageId,
	zUserId,
} from "~/server/utils/zod_types";
import { getKeydbClient, getSubClient } from "~/utils/keydb";

const getChatChannel = (chatId: string) => `chat:${chatId}:messages`;
const getTypingChannel = (chatId: string) => `chat:${chatId}:typing`;
const getEventsChannel = (chatId: string) => `chat:${chatId}:events`;
const getTypingKey = (chatId: string) => `typing:${chatId}`;

const SIDEBAR_CHANNEL = "sidebar_updates";

type SidebarUpdate = z.infer<typeof SidebarUpdateSchema>;
export type Message = z.infer<typeof MessageDataSchema>;
export type TypingData = z.infer<typeof TypingDataSchema>;

export const chatSSERouter = createTRPCRouter({
	createChat: protectedProcedure
		.input(z.object({ userId: zUserId }))
		.mutation(async ({ input }) => {
			const newChat = await add({
				db,
				userId: input.userId,
			});

			return newChat.chatid;
		}),

	getChat: protectedProcedure
		.input(
			z.object({
				chatId: zChatId,
				cursor: zMessageId.nullish(),
				take: z.number().min(1).max(50).nullish(),
			}),
		)
		.query(async ({ input, ctx }) => {
			try {
				const didUpdate = await setReadMessage({
					chatId: input.chatId,
					db,
					reader: ctx.session.id,
				});

				if (didUpdate) {
					const kdb = getKeydbClient();
					if (kdb) {
						await kdb.publish(
							getEventsChannel(input.chatId),
							JSON.stringify({ eventType: "read_receipt", data: null }),
						);
					}
				}

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

				let qry = db
					.selectFrom("messages")
					.innerJoin("users", "users.id", "messages.sender")
					.select([
						"messages.message",
						"messages.messageid",
						"messages.chatid",
						(eb) =>
							eb
								.cast<string>(
									eb.fn("to_char", [
										"messages.time",
										eb.val("YYYY-MM-DD HH24:MI:SS"),
									]),
									"text",
								)
								.as("time"),
						"messages.sender",
						"messages.isread",
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
					])
					.where("users.nome", "is not", null)
					.where("messages.chatid", "=", input.chatId)
					.orderBy("messages.messageid", "desc")
					.limit(take + 1);

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

				const page = await qry.execute();

				const items = page.reverse();
				let nextCursor: typeof cursor | null = null;
				if (items.length > take) {
					const prev = items.shift();
					nextCursor = prev?.messageid ?? null;
				}
				return {
					items,
					nextCursor,
				};
			} catch (e) {
				throw new TRPCError({
					code: "INTERNAL_SERVER_ERROR",
					message: `Error while fetching messages: ${(e as Error).message}`,
				});
			}
		}),
	sendMessage: 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(
					() => new Error("Failed to insert message into database"),
				);

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

			if (!userInfos) {
				throw new TRPCError({
					code: "NOT_FOUND",
					message: "User not found or no name",
				});
			}
			const payload: Message = {
				...newMessage,
				...userInfos,
				time: newMessage.time.toISOString(),
			};

			const kdb = getKeydbClient();
			const msgChannel = getChatChannel(input.chatId);
			const typingChannel = getTypingChannel(input.chatId);
			const typingKey = getTypingKey(input.chatId);
			if (kdb) {
				await kdb.publish(msgChannel, JSON.stringify(payload));
				const member = JSON.stringify({
					userId: input.sender,
					username: userInfos.sender_nome,
				});
				await kdb.zrem(typingKey, member);
				const typingList = await kdb.zrange(typingKey, 0, -1);

				await kdb.publish(
					typingChannel,
					JSON.stringify(typingList.map((m) => JSON.parse(m))),
				);
				const sidebarPayload: SidebarUpdate = {
					chatId: input.chatId,
					lastMessage: input.message,
					senderName: userInfos.sender_nome,
				};
				await kdb.publish(SIDEBAR_CHANNEL, JSON.stringify(sidebarPayload));
			}

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

				const kdb = getKeydbClient();
				if (kdb) {
					await kdb.publish(
						getEventsChannel(input.chatId),
						JSON.stringify({ eventType: "delete", data: input.messageId }),
					);
				}
			} catch (err) {
				console.error("Failed to delete message:", err);
				throw new TRPCError({
					code: "INTERNAL_SERVER_ERROR",
					message: "Failed to delete message",
				});
			}
		}),

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

				const kdb = getKeydbClient();
				if (kdb) {
					await kdb.publish(
						getEventsChannel(input.chatId),
						JSON.stringify({
							eventType: "edit",
							data: { messageId: input.messageId, message: input.message },
						}),
					);
				}
			} catch (err) {
				console.error("Failed to edit message:", err);
				throw new TRPCError({
					code: "INTERNAL_SERVER_ERROR",
					message: "Failed to edit message",
				});
			}
		}),
	onMessageUpdates: publicProcedure
		.input(z.object({ chatId: zChatId.nullable() }))
		.output(
			zAsyncIterable({
				yield: MessageUpdateEventSchema,
				tracked: false,
			}),
		)
		.subscription(async function* ({ input, signal }) {
			if (!input.chatId) return;
			const sub = getSubClient();
			if (!sub) return; // Build-time safety

			const channel = getEventsChannel(input.chatId);

			try {
				await sub.subscribe(channel);
				const iterable = on(sub, "message");

				for await (const [chan, message] of iterable) {
					if (chan === channel) {
						if (signal?.aborted) break; // Extra safety check
						const parsed = MessageUpdateEventSchema.safeParse(
							JSON.parse(message),
						);
						if (parsed.success) {
							yield parsed.data;
						} else {
							console.warn(
								"Received invalid message update event:",
								parsed.error,
							);
						}
					}
				}
			} finally {
				await sub.unsubscribe(channel);
				await sub.quit();
			}
		}),
	onSidebarUpdate: publicProcedure
		.output(
			zAsyncIterable({
				yield: SidebarUpdateSchema,
				tracked: false,
			}),
		)
		.subscription(async function* () {
			const sub = getSubClient();
			if (!sub) return;
			try {
				await sub.subscribe(SIDEBAR_CHANNEL);
				const iterable = on(sub, "message");
				for await (const [_chan, message] of iterable) {
					const parsed = SidebarUpdateSchema.safeParse(JSON.parse(message));
					if (parsed.success) {
						yield parsed.data;
					} else {
						console.warn("Received invalid message event:", parsed.error);
					}
				}
			} finally {
				await sub.unsubscribe(SIDEBAR_CHANNEL);
				await sub.quit();
			}
		}),

	toggleTyping: publicProcedure
		.input(
			z.object({
				chatId: zChatId,
				isTyping: z.boolean(),
				userId: zUserId,
				username: z.string(),
			}),
		)
		.mutation(async ({ input }) => {
			const kdb = getKeydbClient();
			if (!kdb) return;

			const key = getTypingKey(input.chatId);
			const now = Date.now();
			const member = JSON.stringify({
				userId: input.userId,
				username: input.username,
			});
			if (input.isTyping) {
				await kdb.zadd(key, now, member);
			} else {
				await kdb.zrem(key, member);
			}

			// Cleanup: Remove users who haven't pinged in 6 seconds
			await kdb.zremrangebyscore(key, "-inf", now - 6000);

			const typingList = await kdb.zrange(key, 0, -1);

			await kdb.publish(
				getTypingChannel(input.chatId),
				JSON.stringify(typingList.map((m) => JSON.parse(m))),
			);

			return { success: true };
		}),
	onMessage: publicProcedure
		.input(
			z.object({
				chatId: zChatId.nullable(),
				lastEventId: z.string().nullish(),
			}),
		)
		.output(
			zAsyncIterable({
				yield: MessageDataSchema,
				tracked: true,
			}),
		)
		.subscription(async function* ({ input, signal }) {
			if (!input.chatId) return;
			const sub = getSubClient();
			if (!sub) return; // Build-time safety

			const channel = getChatChannel(input.chatId);

			try {
				// ON CONNECT
				await sub.subscribe(channel);
				const iterable = on(sub, "message");

				let lastMessageId: MessagesMessageid | null =
					(() => {
						if (!input.lastEventId) return null;

						const parsed = zMessageId.safeParse(input.lastEventId);
						if (!parsed.success) {
							return null;
						}
						return parsed.data;
					})() ?? null;

				const newMessagesSinceLast = await db
					.selectFrom("messages")
					.innerJoin("users", "users.id", "messages.sender")
					.select([
						"messages.message",
						"messages.messageid",
						"messages.chatid",
						(eb) =>
							eb
								.cast<string>(
									eb.fn("to_char", [
										"messages.time",
										eb.val("YYYY-MM-DD HH24:MI:SS"),
									]),
									"text",
								)
								.as("time"),
						"messages.sender",
						"messages.isread",
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
					])
					.where("messages.chatid", "=", input.chatId)
					.where("messages.messageid", ">", lastMessageId)
					.orderBy("messages.messageid", "asc")
					.execute();

				function* maybeYield(msg: Message) {
					if (msg.chatid !== input.chatId) {
						// ignore posts from other channels - the event emitter can emit from other channels
						return;
					}
					if (lastMessageId && msg.messageid <= lastMessageId) {
						// ignore posts 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
					lastMessageId = msg.messageid;
				}
				for (const msg of newMessagesSinceLast) {
					yield* maybeYield(msg);
				}

				for await (const [chan, message] of iterable) {
					if (chan === channel) {
						if (signal?.aborted) break; // Extra safety check
						const parsed = MessageDataSchema.safeParse(JSON.parse(message));
						if (parsed.success) {
							yield* maybeYield(parsed.data);
						} else {
							console.warn("Received invalid chat event:", parsed.error);
						}
					}
				}
			} finally {
				await sub.unsubscribe(channel);
				await sub.quit();
			}
		}),
	onTyping: publicProcedure
		.input(
			z.object({
				chatId: zChatId.nullable(),
				username: z.string(),
				userId: zUserId,
			}),
		)
		.output(
			zAsyncIterable({
				yield: TypingDataSchema,
				tracked: false,
			}),
		)
		.subscription(async function* ({ input, signal }) {
			if (!input.chatId) return;
			const sub = getSubClient();
			const kdb = getKeydbClient();
			if (!sub || !kdb) return; // Build-time safety

			const channel = getTypingChannel(input.chatId);
			const typingKey = getTypingKey(input.chatId);

			try {
				// ON CONNECT
				await sub.subscribe(channel);
				const iterable = on(sub, "message");

				const current = await kdb.zrange(typingKey, 0, -1);

				const parsed = TypingDataSchema.safeParse(
					current.map((m) => JSON.parse(m)),
				);
				if (parsed.success) {
					yield parsed.data;
				}

				for await (const [chan, message] of iterable) {
					if (chan === channel) {
						if (signal?.aborted) break; // Extra safety check
						const parsed = TypingDataSchema.safeParse(JSON.parse(message));
						if (parsed.success) {
							yield parsed.data;
						} else {
							console.warn("Received invalid chat event:", parsed.error);
						}
					}
				}
			} finally {
				// ON DISCONNECT (Tab closed, refreshed, or network dropped)
				const member = JSON.stringify({
					userId: input.userId,
					username: input.username,
				});
				await kdb.zrem(typingKey, member);
				const remainingT = await kdb.zrange(typingKey, 0, -1);
				await kdb.publish(
					channel,
					JSON.stringify(remainingT.map((m) => JSON.parse(m))),
				);

				await sub.unsubscribe(channel);
				await sub.quit();
			}
		}),
	setRead: protectedProcedure
		.input(z.object({ chatId: zChatId }))
		.mutation(async ({ input, ctx }) => {
			const didUpdate = await setReadMessage({
				chatId: input.chatId,
				db,
				reader: ctx.session.id,
			});

			if (didUpdate) {
				const kdb = getKeydbClient();
				if (kdb) {
					await kdb.publish(
						getEventsChannel(input.chatId),
						JSON.stringify({ eventType: "read_receipt", data: null }),
					);
				}
			}

			return true;
		}),
});