ChatApp/apps/desktop/src/lib/conversationKeySync.ts

import { loadDevicePrivateKey } from '@chat-app/shared/auth';
import { type OwnDeviceCtx, shareConvKeyToDevice } from '@chat-app/shared/chat';
import { pgHexToBytes } from '@chat-app/shared/supabase';

import { devLocalSecretStore } from './secretStore';
import { supabase } from './supabase';

// Watches the `devices` table for new entries AND, on mount, scans every
// conversation we participate in for missing key bundles. Fills gaps by
// re-wrapping our active conv-key for the missing recipient devices.
//
// This fixes the "cannot decrypt" cliff for devices that registered while
// no other participant device was online to share the key with them.

interface SyncCtx {
  myUserId: string;
  myDeviceId: string;
  priv: Uint8Array;
}

// db-types is stale for `conversation_keys`/`active_key_version`; bypass.
function rawFrom(table: string) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  return (supabase as unknown as { from: (t: string) => any }).from(table);
}

export function startConversationKeySync(
  ownUserId: string,
  ownDeviceId: string,
): () => void {
  let cancelled = false;
  let priv: Uint8Array | null = null;

  void loadDevicePrivateKey(devLocalSecretStore, ownUserId, ownDeviceId).then(async (pk) => {
    if (cancelled) return;
    priv = pk;
    if (!priv) return;
    // Run a full backfill once we have the private key — covers devices that
    // registered while we were offline.
    await syncAllExistingGaps({ myUserId: ownUserId, myDeviceId: ownDeviceId, priv });
  });

  const channel = supabase
    .channel('device-key-sync:' + ownDeviceId)
    .on(
      'postgres_changes',
      { event: 'INSERT', schema: 'public', table: 'devices' },
      (payload: { new: { id?: string; user_id?: string; public_key?: string } }) => {
        if (cancelled) return;
        const row = payload.new;
        if (!row?.id || !row.user_id || !row.public_key) return;
        if (row.user_id === ownUserId && row.id === ownDeviceId) return;
        if (!priv) return; // backfill on mount will catch it later
        void wrapForOneDevice(
          { myUserId: ownUserId, myDeviceId: ownDeviceId, priv },
          row.id,
          row.user_id,
          row.public_key,
        );
      },
    )
    .subscribe();

  return () => {
    cancelled = true;
    void supabase.removeChannel(channel);
  };
}

async function listMyConversationIds(myUserId: string): Promise<string[]> {
  const { data, error } = await supabase
    .from('conversation_members')
    .select('conversation_id')
    .eq('user_id', myUserId)
    .eq('accepted', true);
  if (error) {
    console.warn('keySync: own-member lookup failed', error);
    return [];
  }
  return (data ?? []).map((r) => r.conversation_id as string);
}

async function listConversationDevices(
  conversationId: string,
): Promise<{ id: string; user_id: string; public_key: string }[]> {
  const { data: members, error: mErr } = await supabase
    .from('conversation_members')
    .select('user_id')
    .eq('conversation_id', conversationId)
    .eq('accepted', true);
  if (mErr) {
    console.warn('keySync: members lookup failed', mErr);
    return [];
  }
  const userIds = (members ?? []).map((m) => m.user_id as string);
  if (userIds.length === 0) return [];

  const { data: devices, error: dErr } = await supabase
    .from('devices')
    .select('id, user_id, public_key')
    .in('user_id', userIds);
  if (dErr) {
    console.warn('keySync: devices lookup failed', dErr);
    return [];
  }
  return (devices ?? []) as { id: string; user_id: string; public_key: string }[];
}

async function listExistingKeyRecipients(
  conversationId: string,
  keyVersion: number,
): Promise<Set<string>> {
  const { data, error } = await rawFrom('conversation_keys')
    .select('recipient_device_id')
    .eq('conversation_id', conversationId)
    .eq('key_version', keyVersion);
  if (error) {
    console.warn('keySync: existing keys lookup failed', error);
    return new Set();
  }
  return new Set((data ?? []).map((r: { recipient_device_id: string }) => r.recipient_device_id));
}

async function getActiveKeyVersion(conversationId: string): Promise<number> {
  const { data, error } = await rawFrom('conversations')
    .select('active_key_version')
    .eq('id', conversationId)
    .single();
  if (error) {
    console.warn('keySync: active key version lookup failed', error);
    return 1;
  }
  return (data as { active_key_version: number }).active_key_version;
}

async function syncAllExistingGaps(ctx: SyncCtx): Promise<void> {
  const convs = await listMyConversationIds(ctx.myUserId);
  for (const convId of convs) {
    try {
      await syncOneConversationGaps(ctx, convId);
    } catch (err: unknown) {
      console.warn('keySync: conv gap sync failed', { convId, err });
    }
  }
}

async function syncOneConversationGaps(ctx: SyncCtx, convId: string): Promise<void> {
  const version = await getActiveKeyVersion(convId);
  const devices = await listConversationDevices(convId);
  if (devices.length === 0) return;

  const recipients = await listExistingKeyRecipients(convId, version);
  const ownCtx: OwnDeviceCtx = {
    userId: ctx.myUserId,
    deviceId: ctx.myDeviceId,
    privateKey: ctx.priv,
  };

  for (const dev of devices) {
    if (recipients.has(dev.id)) continue;
    // Skip our own device — we already have the bundle if we're capable of
    // sharing (or don't need it if we ourselves haven't been wrapped yet).
    if (dev.id === ctx.myDeviceId) continue;
    try {
      await shareConvKeyToDevice(supabase, convId, dev.id, pgHexToBytes(dev.public_key), ownCtx);
    } catch (err: unknown) {
      // Most common: this device hasn't been wrapped for us yet either, so
      // tryGetConvKey couldn't unwrap. Another peer with the key will fill
      // the gap when they hit syncAllExistingGaps.
      console.warn('keySync: shareConvKeyToDevice gap-fill failed', {
        convId,
        recipient: dev.id,
        err,
      });
    }
  }
}

async function wrapForOneDevice(
  ctx: SyncCtx,
  newDeviceId: string,
  newDeviceUserId: string,
  newDevicePubHex: string,
): Promise<void> {
  const myConvs = new Set(await listMyConversationIds(ctx.myUserId));
  const { data: peerMember, error: pErr } = await supabase
    .from('conversation_members')
    .select('conversation_id')
    .eq('user_id', newDeviceUserId);
  if (pErr) {
    console.warn('keySync: peer-member lookup failed', pErr);
    return;
  }
  const sharedConvs = (peerMember ?? [])
    .map((r) => r.conversation_id as string)
    .filter((id) => myConvs.has(id));
  if (sharedConvs.length === 0) return;

  const newPub = pgHexToBytes(newDevicePubHex);
  const ownCtx: OwnDeviceCtx = {
    userId: ctx.myUserId,
    deviceId: ctx.myDeviceId,
    privateKey: ctx.priv,
  };
  for (const convId of sharedConvs) {
    try {
      await shareConvKeyToDevice(supabase, convId, newDeviceId, newPub, ownCtx);
    } catch (err: unknown) {
      console.warn('keySync: shareConvKeyToDevice failed', { convId, err });
    }
  }
}