ChatApp/packages/shared/src/chat/messages.ts

import { bytesToUtf8, utf8ToBytes } from '../crypto/index.js';
import { bytesToPgHex, pgHexToBytes } from '../supabase/bytea.js';
import type { AppSupabaseClient } from '../supabase/client.js';
import {
  decryptWithConvKey,
  encryptWithConvKey,
  getOrCreateConvKey,
  type OwnDeviceCtx,
  tryGetConvKey,
} from './convKeys.js';
import type { ChatMessage, DecryptedMessage } from './types.js';

const MESSAGE_COLS =
  'id, conversation_id, sender_id, sender_device_id, reply_to_id, edited_at, deleted_at, created_at, ciphertext, nonce, key_version';

interface MessageRow {
  id: string;
  conversation_id: string;
  sender_id: string;
  sender_device_id: string | null;
  reply_to_id: string | null;
  edited_at: string | null;
  deleted_at: string | null;
  created_at: string;
  ciphertext: string;
  nonce: string;
  key_version: number;
}

interface MessageWithCipher extends ChatMessage {
  ciphertext: Uint8Array;
  nonce: Uint8Array;
  keyVersion: number;
}

function mapMessage(row: MessageRow): MessageWithCipher {
  return {
    id: row.id,
    conversationId: row.conversation_id,
    senderId: row.sender_id,
    senderDeviceId: row.sender_device_id,
    replyToId: row.reply_to_id,
    editedAt: row.edited_at,
    deletedAt: row.deleted_at,
    createdAt: row.created_at,
    ciphertext: pgHexToBytes(row.ciphertext),
    nonce: pgHexToBytes(row.nonce),
    keyVersion: row.key_version,
  };
}

interface DeviceKey {
  deviceId: string;
  userId: string;
  publicKey: Uint8Array;
}

// All device public keys for accepted members of a conversation.
export async function listConversationDeviceKeys(
  client: AppSupabaseClient,
  conversationId: string,
): Promise<DeviceKey[]> {
  // 1. Members
  const { data: members, error: mErr } = await client
    .from('conversation_members')
    .select('user_id, accepted')
    .eq('conversation_id', conversationId);
  if (mErr) throw mErr;

  const memberIds = (members ?? []).filter((m) => m.accepted).map((m) => m.user_id);
  if (memberIds.length === 0) return [];

  // 2. Devices for those members
  const { data: devices, error: dErr } = await client
    .from('devices')
    .select('id, user_id, public_key')
    .in('user_id', memberIds);
  if (dErr) throw dErr;

  return (devices ?? []).map((d) => ({
    deviceId: d.id,
    userId: d.user_id,
    publicKey: pgHexToBytes(d.public_key),
  }));
}

export interface SendMessageParams {
  client: AppSupabaseClient;
  conversationId: string;
  plaintext: string;
  senderUserId: string;
  senderDeviceId: string;
  senderPrivateKey: Uint8Array;
  replyToId?: string;
  // Optional encrypted attachments — their handles are already materialised
  // via `encryptAndUploadAttachment`. The caller is responsible for creating
  // the corresponding message_attachments rows (see insertAttachmentRow) once
  // the returned message id is known.
  attachmentHandles?: import('./attachments.js').AttachmentHandle[];
}

