ChatApp/packages/shared/src/auth/userKey.ts

import type { AppSupabaseClient } from '../supabase/client';
import type { KdfParams } from '../crypto/userKey';

export type { KdfParams };

export interface UserKeyBlob {
  exists: true;
  sealedPrivateKey: Uint8Array;
  salt: Uint8Array;
  kdfParams: KdfParams;
  recoverySealedPrivateKey: Uint8Array | null;
  recoverySalt: Uint8Array | null;
  failedAttempts: number;
  failedRecoveryAttempts: number;
  recoveryLockedUntil: string | null;
  keyVersion: number;
}

export type UnlockResult =
  | { exists: false }
  | ({ exists: true; locked: false } & UserKeyBlob)
  | { exists: true; locked: true; lockedUntil: string };

function b64ToBytes(s: string): Uint8Array {
  const bin = atob(s);
  const out = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
  return out;
}

function bytesToB64(b: Uint8Array): string {
  let s = '';
  for (const v of b) s += String.fromCharCode(v);
  return btoa(s);
}

interface RpcCapable {
  rpc: (name: string, params: unknown) => Promise<{ data: unknown; error: unknown }>;
}

function rpc(client: AppSupabaseClient): RpcCapable {
  return client as unknown as RpcCapable;
}

export async function tryUnlockUserKey(
  client: AppSupabaseClient,
  userId: string,
): Promise<UnlockResult> {
  const { data, error } = await rpc(client).rpc('try_unlock_user_key', { p_user_id: userId });
  if (error) throw error;
  const d = data as Record<string, unknown>;
  if (!d?.exists) return { exists: false };
  if (d.locked) {
    return { exists: true, locked: true, lockedUntil: String(d.locked_until ?? '') };
  }
  return {
    exists: true,
    locked: false,
    sealedPrivateKey: b64ToBytes(String(d.sealed_private_key)),
    salt: b64ToBytes(String(d.salt)),
    kdfParams: d.kdf_params as KdfParams,
    recoverySealedPrivateKey: d.recovery_sealed_private_key
      ? b64ToBytes(String(d.recovery_sealed_private_key)) : null,
    recoverySalt: d.recovery_salt ? b64ToBytes(String(d.recovery_salt)) : null,
    failedAttempts: Number(d.failed_attempts ?? 0),
    failedRecoveryAttempts: Number(d.failed_recovery_attempts ?? 0),
    recoveryLockedUntil: (d.recovery_locked_until as string | null) ?? null,
    keyVersion: Number(d.key_version ?? 1),
  };
}

export async function fetchUserKeyBlob(
  client: AppSupabaseClient,
  userId: string,
): Promise<UnlockResult | null> {
  const res = await tryUnlockUserKey(client, userId);
  if (!res.exists) return null;
  return res;
}

export interface UploadParams {
  userId: string;
  publicKey: Uint8Array;
  sealedPrivateKey: Uint8Array;
  salt: Uint8Array;
  kdfParams: KdfParams;
  recoverySealedPrivateKey?: Uint8Array | null;
  recoverySalt?: Uint8Array | null;
}

export async function uploadUserKeyBlob(
  client: AppSupabaseClient,
  params: UploadParams,
): Promise<void> {
  // Non-destructive UPSERT — must NOT touch conversation_keys. Used for the
  // first-time PIN setup, PIN change, and recovery-code regeneration. The
  // 0.18.0–0.18.2 builds wired this to `reset_user_key` which DELETED every
  // legacy conv-key bundle for the user before the migration could re-wrap
  // them, leaving people unable to read or send. `upsert_user_key` writes
  // only the user_keys row and leaves conversation_keys alone.
  const { error } = await rpc(client).rpc('upsert_user_key', {
    p_user_id: params.userId,
    p_public_key_b64: bytesToB64(params.publicKey),
    p_sealed_private_b64: bytesToB64(params.sealedPrivateKey),
    p_salt_b64: bytesToB64(params.salt),
    p_kdf_params: params.kdfParams,
    p_recovery_sealed_b64: params.recoverySealedPrivateKey
      ? bytesToB64(params.recoverySealedPrivateKey) : null,
    p_recovery_salt_b64: params.recoverySalt ? bytesToB64(params.recoverySalt) : null,
  });
  if (error) throw error;
}

export async function resetUserKey(
  client: AppSupabaseClient,
  params: UploadParams,
): Promise<number> {
  const { data, error } = await rpc(client).rpc('reset_user_key', {
    p_user_id: params.userId,
    p_public_key_b64: bytesToB64(params.publicKey),
    p_sealed_private_b64: bytesToB64(params.sealedPrivateKey),
    p_salt_b64: bytesToB64(params.salt),
    p_kdf_params: params.kdfParams,
    p_recovery_sealed_b64: params.recoverySealedPrivateKey
      ? bytesToB64(params.recoverySealedPrivateKey) : null,
    p_recovery_salt_b64: params.recoverySalt ? bytesToB64(params.recoverySalt) : null,
  });
  if (error) throw error;
  return Number(data ?? 0);
}

export interface AttemptResult { failedAttempts: number; lockedUntil: string | null }

export async function recordPinAttempt(
  client: AppSupabaseClient,
  userId: string,
  success: boolean,
  recovery: boolean,
): Promise<AttemptResult> {
  const { data, error } = await rpc(client).rpc('record_pin_attempt', {
    p_user_id: userId, p_success: success, p_recovery: recovery,
  });
  if (error) throw error;
  const d = (data ?? {}) as Record<string, unknown>;
  return {
    failedAttempts: Number(d.failed_attempts ?? 0),
    lockedUntil: (d.locked_until as string | null) ?? null,
  };
}

export interface PeerPublicKey {
  userId: string;
  publicKey: Uint8Array;
  keyVersion: number;
}

export async function fetchPeerPublicKeys(
  client: AppSupabaseClient,
  userIds: string[],
): Promise<PeerPublicKey[]> {
  if (userIds.length === 0) return [];
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const { data, error } = await (client as any)
    .from('user_public_keys')
    .select('user_id, public_key, key_version')
    .in('user_id', userIds);
  if (error) throw error;
  return (data ?? []).map((row: { user_id: string; public_key: string; key_version: number }) => ({
    userId: row.user_id,
    publicKey: pgHexToBytes(row.public_key),
    keyVersion: row.key_version,
  }));
}

function pgHexToBytes(hex: string): Uint8Array {
  const s = hex.startsWith('\\x') ? hex.slice(2) : hex;
  const out = new Uint8Array(s.length / 2);
  for (let i = 0; i < out.length; i++) out[i] = parseInt(s.substr(i * 2, 2), 16);
  return out;
}