DiscordClone/convex/messages.ts

import { query, mutation } from "./_generated/server";
import { paginationOptsValidator } from "convex/server";
import { v } from "convex/values";
import { getPublicStorageUrl } from "./storageUrl";
import { getRolesForUser } from "./roles";

async function enrichMessage(ctx: any, msg: any, userId?: any) {
  const sender = await ctx.db.get(msg.senderId);

  let avatarUrl: string | null = null;
  if (sender?.avatarStorageId) {
    avatarUrl = await getPublicStorageUrl(ctx, sender.avatarStorageId);
  }

  const reactionDocs = await ctx.db
    .query("messageReactions")
    .withIndex("by_message", (q: any) => q.eq("messageId", msg._id))
    .collect();

  const reactions: Record<string, { count: number; me: boolean }> = {};
  for (const r of reactionDocs) {
    const entry = (reactions[r.emoji] ??= { count: 0, me: false });
    entry.count++;
    if (userId && r.userId === userId) {
      entry.me = true;
    }
  }

  let replyToUsername: string | null = null;
  let replyToDisplayName: string | null = null;
  let replyToContent: string | null = null;
  let replyToNonce: string | null = null;
  let replyToAvatarUrl: string | null = null;
  if (msg.replyTo) {
    const repliedMsg = await ctx.db.get(msg.replyTo);
    if (repliedMsg) {
      const repliedSender = await ctx.db.get(repliedMsg.senderId);
      replyToUsername = repliedSender?.username || "Unknown";
      replyToDisplayName = repliedSender?.displayName || null;
      replyToContent = repliedMsg.ciphertext;
      replyToNonce = repliedMsg.nonce;
      if (repliedSender?.avatarStorageId) {
        replyToAvatarUrl = await getPublicStorageUrl(ctx, repliedSender.avatarStorageId);
      }
    }
  }

  return {
    id: msg._id,
    channel_id: msg.channelId,
    sender_id: msg.senderId,
    ciphertext: msg.ciphertext,
    nonce: msg.nonce,
    signature: msg.signature,
    key_version: msg.keyVersion,
    created_at: new Date(msg._creationTime).toISOString(),
    username: sender?.username || "Unknown",
    displayName: sender?.displayName || null,
    public_signing_key: sender?.publicSigningKey || "",
    avatarUrl,
    reactions: Object.keys(reactions).length > 0 ? reactions : null,
    replyToId: msg.replyTo || null,
    replyToUsername,
    replyToDisplayName,
    replyToContent,
    replyToNonce,
    replyToAvatarUrl,
    editedAt: msg.editedAt || null,
    pinned: msg.pinned || false,
  };
}

export const list = query({
  args: {
    paginationOpts: paginationOptsValidator,
    channelId: v.id("channels"),
    userId: v.optional(v.id("userProfiles")),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const result = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .order("desc")
      .paginate(args.paginationOpts);

    const enrichedPage = await Promise.all(
      result.page.map((msg) => enrichMessage(ctx, msg, args.userId))
    );

    return { ...result, page: enrichedPage };
  },
});

export const send = mutation({
  args: {
    channelId: v.id("channels"),
    senderId: v.id("userProfiles"),
    ciphertext: v.string(),
    nonce: v.string(),
    signature: v.string(),
    keyVersion: v.number(),
    replyTo: v.optional(v.id("messages")),
  },
  returns: v.object({ id: v.id("messages") }),
  handler: async (ctx, args) => {
    const id = await ctx.db.insert("messages", {
      channelId: args.channelId,
      senderId: args.senderId,
      ciphertext: args.ciphertext,
      nonce: args.nonce,
      signature: args.signature,
      keyVersion: args.keyVersion,
      replyTo: args.replyTo,
    });
    return { id };
  },
});

export const sendBatch = mutation({
  args: {
    messages: v.array(v.object({
      channelId: v.id("channels"),
      senderId: v.id("userProfiles"),
      ciphertext: v.string(),
      nonce: v.string(),
      signature: v.string(),
      keyVersion: v.number(),
    })),
  },
  returns: v.object({ count: v.number() }),
  handler: async (ctx, args) => {
    for (const msg of args.messages) {
      await ctx.db.insert("messages", { ...msg });
    }
    return { count: args.messages.length };
  },
});

