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

import {
  decryptWithConvKey,
  encryptWithConvKey,
  generateConvKey,
  unwrapConvKey,
  wrapConvKeyForRecipient,
} from '../crypto/sessionKeys';
import { fetchPeerPublicKeys } from '../auth/userKey';
import { bytesToPgHex, pgHexToBytes } from '../supabase/bytea';
import type { AppSupabaseClient } from '../supabase/client';

function rawFrom(client: AppSupabaseClient, table: string) {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  return (client as unknown as { from: (t: string) => any }).from(table);
}

export interface OwnUserCtx {
  userId: string;
  privateKey: Uint8Array;
}

export interface ConvKeyHandle {
  conversationId: string;
  keyVersion: number;
  key: Uint8Array;
}

const cache = new Map<string, ConvKeyHandle>();
const cacheKey = (convId: string, v: number) => convId + '@' + v;

export function clearConvKeyCache(): void { cache.clear(); }

async function listMemberPublicKeys(
  client: AppSupabaseClient,
  conversationId: string,
): Promise<{ userId: string; publicKey: Uint8Array }[]> {
  const { data: members, error } = await client
    .from('conversation_members')
    .select('user_id, accepted')
    .eq('conversation_id', conversationId);
  if (error) throw error;
  const ids = (members ?? []).filter((m) => m.accepted).map((m) => m.user_id);
  return fetchPeerPublicKeys(client, ids);
}

async function fetchActiveKeyVersion(
  client: AppSupabaseClient,
  conversationId: string,
): Promise<number> {
  const { data, error } = await rawFrom(client, 'conversations')
    .select('active_key_version').eq('id', conversationId).single();
  if (error) throw error;
  return (data as { active_key_version: number }).active_key_version;
}

interface SenderInfo { senderUserId: string; senderPublicKey: Uint8Array }

async function fetchKeyBundle(
  client: AppSupabaseClient,
  conversationId: string,
  ownUserId: string,
  keyVersion: number,
): Promise<{ encryptedKey: Uint8Array; nonce: Uint8Array; sender: SenderInfo } | null> {
  const { data, error } = await rawFrom(client, 'conversation_keys')
    .select('encrypted_key, nonce, sender_user_id')
    .eq('conversation_id', conversationId)
    .eq('recipient_user_id', ownUserId)
    .eq('key_version', keyVersion)
    .maybeSingle();
  if (error) throw error;
  if (!data) return null;
  const row = data as { encrypted_key: string; nonce: string; sender_user_id: string };
  const peers = await fetchPeerPublicKeys(client, [row.sender_user_id]);
  const sender = peers[0];
  if (!sender) throw new Error('sender public key missing');
  return {
    encryptedKey: pgHexToBytes(row.encrypted_key),
    nonce: pgHexToBytes(row.nonce),
    sender: { senderUserId: sender.userId, senderPublicKey: sender.publicKey },
  };
}

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

export async function bootstrapConvKey(
  client: AppSupabaseClient,
  conversationId: string,
  own: OwnUserCtx,
  keyVersion: number,
): Promise<ConvKeyHandle> {
  const convKey = generateConvKey();
  const recipients = await listMemberPublicKeys(client, conversationId);
  if (recipients.length === 0) throw new Error('cannot bootstrap conv key — no recipients');
  const bundles: Array<{ recipient_user_id: string; encrypted_key: string; nonce: string }> = [];
  for (const r of recipients) {
    const wrapped = await wrapConvKeyForRecipient(convKey, r.publicKey, own.privateKey);
    bundles.push({
      recipient_user_id: r.userId,
      encrypted_key: hexNoPrefix(wrapped.ciphertext),
      nonce: hexNoPrefix(wrapped.nonce),
    });
  }
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const rpc = (client as unknown as { rpc: (n: string, p: object) => Promise<{ error: any }> }).rpc;
  const { error } = await rpc.call(client, 'share_conv_keys', {
    p_conv_id: conversationId,
    p_sender_device_id: null,
    p_sender_user_id: own.userId,
    p_key_version: keyVersion,
    p_bundles: bundles,
  });
  if (error) throw error;
  const handle = { conversationId, keyVersion, key: convKey };
  cache.set(cacheKey(conversationId, keyVersion), handle);
  return handle;
}

export async function getOrCreateConvKey(
  client: AppSupabaseClient,
  conversationId: string,
  own: OwnUserCtx,
): Promise<ConvKeyHandle> {
  const version = await fetchActiveKeyVersion(client, conversationId);
  const cached = cache.get(cacheKey(conversationId, version));
  if (cached) return cached;
  const bundle = await fetchKeyBundle(client, conversationId, own.userId, version);
  if (bundle) {
    const key = await unwrapConvKey(
      bundle.encryptedKey, bundle.nonce, bundle.sender.senderPublicKey, own.privateKey,
    );
    const handle = { conversationId, keyVersion: version, key };
    cache.set(cacheKey(conversationId, version), handle);
    return handle;
  }
  const { count, error: cntErr } = await rawFrom(client, 'conversation_keys')
    .select('recipient_user_id', { count: 'exact', head: true })
    .eq('conversation_id', conversationId)
    .eq('key_version', version);
  if (cntErr) throw cntErr;
  if ((count ?? 0) > 0) {
    // Rows exist for this version, but none for me. Either I lost the device-key
    // that originally received my bundle, or my own bundle was wiped by the
    // 0.18.0 reset_user_key bug. Either way, the only way out is to mint a fresh
    // conv-key at version+1 and wrap it for everyone we can. Old messages stay
    // unreadable for me; new ones flow.
    console.info('[conv-key] no bundle for me at v' + version + ' — auto-rotating');
    return rotateConvKey(client, conversationId, own);
  }
  return bootstrapConvKey(client, conversationId, own, version);
}

