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

import { decryptFrom, encryptFor } from '../crypto/box';
import { bytesToPgHex, pgHexToBytes } from '../supabase/bytea';
import type { AppSupabaseClient } from '../supabase/client';

export interface MigrateParams {
  client: AppSupabaseClient;
  ownUserId: string;
  ownNewPublicKey: Uint8Array;
  ownNewPrivateKey: Uint8Array;
  ownLegacyDeviceIds: string[];
  ownLegacyDevicePrivateKeys: Record<string, Uint8Array>;
}

export interface MigrateResult {
  migratedConversations: number;
  errors: { conversationId: string; reason: string }[];
  // Per-pair detail collected even on success so the SecurityCenter "retry"
  // button can show "X / Y bundles re-wrapped, Z stuck because the device
  // private key is no longer in the local vault."
  attempted: number;
  noStrongholdKey: number;
  decryptFailed: number;
  rpcFailed: number;
}

interface LegacyRow {
  conversation_id: string;
  key_version: number;
  recipient_device_id: string;
  sender_device_id: string;
  sender_user_id: string | null;
  encrypted_key: string;
  nonce: string;
}

function hexNoPrefix(bytes: Uint8Array): string { return bytesToPgHex(bytes).slice(2); }

async function fetchSenderPubKeys(
  client: AppSupabaseClient,
  ids: string[],
): Promise<Map<string, { userId: string; publicKey: Uint8Array }>> {
  if (ids.length === 0) return new Map();
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const { data, error } = await (client as any)
    .from('devices').select('id, user_id, public_key').in('id', ids);
  if (error) throw error;
  const out = new Map<string, { userId: string; publicKey: Uint8Array }>();
  for (const row of (data ?? []) as { id: string; user_id: string; public_key: string }[]) {
    out.set(row.id, { userId: row.user_id, publicKey: pgHexToBytes(row.public_key) });
  }
  return out;
}

export async function migrateOwnLegacyBundles(params: MigrateParams): Promise<MigrateResult> {
  const result: MigrateResult = {
    migratedConversations: 0, errors: [],
    attempted: 0, noStrongholdKey: 0, decryptFailed: 0, rpcFailed: 0,
  };
  if (params.ownLegacyDeviceIds.length === 0) {
    console.info('[crypto-migration] no legacy device-ids to consider — skipping');
    return result;
  }

  // RLS already filters this query down to rows where one of the recipient
  // device-ids belongs to us. We do NOT pre-filter by `ownLegacyDeviceIds`
  // anymore — historically, users can have device rows server-side whose
  // private key is no longer in the local vault (fresh OS install, vault
  // wiped, etc.). Conversely, the vault may hold a key for a device-id the
  // server forgot. The decisive check is "do we have the matching private
  // key in stronghold?" — which we evaluate per row below.
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const { data: rowsRaw, error } = await (params.client as any)
    .from('conversation_keys')
    .select('conversation_id, key_version, recipient_device_id, sender_device_id, sender_user_id, encrypted_key, nonce')
    .is('recipient_user_id', null)
    .not('recipient_device_id', 'is', null);
  if (error) throw error;
  const rows = (rowsRaw ?? []) as LegacyRow[];
  console.info('[crypto-migration] legacy rows visible to me: ' + rows.length);
  if (rows.length === 0) return result;

  const senderDeviceIds = Array.from(new Set(rows.map((r) => r.sender_device_id).filter(Boolean)));
  const senderMap = await fetchSenderPubKeys(params.client, senderDeviceIds);

  for (const row of rows) {
    result.attempted += 1;
    const ownPriv = params.ownLegacyDevicePrivateKeys[row.recipient_device_id];
    if (!ownPriv) {
      result.noStrongholdKey += 1;
      result.errors.push({
        conversationId: row.conversation_id,
        reason: 'no private key in stronghold for device ' + row.recipient_device_id.slice(0, 8),
      });
      continue;
    }
    const sender = senderMap.get(row.sender_device_id);
    if (!sender) {
      result.decryptFailed += 1;
      result.errors.push({ conversationId: row.conversation_id, reason: 'sender device not found' });
      continue;
    }
    let convKey: Uint8Array;
    try {
      convKey = await decryptFrom(
        pgHexToBytes(row.encrypted_key), pgHexToBytes(row.nonce),
        sender.publicKey, ownPriv,
      );
    } catch (err) {
      result.decryptFailed += 1;
      result.errors.push({
        conversationId: row.conversation_id,
        reason: err instanceof Error ? err.message : String(err),
      });
      continue;
    }
    const wrapped = await encryptFor(convKey, params.ownNewPublicKey, params.ownNewPrivateKey);
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const { error: rpcError } = await (params.client as any).rpc('migrate_user_key_recipients', {
      p_conv_id: row.conversation_id,
      p_user_id: params.ownUserId,
      p_key_version: row.key_version,
      p_bundles: [{
        encrypted_key: hexNoPrefix(wrapped.ciphertext),
        nonce: hexNoPrefix(wrapped.nonce),
        sender_user_id: row.sender_user_id ?? params.ownUserId,
      }],
    });
    if (rpcError) {
      result.rpcFailed += 1;
      result.errors.push({
        conversationId: row.conversation_id,
        reason: (rpcError as Error).message ?? String(rpcError),
      });
      continue;
    }
    result.migratedConversations += 1;
  }

  console.info(
    '[crypto-migration] result:',
    'attempted=' + result.attempted,
    'migrated=' + result.migratedConversations,
    'noKey=' + result.noStrongholdKey,
    'decryptFail=' + result.decryptFailed,
    'rpcFail=' + result.rpcFailed,
  );
  return result;
}