export const edit = mutation({
  args: {
    id: v.id("messages"),
    ciphertext: v.string(),
    nonce: v.string(),
    signature: v.string(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    await ctx.db.patch(args.id, {
      ciphertext: args.ciphertext,
      nonce: args.nonce,
      signature: args.signature,
      editedAt: Date.now(),
    });
    return null;
  },
});

export const pin = mutation({
  args: {
    id: v.id("messages"),
    pinned: v.boolean(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    await ctx.db.patch(args.id, { pinned: args.pinned });
    return null;
  },
});

export const listPinned = query({
  args: {
    channelId: v.id("channels"),
    userId: v.optional(v.id("userProfiles")),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const pinned = await ctx.db
      .query("messages")
      .withIndex("by_channel_pinned", (q) =>
        q.eq("channelId", args.channelId).eq("pinned", true)
      )
      .collect();

    return Promise.all(
      pinned.map((msg) => enrichMessage(ctx, msg, args.userId))
    );
  },
});

// Slim paginated query for bulk search index rebuilding.
// Skips reactions, avatars, reply resolution — only resolves sender username.
export const fetchBulkPage = query({
  args: {
    channelId: v.id("channels"),
    paginationOpts: paginationOptsValidator,
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const result = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .order("asc")
      .paginate(args.paginationOpts);

    const enrichedPage = await Promise.all(
      result.page.map(async (msg) => {
        const sender = await ctx.db.get(msg.senderId);
        return {
          id: msg._id,
          channel_id: msg.channelId,
          sender_id: msg.senderId,
          username: sender?.username || "Unknown",
          ciphertext: msg.ciphertext,
          nonce: msg.nonce,
          created_at: new Date(msg._creationTime).toISOString(),
          pinned: msg.pinned || false,
          replyToId: msg.replyTo || null,
        };
      })
    );

    return { ...result, page: enrichedPage };
  },
});

export const listAround = query({
  args: {
    channelId: v.id("channels"),
    messageId: v.id("messages"),
    userId: v.optional(v.id("userProfiles")),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const target = await ctx.db.get(args.messageId);
    if (!target || target.channelId !== args.channelId) {
      return { messages: [], hasOlder: false, hasNewer: false, targetFound: false };
    }

    const targetTime = target._creationTime;

    const before = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) =>
        q.eq("channelId", args.channelId).lt("_creationTime", targetTime)
      )
      .order("desc")
      .take(26);

    const after = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) =>
        q.eq("channelId", args.channelId).gt("_creationTime", targetTime)
      )
      .order("asc")
      .take(26);

    const hasOlder = before.length > 25;
    const hasNewer = after.length > 25;
    const olderMessages = before.slice(0, 25).reverse();
    const newerMessages = after.slice(0, 25);

    const allRaw = [...olderMessages, target, ...newerMessages];
    const messages = await Promise.all(
      allRaw.map((msg) => enrichMessage(ctx, msg, args.userId))
    );

    return { messages, hasOlder, hasNewer, targetFound: true };
  },
});

export const listBefore = query({
  args: {
    channelId: v.id("channels"),
    beforeTimestamp: v.number(),
    userId: v.optional(v.id("userProfiles")),
    limit: v.optional(v.number()),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const limit = args.limit ?? 50;
    const rows = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) =>
        q.eq("channelId", args.channelId).lt("_creationTime", args.beforeTimestamp)
      )
      .order("desc")
      .take(limit + 1);

    const hasMore = rows.length > limit;
    const page = rows.slice(0, limit);

    const messages = await Promise.all(
      page.reverse().map((msg) => enrichMessage(ctx, msg, args.userId))
    );

    return { messages, hasMore };
  },
});

export const listAfter = query({
  args: {
    channelId: v.id("channels"),
    afterTimestamp: v.number(),
    userId: v.optional(v.id("userProfiles")),
    limit: v.optional(v.number()),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const limit = args.limit ?? 50;
    const rows = await ctx.db
      .query("messages")
      .withIndex("by_channel", (q) =>
        q.eq("channelId", args.channelId).gt("_creationTime", args.afterTimestamp)
      )
      .order("asc")
      .take(limit + 1);

    const hasMore = rows.length > limit;
    const page = rows.slice(0, limit);

    const messages = await Promise.all(
      page.map((msg) => enrichMessage(ctx, msg, args.userId))
    );

    return { messages, hasMore };
  },
});

export const remove = mutation({
  args: { id: v.id("messages"), userId: v.id("userProfiles") },
  returns: v.null(),
  handler: async (ctx, args) => {
    const message = await ctx.db.get(args.id);
    if (!message) throw new Error("Message not found");

    const isSender = message.senderId === args.userId;
    if (!isSender) {
      const roles = await getRolesForUser(ctx, args.userId);
      const canManage = roles.some(
        (role) => (role.permissions as Record<string, boolean>)?.manage_messages
      );
      if (!canManage) {
        throw new Error("Not authorized to delete this message");
      }
    }

    const reactions = await ctx.db
      .query("messageReactions")
      .withIndex("by_message", (q) => q.eq("messageId", args.id))
      .collect();
    for (const r of reactions) {
      await ctx.db.delete(r._id);
    }

    await ctx.db.delete(args.id);
    return null;
  },
});