ChatApp/apps/desktop/src/lib/userIdentity.ts

import {
  fetchUserKeyBlob, listOwnDevices, recordPinAttempt,
  resetUserKey, tryUnlockUserKey, uploadUserKeyBlob,
} from '@chat-app/shared/auth';
import {
  generateRecoveryCode, generateUserKeyPair, normalizeRecoveryCode,
  openUserKey, sealUserKey,
} from '@chat-app/shared/crypto';
import { migrateOwnLegacyBundles } from '@chat-app/shared/chat';

import { devLocalSecretStore } from './secretStore';
import { supabase } from './supabase';

const cacheKey = (userId: string) => `chatapp.userpriv.${userId}`;

export interface SetupParams { userId: string; pin: string; withRecovery: boolean }
export interface SetupResult { publicKey: Uint8Array; recoveryCode: string | null }

export async function setupNewUserIdentity(p: SetupParams): Promise<SetupResult> {
  const kp = await generateUserKeyPair();
  const sealed = await sealUserKey({ privateKey: kp.privateKey, pin: p.pin });
  let recoveryCode: string | null = null;
  let recoverySealed: { sealedPrivateKey: Uint8Array; salt: Uint8Array } | null = null;
  if (p.withRecovery) {
    recoveryCode = await generateRecoveryCode();
    const r = await sealUserKey({ privateKey: kp.privateKey, pin: normalizeRecoveryCode(recoveryCode) });
    recoverySealed = { sealedPrivateKey: r.sealedPrivateKey, salt: r.salt };
  }
  await uploadUserKeyBlob(supabase, {
    userId: p.userId,
    publicKey: kp.publicKey,
    sealedPrivateKey: sealed.sealedPrivateKey,
    salt: sealed.salt,
    kdfParams: sealed.kdfParams,
    recoverySealedPrivateKey: recoverySealed?.sealedPrivateKey ?? null,
    recoverySalt: recoverySealed?.salt ?? null,
  });
  await devLocalSecretStore.setSecret(cacheKey(p.userId), kp.privateKey);
  void runLegacyMigration(p.userId, kp.privateKey, kp.publicKey).catch((err) => {
    console.warn('legacy conv-key migration failed', err);
  });
  return { publicKey: kp.publicKey, recoveryCode };
}

export interface UnlockParams { userId: string; pin: string; isRecoveryCode?: boolean }

export type UnlockOutcome =
  | { kind: 'unlocked' }
  | { kind: 'locked'; lockedUntil: string }
  | { kind: 'missing' };

export async function loadOrUnlockUserKey(p: UnlockParams): Promise<UnlockOutcome> {
  const remote = await tryUnlockUserKey(supabase, p.userId);
  if (!remote.exists) return { kind: 'missing' };
  if (remote.locked) return { kind: 'locked', lockedUntil: remote.lockedUntil };
  const secret = p.isRecoveryCode ? normalizeRecoveryCode(p.pin) : p.pin;
  const sealed = p.isRecoveryCode ? remote.recoverySealedPrivateKey : remote.sealedPrivateKey;
  const salt = p.isRecoveryCode ? remote.recoverySalt : remote.salt;
  if (!sealed || !salt) throw new Error('no recovery blob configured');
  let priv: Uint8Array;
  try {
    priv = await openUserKey({ sealed, pin: secret, salt, kdfParams: remote.kdfParams });
  } catch (err) {
    await recordPinAttempt(supabase, p.userId, false, p.isRecoveryCode === true).catch(() => {});
    throw err;
  }
  await recordPinAttempt(supabase, p.userId, true, p.isRecoveryCode === true).catch(() => {});
  await devLocalSecretStore.setSecret(cacheKey(p.userId), priv);
  return { kind: 'unlocked' };
}

export async function cachedUserKey(userId: string): Promise<Uint8Array | null> {
  return devLocalSecretStore.getSecret(cacheKey(userId));
}

