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;

// Clear the in-memory conv-key cache. Three modes:
//   * no args     → clear everything (e.g. on logout)
//   * convId only → clear all key-version entries for this conversation
//   * convId + v  → clear just the specific (conv, version) entry
//
// Callers that observe a peer rotation or a server-side conv-keys mutation
// MUST invalidate the affected entries so subsequent `getOrCreateConvKey` /
// `tryGetConvKey` calls re-fetch the canonical bundle from the server
// instead of returning a now-stale cached key.
export function clearConvKeyCache(conversationId?: string, keyVersion?: number): void {
  if (conversationId === undefined) {
    cache.clear();
    return;
  }
  if (keyVersion !== undefined) {
    cache.delete(cacheKey(conversationId, keyVersion));
    return;
  }
  for (const key of Array.from(cache.keys())) {
    if (key.startsWith(conversationId + '@')) cache.delete(key);
  }
}

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;

  // `share_conv_keys` uses `ON CONFLICT (conv, recipient_user_id, key_version)
  // DO NOTHING`. If a concurrent peer bootstrapped first at the same version,
  // OUR INSERTs were silently skipped server-side and the row on the server
  // holds THEIR conv-key, not ours. Trusting the locally-generated key here
  // would leave both clients with mutually un-decryptable bundles (each
  // encrypting/decrypting with its own key — exactly the bug that broke
  // conv aae12d84). Re-fetch our own bundle and unwrap to get the CANONICAL
  // server key. Whoever wrote first wins; the loser converges.
  const ownBundle = await fetchKeyBundle(client, conversationId, own.userId, keyVersion);
  if (!ownBundle) {
    throw new Error('bootstrapConvKey: own bundle missing after share_conv_keys');
  }
  const canonicalKey = await unwrapConvKey(
    ownBundle.encryptedKey,
    ownBundle.nonce,
    ownBundle.sender.senderPublicKey,
    own.privateKey,
  );
  const handle = { conversationId, keyVersion, key: canonicalKey };
  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) {
    try {
      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;
    } catch (err) {
      // A bundle exists for us but our current private key cannot unwrap it.
      // The most common cause is `reset_user_key`: a fresh user-key pair was
      // generated locally while the on-server bundle is still wrapped against
      // the previous public key. Treat this the same as "no bundle for me" —
      // mint a fresh conv-key at version+1 wrapped to our CURRENT key. Old
      // messages stay unreadable for us; new ones flow.
      console.warn(
        '[conv-key] unwrap own bundle failed at v' + version + ' — auto-rotating',
        err,
      );
    }
  }
  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 usable for me. Either I lost the
    // device-key that originally received my bundle, my own bundle was wiped
    // by the 0.18.0 reset_user_key bug, or my key was reset and the existing
    // bundle is unwrappable (handled in the try/catch above). 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 usable 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;
  let key: Uint8Array;
  try {
    key = await unwrapConvKey(
      bundle.encryptedKey, bundle.nonce, bundle.sender.senderPublicKey, ownPrivateKey,
    );
  } catch (err) {
    // Bundle exists but the current private key doesn't unwrap it (typically
    // after `reset_user_key`). Return null so the caller treats the message
    // as un-decryptable instead of throwing and killing the whole batch.
    // The conversation will be auto-rotated to a fresh key on the next send
    // or chat open via `getOrCreateConvKey`'s own recovery path.
    console.warn(
      '[conv-key] tryGetConvKey unwrap failed at v' + keyVersion +
        ' (conv=' + conversationId.slice(0, 8) + ') — marking as un-decryptable',
      err,
    );
    return null;
  }
  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 };