BRETT/server/utils/notifications.ts

import type { Peer } from 'crossws'
import { and, countDistinct, desc, eq, ne, sql } from 'drizzle-orm'
import { mentionedUserIds } from '../../shared/utils/mentions'
import { notifications, sessions, threadSubscriptions, uploads, users } from '../db/schema'

const peersByUser = new Map<number, Set<Peer>>()
const allPeers = new Set<Peer>()

export function registerGlobalPeer(peer: Peer) {
  allPeers.add(peer)
}

export function unregisterGlobalPeer(peer: Peer) {
  allPeers.delete(peer)
}

export function broadcastToAll(payload: unknown) {
  const message = JSON.stringify(payload)
  for (const peer of allPeers) {
    try {
      peer.send(message)
    }
    catch {
      allPeers.delete(peer)
    }
  }
}

export function registerPeer(userId: number, peer: Peer) {
  let set = peersByUser.get(userId)
  if (!set) {
    set = new Set()
    peersByUser.set(userId, set)
  }
  set.add(peer)
}

export function unregisterPeer(userId: number | undefined, peer: Peer) {
  if (userId === undefined) return
  const set = peersByUser.get(userId)
  if (!set) return
  set.delete(peer)
  if (!set.size) peersByUser.delete(userId)
}

export function pushToUser(userId: number, payload: unknown) {
  const set = peersByUser.get(userId)
  if (!set) return
  const message = JSON.stringify(payload)
  for (const peer of set) {
    try {
      peer.send(message)
    }
    catch {
      set.delete(peer)
    }
  }
}

export async function getUserIdFromCookieHeader(cookieHeader: string | undefined) {
  const token = /(?:^|;\s*)session=([^;]+)/.exec(cookieHeader ?? '')?.[1]
  if (!token) return null
  const rows = await useDb()
    .select({ userId: sessions.userId, expiresAt: sessions.expiresAt })
    .from(sessions)
    .where(eq(sessions.id, token))
    .limit(1)
  const row = rows[0]
  if (!row || row.expiresAt < new Date()) return null
  return row.userId
}

export async function notifyComment(actorId: number, commentId: number, uploadId: string, content: string) {
  const db = useDb()
  const [upload] = await db
    .select({ userId: uploads.userId, type: uploads.type, fileName: uploads.fileName, externalUrl: uploads.externalUrl, previewImage: uploads.previewImage })
    .from(uploads)
    .where(eq(uploads.id, uploadId))
    .limit(1)
  if (!upload) return

  const mentionUsers = content.includes('@')
    ? await db.select({ id: users.id, name: users.name, color: users.color }).from(users)
    : []
  const mentioned = mentionedUserIds(content, mentionUsers)
  const subscriptions = await db.select({ userId: threadSubscriptions.userId, enabled: threadSubscriptions.enabled })
    .from(threadSubscriptions).where(eq(threadSubscriptions.uploadId, uploadId))
  const recipients = new Set([...subscriptions.filter(s => s.enabled).map(s => s.userId), ...mentioned])
  for (const subscription of subscriptions) {
    if (!subscription.enabled) recipients.delete(subscription.userId)
  }
  recipients.delete(actorId)
  if (!recipients.size) return

  const [actor] = await db
    .select({ name: users.name, color: users.color })
    .from(users)
    .where(eq(users.id, actorId))
    .limit(1)

  const url = upload.externalUrl ?? `/uploads/${upload.fileName}`
  const embedId = upload.type === 'youtube' && upload.externalUrl ? getYoutubeId(upload.externalUrl) : null
  const thumbUrl = upload.type === 'youtube' && embedId
    ? youtubeThumbUrl(embedId)
    : upload.type === 'image' ? url : upload.type === 'link' || upload.type === 'twitch' ? upload.previewImage : null

  for (const userId of recipients) {
    const id = await db.transaction(async (tx) => {
      await tx.select({ id: users.id }).from(users).where(eq(users.id, userId)).for('update')
      const [subscription] = await tx.select({ enabled: threadSubscriptions.enabled }).from(threadSubscriptions)
        .where(and(eq(threadSubscriptions.uploadId, uploadId), eq(threadSubscriptions.userId, userId))).limit(1).for('update')
      if (subscription && !subscription.enabled) return null
      const thread = and(eq(notifications.userId, userId), eq(notifications.uploadId, uploadId))
      const [existing] = await tx.select({ id: notifications.id, commentId: notifications.commentId })
        .from(notifications).where(thread).orderBy(desc(notifications.id)).limit(1).for('update')
      if (existing) {
        if (existing.commentId > commentId) return null
        await tx.update(notifications)
          .set({ actorId, commentId, isRead: false, createdAt: sql`CURRENT_TIMESTAMP` })
          .where(eq(notifications.id, existing.id))
        await tx.update(notifications).set({ isRead: true }).where(and(thread, ne(notifications.id, existing.id)))
        return existing.id
      }
      const [result] = await tx.insert(notifications).values({ userId, actorId, uploadId, commentId })
      return Number(result.insertId)
    })
    if (id === null) continue
    const [unread] = await db.select({ count: countDistinct(notifications.uploadId) }).from(notifications)
      .where(and(eq(notifications.userId, userId), eq(notifications.isRead, false)))
    pushToUser(userId, {
      type: 'notification',
      unreadCount: Number(unread?.count ?? 0),
      notification: {
        id,
        actorName: actor?.name ?? '?',
        actorColor: actor?.color ?? null,
        uploadId,
        commentId,
        content: content.slice(0, 140),
        isMention: mentioned.has(userId),
        isRead: false,
        createdAt: new Date().toISOString(),
        uploadType: upload.type,
        url,
        thumbUrl,
      },
    })
  }
}