export async function clearUserKeyCache(userId: string): Promise<void> {
  await devLocalSecretStore.removeSecret(cacheKey(userId));
}

export async function userKeyExistsRemotely(userId: string): Promise<boolean> {
  const blob = await fetchUserKeyBlob(supabase, userId);
  return blob !== null;
}

export async function changePin(params: {
  userId: string; oldPin: string; newPin: string;
}): Promise<void> {
  const cached = await cachedUserKey(params.userId);
  if (!cached) throw new Error('user key not cached locally — re-login required');
  const fresh = await sealUserKey({ privateKey: cached, pin: params.newPin });
  await uploadUserKeyBlob(supabase, {
    userId: params.userId,
    publicKey: await derivePublicKey(cached),
    sealedPrivateKey: fresh.sealedPrivateKey,
    salt: fresh.salt,
    kdfParams: fresh.kdfParams,
  });
  void params.oldPin; // unused: cached key already proves old PIN was correct
}

export async function regenerateRecoveryCode(params: { userId: string }): Promise<string> {
  const cached = await cachedUserKey(params.userId);
  if (!cached) throw new Error('user key not cached locally');
  const blob = await fetchUserKeyBlob(supabase, params.userId);
  if (!blob || !blob.exists || blob.locked) throw new Error('cannot regenerate recovery while locked');
  const recoveryCode = await generateRecoveryCode();
  const sealed = await sealUserKey({ privateKey: cached, pin: normalizeRecoveryCode(recoveryCode) });
  await uploadUserKeyBlob(supabase, {
    userId: params.userId,
    publicKey: await derivePublicKey(cached),
    sealedPrivateKey: blob.sealedPrivateKey,
    salt: blob.salt,
    kdfParams: blob.kdfParams,
    recoverySealedPrivateKey: sealed.sealedPrivateKey,
    recoverySalt: sealed.salt,
  });
  return recoveryCode;
}

export async function resetIdentity(params: { userId: string; pin: string }): Promise<string> {
  await clearUserKeyCache(params.userId);
  // resetUserKey already deletes legacy bundles; setupNewUserIdentity uploads
  // the brand-new blob via the same RPC (UPSERT). After this, no migration
  // pass runs because all legacy conv-keys are gone.
  await resetUserKey(supabase, {
    userId: params.userId,
    publicKey: new Uint8Array(32), // overwritten by next upload
    sealedPrivateKey: new Uint8Array(40),
    salt: new Uint8Array(16),
    kdfParams: { algo: 'argon2id', preset: 'moderate', opslimit: 1, memlimit: 1 },
  });
  const setup = await setupNewUserIdentity({ userId: params.userId, pin: params.pin, withRecovery: true });
  return setup.recoveryCode ?? '';
}

async function runLegacyMigration(
  userId: string,
  ownNewPriv: Uint8Array,
  ownNewPub: Uint8Array,
): Promise<void> {
  const devices = await listOwnDevices(supabase);
  if (devices.length === 0) return;
  const ownLegacyDevicePrivateKeys: Record<string, Uint8Array> = {};
  for (const d of devices) {
    const k = await devLocalSecretStore.getSecret(`chatapp.priv.${userId}.${d.id}`);
    if (k) ownLegacyDevicePrivateKeys[d.id] = k;
  }
  const ids = Object.keys(ownLegacyDevicePrivateKeys);
  if (ids.length === 0) return;
  await migrateOwnLegacyBundles({
    client: supabase,
    ownUserId: userId,
    ownNewPublicKey: ownNewPub,
    ownNewPrivateKey: ownNewPriv,
    ownLegacyDeviceIds: ids,
    ownLegacyDevicePrivateKeys,
  });
}

async function derivePublicKey(privateKey: Uint8Array): Promise<Uint8Array> {
  const sodium = (await import('libsodium-wrappers-sumo')).default;
  await sodium.ready;
  return sodium.crypto_scalarmult_base(privateKey);
}