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 ensureLegacyMigrated(p.userId).catch((err) => {
    console.warn('legacy conv-key migration failed', err);
  });
  return { publicKey: kp.publicKey, recoveryCode };
}

// Background, idempotent re-wrap of own legacy conv-key bundles for the new
// per-user identity. Safe to call repeatedly: the underlying RPC uses
// ON CONFLICT DO NOTHING. Triggered on every successful unlock so users who
// upgraded to 0.18.0 (where the .eq(null) bug made setup-time migration a
// no-op) auto-recover on the next launch.
export async function ensureLegacyMigrated(userId: string): Promise<void> {
  const priv = await devLocalSecretStore.getSecret(cacheKey(userId));
  if (!priv) return;
  const pub = await derivePublicKey(priv);
  await runLegacyMigration(userId, priv, pub);
}

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);
  void ensureLegacyMigrated(p.userId).catch((err) => {
    console.warn('legacy conv-key migration failed', err);
  });
  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 ?? '';
}

// Returned by ensureLegacyMigrated and SecurityCenter's manual retry. Lets
// the UI surface "X conv-keys re-wrapped, Y stuck because no key in vault."
export interface LegacyMigrationReport {
  serverDevices: number;
  strongholdKeysFromServerDevices: number;
  strongholdKeysFromBundleScan: number;
  attempted: number;
  migrated: number;
  noStrongholdKey: number;
  decryptFailed: number;
  rpcFailed: number;
}

async function runLegacyMigration(
  userId: string,
  ownNewPriv: Uint8Array,
  ownNewPub: Uint8Array,
): Promise<LegacyMigrationReport> {
  const report: LegacyMigrationReport = {
    serverDevices: 0,
    strongholdKeysFromServerDevices: 0,
    strongholdKeysFromBundleScan: 0,
    attempted: 0,
    migrated: 0,
    noStrongholdKey: 0,
    decryptFailed: 0,
    rpcFailed: 0,
  };

  const devices = await listOwnDevices(supabase);
  report.serverDevices = devices.length;
  const ownLegacyDevicePrivateKeys: Record<string, Uint8Array> = {};

  // 1) Try every server-listed device first.
  for (const d of devices) {
    const k = await devLocalSecretStore.getSecret(`chatapp.priv.${userId}.${d.id}`);
    if (k) ownLegacyDevicePrivateKeys[d.id] = k;
  }
  report.strongholdKeysFromServerDevices = Object.keys(ownLegacyDevicePrivateKeys).length;

  // 2) Scan our visible un-migrated conversation_keys for distinct
  // recipient_device_ids and probe stronghold for each. This catches the case
  // where a device row was deleted server-side but its key remains locally,
  // OR where listOwnDevices missed a device because of an RLS edge.
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  const { data: scanRowsRaw } = await (supabase as any)
    .from('conversation_keys')
    .select('recipient_device_id')
    .is('recipient_user_id', null)
    .not('recipient_device_id', 'is', null);
  const scanIds = Array.from(new Set(
    ((scanRowsRaw ?? []) as { recipient_device_id: string }[])
      .map((r) => r.recipient_device_id)
      .filter((id): id is string => !!id),
  ));
  for (const id of scanIds) {
    if (ownLegacyDevicePrivateKeys[id]) continue;
    const k = await devLocalSecretStore.getSecret(`chatapp.priv.${userId}.${id}`);
    if (k) {
      ownLegacyDevicePrivateKeys[id] = k;
      report.strongholdKeysFromBundleScan += 1;
    }
  }

  console.info(
    '[crypto-migration] vault scan:',
    'serverDevices=' + report.serverDevices,
    'keysFromServerList=' + report.strongholdKeysFromServerDevices,
    'extraKeysFromBundleScan=' + report.strongholdKeysFromBundleScan,
  );

  const ids = Object.keys(ownLegacyDevicePrivateKeys);
  if (ids.length === 0) {
    console.warn('[crypto-migration] no legacy private keys in vault — nothing to migrate');
    return report;
  }

  const result = await migrateOwnLegacyBundles({
    client: supabase,
    ownUserId: userId,
    ownNewPublicKey: ownNewPub,
    ownNewPrivateKey: ownNewPriv,
    ownLegacyDeviceIds: ids,
    ownLegacyDevicePrivateKeys,
  });

  report.attempted = result.attempted;
  report.migrated = result.migratedConversations;
  report.noStrongholdKey = result.noStrongholdKey;
  report.decryptFailed = result.decryptFailed;
  report.rpcFailed = result.rpcFailed;
  return report;
}

// Public wrapper for SecurityCenter's "Migration erneut versuchen" button.
// Returns a structured report so the UI can render numbers and reasons.
export async function retryLegacyMigration(userId: string): Promise<LegacyMigrationReport> {
  const priv = await devLocalSecretStore.getSecret(cacheKey(userId));
  if (!priv) throw new Error('user key not cached locally — re-login required');
  const pub = await derivePublicKey(priv);
  return runLegacyMigration(userId, priv, pub);
}

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