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

import { on } from "node:events";
import { TRPCError, tracked } from "@trpc/server";
import { sql } from "kysely";
import { jsonObjectFrom } from "kysely/helpers/postgres";
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,
	zUserStorageId,
} from "~/server/utils/zod_types";
import { getKeydbClient, getSubClient } from "~/utils/keydb";
import { withUserColorStr } from "~/utils/kysely-helper";

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(),
				targetMsgId: zMessageId.nullish(),
			}),
		)
		.output(
			z.object({
				items: z.array(MessageDataSchema),
				nextCursor: zMessageId.nullable(),
			}),
		)
		.query(async ({ input, ctx }) => {
			const take = input.take ?? 30;
			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 }),
						);
					}
				}

				if (input.targetMsgId) {
					// If a targetMsgId is provided, we want to fetch messages around that message
					const msgs = await db
						.selectFrom("messages")
						.innerJoin("users", "users.id", "messages.sender")
						.select((eb) => [
							"messages.message",
							"messages.messageid",
							"messages.chatid",
							eb
								.cast<string>(
									eb.fn("to_char", [
										"messages.time",
										eb.val("YYYY-MM-DD HH24:MI:SS"),
									]),
									"text",
								)
								.as("time"),
							"messages.sender",
							"messages.isread",
							"messages.is_forwarded",
							"messages.reply_to_id",
							"users.nome as sender_nome",
							"users.isAdmin as sender_isAdmin",

							withUserColorStr(),
							jsonObjectFrom(
								eb
									.selectFrom("messages as reply_source")
									.select([
										"reply_source.messageid",
										"reply_source.message",
										"reply_source.sender",
									])
									.innerJoin("users", "users.id", "reply_source.sender")
									.select(["users.nome", withUserColorStr()])
									.whereRef(
										"messages.reply_to_id",
										"=",
										"reply_source.messageid",
									)
									.limit(1),
							).as("reply_source"),
							eb
								.exists(
									eb
										.selectFrom("messages_attachments")
										.select("messages_attachments.userStorageId")
										.whereRef(
											"messages_attachments.messageId",
											"=",
											"messages.messageid",
										),
								)
								.$castTo<boolean>()
								.as("has_attachments"),
						])
						.where("messages.chatid", "=", input.chatId)
						.where("messages.messageid", ">=", input.targetMsgId)
						.orderBy("messages.messageid", "asc")
						.execute();
					const prevCheck = await db
						.selectFrom("messages")
						.where("chatid", "=", input.chatId)
						.where("messageid", "<", input.targetMsgId)
						.select("messageid")
						.orderBy("messageid", "desc")
						.limit(1)
						.executeTakeFirst();

					return {
						items: msgs,
						nextCursor: prevCheck?.messageid ?? null,
					};
				}

				const cursor = input.cursor;

				let qry = db
					.selectFrom("messages")
					.innerJoin("users", "users.id", "messages.sender")
					.select((eb) => [
						"messages.message",
						"messages.messageid",
						"messages.chatid",

						eb
							.cast<string>(
								eb.fn("to_char", [
									"messages.time",
									eb.val("YYYY-MM-DD HH24:MI:SS"),
								]),
								"text",
							)
							.as("time"),
						"messages.sender",
						"messages.isread",
						"messages.is_forwarded",
						"messages.reply_to_id",
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
						withUserColorStr(),
						jsonObjectFrom(
							eb
								.selectFrom("messages as reply_source")
								.select([
									"reply_source.messageid",
									"reply_source.message",
									"reply_source.sender",
								])
								.innerJoin("users", "users.id", "reply_source.sender")
								.select(["users.nome", withUserColorStr()])
								.whereRef("messages.reply_to_id", "=", "reply_source.messageid")
								.limit(1),
						).as("reply_source"),
						eb
							.exists(
								eb
									.selectFrom("messages_attachments")
									.select("messages_attachments.userStorageId")
									.whereRef(
										"messages_attachments.messageId",
										"=",
										"messages.messageid",
									),
							)
							.$castTo<boolean>()
							.as("has_attachments"),
					])

					.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,
				is_forwarded: z.boolean().default(false),
				reply_to_id: zMessageId.nullable().default(null),
				attachments: z
					.object({
						userStorageId: zUserStorageId,
						storageId: z.string(),
					})
					.array()
					.default([]),
			}),
		)
		.mutation(async ({ input }) => {
			const newMessage = await db
				.insertInto("messages")
				.values({
					chatid: input.chatId,
					message: input.message,
					sender: input.sender,
					time: new Date(),
					is_forwarded: input.is_forwarded,
					reply_to_id: input.reply_to_id,
				})
				.returningAll()
				.executeTakeFirstOrThrow(
					() => new Error("Failed to insert message into database"),
				);
			let hasAttachments = false;
			if (input.attachments.length > 0) {
				const newAttachments = await db
					.insertInto("messages_attachments")
					.values(
						input.attachments.map(({ userStorageId }) => ({
							messageId: newMessage.messageid,
							userStorageId,
						})),
					)
					.returning("userStorageId")
					.execute();
				hasAttachments = newAttachments.length > 0;
			}
			const rs = await db
				.selectFrom("messages")
				.innerJoin("users", "users.id", "messages.sender")
				.select([
					"messages.messageid",
					"messages.message",
					"messages.sender",
					"users.nome",
					withUserColorStr(),
				])
				.where("messages.messageid", "=", newMessage.reply_to_id)
				.executeTakeFirst();

			const userInfos = await db
				.selectFrom("users")
				.select([
					"nome as sender_nome",
					"isAdmin as sender_isAdmin",
					sql<string>`concat(username, '', id)`.as("color_str"),
				])
				.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,
				has_attachments: hasAttachments,
				reply_source: rs
					? {
							messageid: rs.messageid,
							message: rs.message,
							nome: rs.nome,
							sender: rs.sender,
							color_str: rs.color_str,
						}
					: null,
				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((eb) => [
						"messages.message",
						"messages.messageid",
						"messages.chatid",
						eb
							.cast<string>(
								eb.fn("to_char", [
									"messages.time",
									eb.val("YYYY-MM-DD HH24:MI:SS"),
								]),
								"text",
							)
							.as("time"),
						"messages.sender",
						"messages.isread",
						"messages.is_forwarded",
						"messages.reply_to_id",
						"users.nome as sender_nome",
						"users.isAdmin as sender_isAdmin",
						withUserColorStr(),
						jsonObjectFrom(
							eb
								.selectFrom("messages as reply_source")
								.select([
									"reply_source.messageid",
									"reply_source.message",
									"reply_source.sender",
								])
								.innerJoin("users", "users.id", "reply_source.sender")
								.select(["users.nome", withUserColorStr()])
								.whereRef("messages.reply_to_id", "=", "reply_source.messageid")
								.limit(1),
						).as("reply_source"),
						eb
							.exists(
								eb
									.selectFrom("messages_attachments")
									.select("messages_attachments.userStorageId")
									.whereRef(
										"messages_attachments.messageId",
										"=",
										"messages.messageid",
									),
							)
							.$castTo<boolean>()
							.as("has_attachments"),
					])
					.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;
		}),
});