// Encrypts and inserts a message using the shared per-conversation key
// (Sender-Key / Signal-style). The conv-key is generated lazily on first
// send and shared with every existing recipient device. New devices that
// register later receive their key bundle through `shareConvKeyToDevice`.
export async function sendEncryptedMessage(params: SendMessageParams): Promise<ChatMessage> {
  const ownCtx: OwnDeviceCtx = {
    userId: params.senderUserId,
    deviceId: params.senderDeviceId,
    privateKey: params.senderPrivateKey,
  };
  const handle = await getOrCreateConvKey(params.client, params.conversationId, ownCtx);

  const attachments = params.attachmentHandles ?? [];
  const payloadString =
    attachments.length === 0
      ? params.plaintext
      : JSON.stringify({ v: 1, text: params.plaintext, attachments });
  const plainBytes = utf8ToBytes(payloadString);

  const cipher = encryptWithConvKey(plainBytes, handle.key);

  const insertPayload: Record<string, unknown> = {
    conversation_id: params.conversationId,
    sender_id: params.senderUserId,
    sender_device_id: params.senderDeviceId,
    ciphertext: bytesToPgHex(cipher.ciphertext),
    nonce: bytesToPgHex(cipher.nonce),
    key_version: handle.keyVersion,
  };
  if (params.replyToId) insertPayload.reply_to_id = params.replyToId;

  const { data: messageRow, error: insertErr } = await params.client
    .from('messages')
    .insert(insertPayload as never)
    .select(MESSAGE_COLS)
    .single();
  if (insertErr) throw insertErr;
  return mapMessage(messageRow as unknown as MessageRow);
}

// Fetch the last `limit` messages of a conversation in ascending order.
export async function fetchConversationMessages(
  client: AppSupabaseClient,
  conversationId: string,
  limit = 100,
): Promise<MessageWithCipher[]> {
  const { data, error } = await client
    .from('messages')
    .select(MESSAGE_COLS)
    .eq('conversation_id', conversationId)
    .order('created_at', { ascending: false })
    .limit(limit);
  if (error) throw error;
  const rows = (data ?? []) as unknown as MessageRow[];
  return rows.map(mapMessage).reverse();
}

export type { MessageWithCipher };

// ---------------------------------------------------------------------------
// Edit + delete
// ---------------------------------------------------------------------------

export interface EditMessageParams {
  client: AppSupabaseClient;
  messageId: string;
  conversationId: string;
  newPlaintext: string;
  senderPrivateKey: Uint8Array;
}

// Re-encrypts the message body with the conv-key and updates the row.
// Server-side trigger enforces 24h window + sender-only rule.
export async function editEncryptedMessage(
  params: EditMessageParams & { senderUserId: string; senderDeviceId: string },
): Promise<void> {
  const ownCtx: OwnDeviceCtx = {
    userId: params.senderUserId,
    deviceId: params.senderDeviceId,
    privateKey: params.senderPrivateKey,
  };
  const handle = await getOrCreateConvKey(params.client, params.conversationId, ownCtx);

  const cipher = encryptWithConvKey(utf8ToBytes(params.newPlaintext), handle.key);

  const { error } = await params.client
    .from('messages')
    .update({
      ciphertext: bytesToPgHex(cipher.ciphertext),
      nonce: bytesToPgHex(cipher.nonce),
      key_version: handle.keyVersion,
      edited_at: new Date().toISOString(),
    } as never)
    .eq('id', params.messageId);
  if (error) throw error;
}

export async function softDeleteMessage(
  client: AppSupabaseClient,
  messageId: string,
): Promise<void> {
  const { error } = await client
    .from('messages')
    .update({ deleted_at: new Date().toISOString() } as never)
    .eq('id', messageId);
  if (error) throw error;
}

// ---------------------------------------------------------------------------
// Read receipts
// ---------------------------------------------------------------------------

// Upsert read rows for each message. Idempotent thanks to the composite PK.
export async function markMessagesRead(
  client: AppSupabaseClient,
  messageIds: string[],
): Promise<void> {
  if (messageIds.length === 0) return;
  const { data: session, error: aErr } = await client.auth.getUser();
  if (aErr) throw aErr;
  if (!session.user) return;
  const myId = session.user.id;
  const rows = messageIds.map((id) => ({ message_id: id, user_id: myId }));
  const { error } = await client
    .from('message_reads')
    .upsert(rows as never, { onConflict: 'message_id,user_id', ignoreDuplicates: true });
  if (error) throw error;
}