// Mints a fresh conv-key at active_key_version + 1 and wraps it for every
// accepted member. Per-user bundles take priority; for members lacking a
// user_keys row we fall back to per-device wrapping (one bundle per device)
// so peers still on the legacy 0.17.x client can decrypt with their device
// private key. Caller must own a copy of their private key in `own`.
export async function rotateConvKey(
  client: AppSupabaseClient,
  conversationId: string,
  own: OwnUserCtx,
): Promise<ConvKeyHandle> {
  const currentVersion = await fetchActiveKeyVersion(client, conversationId);
  const newVersion = currentVersion + 1;

  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) throw new Error('cannot rotate — no accepted members');

  const userKeys = await fetchPeerPublicKeys(client, memberIds);
  const userKeyByUserId = new Map(userKeys.map((k) => [k.userId, k.publicKey]));
  const missingUserKeyMembers = memberIds.filter((id) => !userKeyByUserId.has(id));

  let legacyDevices: { userId: string; deviceId: string; publicKey: Uint8Array }[] = [];
  if (missingUserKeyMembers.length > 0) {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const { data: devs, error: dErr } = await (client as any)
      .from('devices')
      .select('id, user_id, public_key')
      .in('user_id', missingUserKeyMembers)
      .not('public_key', 'is', null);
    if (dErr) throw dErr;
    legacyDevices = (devs ?? []).map((d: { id: string; user_id: string; public_key: string }) => ({
      userId: d.user_id,
      deviceId: d.id,
      publicKey: pgHexToBytes(d.public_key),
    }));
  }

  const convKey = generateConvKey();
  const bundles: Array<{
    recipient_user_id?: string;
    recipient_device_id?: string;
    encrypted_key: string;
    nonce: string;
  }> = [];
  for (const k of userKeys) {
    const wrapped = await wrapConvKeyForRecipient(convKey, k.publicKey, own.privateKey);
    bundles.push({
      recipient_user_id: k.userId,
      encrypted_key: hexNoPrefix(wrapped.ciphertext),
      nonce: hexNoPrefix(wrapped.nonce),
    });
  }
  for (const d of legacyDevices) {
    const wrapped = await wrapConvKeyForRecipient(convKey, d.publicKey, own.privateKey);
    bundles.push({
      recipient_device_id: d.deviceId,
      encrypted_key: hexNoPrefix(wrapped.ciphertext),
      nonce: hexNoPrefix(wrapped.nonce),
    });
  }
  if (bundles.length === 0) {
    throw new Error('cannot rotate — no peers have a public key (no user_keys, no devices)');
  }

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const rpc = (client as unknown as { rpc: (n: string, p: object) => Promise<{ error: any }> }).rpc;
  const { error } = await rpc.call(client, 'rotate_conv_key', {
    p_conv_id: conversationId,
    p_sender_user_id: own.userId,
    p_new_version: newVersion,
    p_bundles: bundles,
  });
  if (error) throw error;

  const handle = { conversationId, keyVersion: newVersion, key: convKey };
  cache.set(cacheKey(conversationId, newVersion), handle);
  console.info(
    '[conv-key] rotated conversation ' + conversationId.slice(0, 8) +
    ' from v' + currentVersion + ' to v' + newVersion +
    ' — wrapped for ' + userKeys.length + ' user-keys + ' + legacyDevices.length + ' legacy devices',
  );
  return handle;
}

export async function tryGetConvKey(
  client: AppSupabaseClient,
  conversationId: string,
  ownUserId: string,
  ownPrivateKey: Uint8Array,
  keyVersion: number,
): Promise<ConvKeyHandle | null> {
  const cached = cache.get(cacheKey(conversationId, keyVersion));
  if (cached) return cached;
  const bundle = await fetchKeyBundle(client, conversationId, ownUserId, keyVersion);
  if (!bundle) return null;
  const key = await unwrapConvKey(
    bundle.encryptedKey, bundle.nonce, bundle.sender.senderPublicKey, ownPrivateKey,
  );
  const handle = { conversationId, keyVersion, key };
  cache.set(cacheKey(conversationId, keyVersion), handle);
  return handle;
}

export async function shareConvKeyToUser(
  client: AppSupabaseClient,
  conversationId: string,
  recipientUserId: string,
  recipientPublicKey: Uint8Array,
  own: OwnUserCtx,
): Promise<void> {
  const version = await fetchActiveKeyVersion(client, conversationId);
  const handle =
    cache.get(cacheKey(conversationId, version)) ??
    (await tryGetConvKey(client, conversationId, own.userId, own.privateKey, version));
  if (!handle) throw new Error('cannot share conv key — own user does not have it yet');
  const wrapped = await wrapConvKeyForRecipient(handle.key, recipientPublicKey, own.privateKey);
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const rpc = (client as unknown as { rpc: (n: string, p: object) => Promise<{ error: any }> }).rpc;
  const { error } = await rpc.call(client, 'share_conv_keys', {
    p_conv_id: conversationId,
    p_sender_device_id: null,
    p_sender_user_id: own.userId,
    p_key_version: version,
    p_bundles: [{
      recipient_user_id: recipientUserId,
      encrypted_key: hexNoPrefix(wrapped.ciphertext),
      nonce: hexNoPrefix(wrapped.nonce),
    }],
  });
  if (error) throw error;
}

export { decryptWithConvKey, encryptWithConvKey };