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

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

export const messagesSSERouter = createTRPCRouter({
  add: protectedProcedure
    .input(
      z.object({
        message: z.string(),
        chatId: zChatId,
        sender: zUserId,
      }),
    )
    .mutation(async ({ input }) => {
      const newMessage = await db
        .insertInto("messages")
        .values({
          message: input.message,
          sender: input.sender,
          chatid: input.chatId,
          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, {
        eventType: "typing",
        data: await getCurrentlyTyping(input.chatId),
      });

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

      return { ...newMessage, ...userInfos };
    }),

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

      ee.emit("messageUpdates", input.chatId, {
        eventType: "edit",
        data: {
          messageId: input.messageId,
          message: input.message,
        },
      });
    }),
  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, {
        eventType: "delete",
        data: input.messageId,
      });
    }),
  read: protectedProcedure
    .input(z.object({ chatId: zChatId }))
    .mutation(async ({ input, ctx }) => {
      await setReadMessage({
        db,
        chatId: input.chatId,
        reader: ctx.session.id,
      });

      return true;
    }),
  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({
          db,
          chatId: input.chatId,
          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();
          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({
          db,
          chatId: input.chatId,
          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({
        yield: z.custom<MessageUpdate>(),
        tracked: true,
      }),
    )
    .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);
      }
    }),
});