// Returns the set of message ids that `peerUserId` has read (among the given ids).
// RLS hides rows when either side has read receipts off.
export async function listPeerReadsForMessages(
  client: AppSupabaseClient,
  messageIds: string[],
  peerUserId: string,
): Promise<Set<string>> {
  if (messageIds.length === 0) return new Set();
  const { data, error } = await client
    .from('message_reads')
    .select('message_id')
    .eq('user_id', peerUserId)
    .in('message_id', messageIds);
  if (error) throw error;
  return new Set((data ?? []).map((r) => r.message_id));
}

// ---------------------------------------------------------------------------
// Reactions
// ---------------------------------------------------------------------------

export interface MessageReaction {
  messageId: string;
  userId: string;
  emoji: string;
  createdAt: string;
}

export async function listReactionsForMessages(
  client: AppSupabaseClient,
  messageIds: string[],
): Promise<MessageReaction[]> {
  if (messageIds.length === 0) return [];
  const { data, error } = await client
    .from('message_reactions')
    .select('message_id, user_id, emoji, created_at')
    .in('message_id', messageIds);
  if (error) throw error;
  return (data ?? []).map((r) => ({
    messageId: r.message_id,
    userId: r.user_id,
    emoji: r.emoji,
    createdAt: r.created_at,
  }));
}

export async function addReaction(
  client: AppSupabaseClient,
  messageId: string,
  emoji: string,
): Promise<void> {
  const { data: session, error: aErr } = await client.auth.getUser();
  if (aErr) throw aErr;
  if (!session.user) throw new Error('not authenticated');
  const { error } = await client.from('message_reactions').insert({
    message_id: messageId,
    user_id: session.user.id,
    emoji,
  } as never);
  if (error) throw error;
}

export async function removeReaction(
  client: AppSupabaseClient,
  messageId: string,
  emoji: string,
): Promise<void> {
  const { data: session, error: aErr } = await client.auth.getUser();
  if (aErr) throw aErr;
  if (!session.user) throw new Error('not authenticated');
  const { error } = await client
    .from('message_reactions')
    .delete()
    .eq('message_id', messageId)
    .eq('user_id', session.user.id)
    .eq('emoji', emoji);
  if (error) throw error;
}

// ---------------------------------------------------------------------------
// Decrypt helpers
// ---------------------------------------------------------------------------

export interface DecryptParams {
  client: AppSupabaseClient;
  messages: MessageWithCipher[];
  ownDeviceId: string;
  ownPrivateKey: Uint8Array;
}

// Decrypts messages using their conv-key (looked up + cached per
// keyVersion). Returns null `plaintext` when this device has no key bundle
// for that version yet (e.g. brand-new device waiting for share).
export async function decryptMessages(opts: DecryptParams): Promise<DecryptedMessage[]> {
  const out: DecryptedMessage[] = [];
  // Group versions to avoid redundant lookups.
  const versions = new Map<string, Map<number, Uint8Array | null>>(); // convId -> version -> key | null

  for (const m of opts.messages) {
    let convCache = versions.get(m.conversationId);
    if (!convCache) {
      convCache = new Map();
      versions.set(m.conversationId, convCache);
    }
    let key: Uint8Array | null;
    if (convCache.has(m.keyVersion)) {
      key = convCache.get(m.keyVersion) ?? null;
    } else {
      const handle = await tryGetConvKey(
        opts.client,
        m.conversationId,
        opts.ownDeviceId,
        opts.ownPrivateKey,
        m.keyVersion,
      );
      key = handle?.key ?? null;
      convCache.set(m.keyVersion, key);
    }

    let plaintext: string | null = null;
    if (key) {
      try {
        const decoded = decryptWithConvKey(m.ciphertext, m.nonce, key);
        plaintext = bytesToUtf8(decoded);
      } catch {
        plaintext = null;
      }
    }
    out.push({ ...m, plaintext });
  }
  return out;
}