DiscordClone/convex/messages.ts

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

const DEFAULT_ROLE_COLOR = "#99aab5";

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);
  }

  // Highest-position role with a non-default colour — mirrors how
  // Discord tints usernames in chat. The Owner role is deliberately
  // skipped so owners fall through to the next non-default role's
  // colour (per product decision: we don't want every owner's name
  // to glow in the bootstrap pink). Default grey (`#99aab5`) is
  // treated as "no colour" so regular users fall back to
  // `--text-primary`.
  let senderRoleColor: string | null = null;
  try {
    const senderRoleDocs = await ctx.db
      .query("userRoles")
      .withIndex("by_user", (q: any) => q.eq("userId", msg.senderId))
      .collect();
    let best: { position: number; color: string } | null = null;
    for (const ur of senderRoleDocs) {
      const role = await ctx.db.get(ur.roleId);
      if (!role) continue;
      if ((role as any).name === "Owner") continue;
      const color = (role as any).color as string | undefined;
      const position = (role as any).position ?? 0;
      if (!color || color.toLowerCase() === DEFAULT_ROLE_COLOR) continue;
      if (!best || position > best.position) {
        best = { position, color };
      }
    }
    senderRoleColor = best?.color ?? null;
  } catch {
    senderRoleColor = null;
  }

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

  // Accumulate into a Map so we don't use emoji surrogates as object
  // field names — Convex's return-value validator rejects non-ASCII
  // field names, which is what caused "Field name 👍 has invalid
  // character" errors on any channel with a unicode reaction.
  //
  // For each emoji we also collect a capped list of reactor profiles
  // (up to MAX_REACTION_USERS) so the client can render the hover
  // tooltip + the full reactions modal without a second round-trip.
  // Reactor profiles are cached per-message so the same user picked
  // for multiple emojis only costs one db lookup.
  const MAX_REACTION_USERS = 100;
  const profileCache = new Map<
    string,
    { userId: string; username: string; displayName: string | null }
  >();
  const resolveProfile = async (reactorUserId: any) => {
    const key = String(reactorUserId);
    const cached = profileCache.get(key);
    if (cached) return cached;
    const profile = await ctx.db.get(reactorUserId);
    const shaped = {
      userId: key,
      username: profile?.username || "Unknown",
      displayName: profile?.displayName || null,
    };
    profileCache.set(key, shaped);
    return shaped;
  };

  interface ReactionAccumulator {
    count: number;
    me: boolean;
    users: Array<{
      userId: string;
      username: string;
      displayName: string | null;
    }>;
  }
  const reactionMap = new Map<string, ReactionAccumulator>();
  for (const r of reactionDocs) {
    let entry = reactionMap.get(r.emoji);
    if (!entry) {
      entry = { count: 0, me: false, users: [] };
      reactionMap.set(r.emoji, entry);
    }
    entry.count++;
    if (userId && r.userId === userId) {
      entry.me = true;
    }
    if (entry.users.length < MAX_REACTION_USERS) {
      entry.users.push(await resolveProfile(r.userId));
    }
  }
  const reactions: Array<{
    emoji: string;
    count: number;
    me: boolean;
    users: Array<{
      userId: string;
      username: string;
      displayName: string | null;
    }>;
  }> = [];
  reactionMap.forEach((info, emoji) => {
    reactions.push({
      emoji,
      count: info.count,
      me: info.me,
      users: info.users,
    });
  });

  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,
    senderRoleColor,
    reactions: 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 };
  },
});

/**
 * Pull the latest N messages from each of the given channel IDs in a
 * single round-trip. Used by SearchPanel to scan across every channel
 * the user can read without spinning up N independent useQuery hooks
 * (React would error on hook-count drift between renders).
 *
 * `perChannelLimit` is clamped to 200 server-side to keep payload
 * bounded even if the client over-asks.
 */
export const searchScan = query({
  args: {
    channelIds: v.array(v.id("channels")),
    perChannelLimit: v.optional(v.number()),
    userId: v.optional(v.id("userProfiles")),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const limit = Math.min(Math.max(args.perChannelLimit ?? 100, 1), 200);
    const out: Array<{
      channelId: string;
      messages: any[];
    }> = [];
    for (const channelId of args.channelIds) {
      const rows = await ctx.db
        .query("messages")
        .withIndex("by_channel", (q) => q.eq("channelId", channelId))
        .order("desc")
        .take(limit);
      const enriched = await Promise.all(
        rows.map((msg) => enrichMessage(ctx, msg, args.userId)),
      );
      out.push({ channelId, messages: enriched });
    }
    return out;
  },
});

// Plain text caps out at 4000 chars in the composer, which encodes to
// well under 8KB of AES-GCM output. Attachments ride a separate
// metadata-only JSON path, so 64KB leaves ~8× headroom for future
// structured content without letting a rogue client write multi-MB
// rows and bloat the database.
const MAX_CIPHERTEXT_CHARS = 64 * 1024;

// Internal write. The public API is `messageActions.send`, which verifies
// the caller's Ed25519 signature over (channelId, senderId, timestamp)
// before invoking this. Never call this from a public mutation — it
// trusts `senderId` completely and doing so would reintroduce the
// spoofing vulnerability.
export const sendInternal = internalMutation({
  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) => {
    if (args.ciphertext.length > MAX_CIPHERTEXT_CHARS) {
      throw new Error("Message too large");
    }
    if (await isBanned(ctx, args.senderId)) {
      throw new Error("You've been banned from this server.");
    }
    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 };
  },
});

// Internal-only: there's no public caller and no way to verify every
// senderId in a batch without a signature per message, which defeats
// the point of batching. Kept as internalMutation so internal seed
// scripts / migrations can still use it.
export const sendBatchInternal = internalMutation({
  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 };
  },
});

// Internal write — public entry point is `messageActions.edit`.
// Enforces that only the original sender can edit. The action layer
// has already verified the caller controls `userId`.
export const editInternal = internalMutation({
  args: {
    id: v.id("messages"),
    userId: v.id("userProfiles"),
    ciphertext: v.string(),
    nonce: v.string(),
    signature: v.string(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    const msg = await ctx.db.get(args.id);
    if (!msg) throw new Error("Message not found");
    if (msg.senderId !== args.userId) {
      throw new Error("Only the author can edit this message");
    }
    await ctx.db.patch(args.id, {
      ciphertext: args.ciphertext,
      nonce: args.nonce,
      signature: args.signature,
      editedAt: Date.now(),
    });
    return null;
  },
});

// Internal write — public entry point is `messageActions.pin`. Only
// `manage_messages` role-holders can pin/unpin; the action layer has
// already verified the caller controls `userId`.
export const pinInternal = internalMutation({
  args: {
    id: v.id("messages"),
    userId: v.id("userProfiles"),
    pinned: v.boolean(),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    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 pin messages");
    }
    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 };
  },
});

// Internal write — public entry point is `messageActions.remove`. The
// existing isSender-or-manage_messages check stays here; the action
// layer verifies the caller actually controls `userId` before we trust
// that arg.
export const removeInternal = internalMutation({
  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;
  },
});