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) {
    throw new Error('Awaiting conversation key — another user must share it with this user.');
  }
  return bootstrapConvKey(client, conversationId, own, version);
}

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 };