/**
 * Zero-knowledge key management with wrapped-DEK recovery (Phase 2).
 *
 * Key model:
 *   - DEK (Data Encryption Key): a random AES-256-GCM key used to encrypt/decrypt
 *     all user data. Lives in memory only.
 *   - KEK (Key Encryption Key): derived from the password (or recovery phrase)
 *     via PBKDF2. Used to wrap (encrypt) the DEK for storage on the server.
 *
 * The server stores two wrapped forms of the DEK:
 *   - password_wrapped_dek: DEK encrypted under KEK(password)
 *   - recovery_wrapped_dek: DEK encrypted under KEK(recovery_phrase)
 *
 * Recovery: derive KEK(recovery_phrase) → unwrap the DEK → re-wrap under the
 * new password's KEK. The server never holds the DEK or any KEK.
 */

import { encrypt, decrypt, type CipherPayload } from './cipher';

const PBKDF2_ITERATIONS = 150_000;
const KEY_BITS = 256;

const encoder = new TextEncoder();

// Domain-separation salts for each PBKDF2 purpose.
export const AUTH_SALT = 'normogen-auth-v1';
export const PASSWORD_KEK_SALT = 'normogen-kek-password-v1';
export const RECOVERY_KEK_SALT = 'normogen-kek-recovery-v1';

// ---------------------------------------------------------------------------
// Low-level: PBKDF2 derivation helpers
// ---------------------------------------------------------------------------

async function pbkdf2DeriveBits(password: string, salt: string): Promise<ArrayBuffer> {
  const key = await crypto.subtle.importKey(
    'raw',
    encoder.encode(password) as BufferSource,
    'PBKDF2',
    false,
    ['deriveBits'],
  );
  return crypto.subtle.deriveBits(
    { name: 'PBKDF2', salt: encoder.encode(salt) as BufferSource, iterations: PBKDF2_ITERATIONS, hash: 'SHA-256' },
    key,
    KEY_BITS,
  );
}

function base64(bytes: Uint8Array): string {
  let bin = '';
  for (const b of bytes) bin += String.fromCharCode(b);
  return btoa(bin);
}

/** Derive the auth secret (base64) from the password — sent to the server. */
export async function deriveAuthSecret(password: string): Promise<string> {
  const bits = await pbkdf2DeriveBits(password, AUTH_SALT);
  return base64(new Uint8Array(bits));
}

/** Derive a KEK (extractable AES-GCM key) from a password or recovery phrase. */
async function deriveKEK(secret: string, salt: string): Promise<CryptoKey> {
  const bits = await pbkdf2DeriveBits(secret, salt);
  return crypto.subtle.importKey(
    'raw',
    bits,
    { name: 'AES-GCM' },
    true, // extractable so we can export for wrapping
    ['encrypt', 'decrypt'],
  );
}

// ---------------------------------------------------------------------------
// DEK generation + wrapping
// ---------------------------------------------------------------------------

/** Generate a random AES-256-GCM DEK (extractable for wrapping). */
export async function generateDek(): Promise<CryptoKey> {
  return crypto.subtle.generateKey(
    { name: 'AES-GCM', length: KEY_BITS },
    true, // extractable so we can export raw bytes for wrapping
    ['encrypt', 'decrypt'],
  );
}

/** Export a DEK to raw bytes, base64-encode, then encrypt (wrap) under a KEK. */
export async function wrapDek(dek: CryptoKey, kek: CryptoKey): Promise<CipherPayload> {
  const rawDek = await crypto.subtle.exportKey('raw', dek);
  // Base64-encode the raw bytes so they survive the encrypt/decrypt string cycle.
  return encrypt(base64(new Uint8Array(rawDek)), kek);
}

/** Decrypt (unwrap) a wrapped DEK and import as an AES-GCM key. */
export async function unwrapDek(payload: CipherPayload, kek: CryptoKey): Promise<CryptoKey> {
  const rawDekB64 = await decrypt(payload, kek);
  // The wrapped DEK was encrypted as a base64 string of the raw key bytes.
  const rawDek = Uint8Array.from(atob(rawDekB64), (c) => c.charCodeAt(0));
  return crypto.subtle.importKey(
    'raw',
    rawDek as BufferSource,
    { name: 'AES-GCM' },
    true, // extractable so rewrapDek can export it for re-wrapping
    ['encrypt', 'decrypt'],
  );
}

/** Re-wrap a DEK under a new KEK (for password change / post-recovery). */
export async function rewrapDek(
  dek: CryptoKey,
  newSecret: string,
  salt: string = PASSWORD_KEK_SALT,
): Promise<CipherPayload> {
  const newKek = await deriveKEK(newSecret, salt);
  return wrapDek(dek, newKek);
}

// ---------------------------------------------------------------------------
// High-level flows
// ---------------------------------------------------------------------------

export interface EncryptionSetup {
  authSecret: string;
  dek: CryptoKey;
  passwordWrappedDek: CipherPayload;
  recoveryWrappedDek?: CipherPayload;
  recoveryKekHash?: string; // base64 hash of the recovery KEK proof (for server verification)
}

/**
 * Full setup at registration: derive auth secret, generate a DEK, wrap it under
 * the password KEK and (optionally) the recovery-phrase KEK.
 */
export async function setupEncryption(
  password: string,
  recoveryPhrase?: string,
): Promise<EncryptionSetup> {
  const authSecret = await deriveAuthSecret(password);
  const dek = await generateDek();

  const passwordKek = await deriveKEK(password, PASSWORD_KEK_SALT);
  const passwordWrappedDek = await wrapDek(dek, passwordKek);

  let recoveryWrappedDek: CipherPayload | undefined;
  let recoveryKekHash: string | undefined;

  if (recoveryPhrase) {
    const recoveryKek = await deriveKEK(recoveryPhrase, RECOVERY_KEK_SALT);
    recoveryWrappedDek = await wrapDek(dek, recoveryKek);
    // A proof the server can verify: derive a separate value from the recovery
    // phrase and hash it. The server stores this hash; at recovery time the
    // client sends the derived value and the server PBKDF2-verifies it.
    // We send the recovery auth proof (base64 of a PBKDF2 derivation).
    recoveryKekHash = await deriveAuthSecret(recoveryPhrase); // reuse the auth derivation as the proof
  }

  return { authSecret, dek, passwordWrappedDek, recoveryWrappedDek, recoveryKekHash };
}

export interface UnlockResult {
  authSecret: string;
  dek: CryptoKey;
}

/**
 * At login: derive the auth secret and unwrap the DEK from the password-wrapped form.
 */
export async function unlockWithPassword(
  password: string,
  passwordWrappedDek: CipherPayload,
): Promise<UnlockResult> {
  const authSecret = await deriveAuthSecret(password);
  const passwordKek = await deriveKEK(password, PASSWORD_KEK_SALT);
  const dek = await unwrapDek(passwordWrappedDek, passwordKek);
  return { authSecret, dek };
}

/**
 * At recovery: unwrap the DEK from the recovery-wrapped form using the recovery phrase.
 */
export async function unlockWithRecovery(
  recoveryPhrase: string,
  recoveryWrappedDek: CipherPayload,
): Promise<CryptoKey> {
  const recoveryKek = await deriveKEK(recoveryPhrase, RECOVERY_KEK_SALT);
  return unwrapDek(recoveryWrappedDek, recoveryKek);
}

// ---------------------------------------------------------------------------
// In-memory key store (unchanged from Phase 1)
// ---------------------------------------------------------------------------

let currentEncKey: CryptoKey | null = null;

export function setEncKey(key: CryptoKey): void {
  currentEncKey = key;
}

export function getEncKey(): CryptoKey | null {
  return currentEncKey;
}

export function clearEncKey(): void {
  currentEncKey = null;
}

export function hasEncKey(): boolean {
  return currentEncKey !== null;
}

// ---------------------------------------------------------------------------
// Backward compat: the old deriveAuthAndEncKeys (Phase 1 direct-from-password
// enc key). Kept for tests; new code uses setupEncryption/unlockWithPassword.
// ---------------------------------------------------------------------------

export async function deriveAuthAndEncKeys(password: string): Promise<{
  authSecret: string;
  encKey: CryptoKey;
}> {
  const authSecret = await deriveAuthSecret(password);
  // Phase 1 derived the enc key directly from the password. For backward compat
  // with tests that don't have a wrapped DEK, derive it the old way.
  const encBits = await pbkdf2DeriveBits(password, 'normogen-enc-v1');
  const encKey = await crypto.subtle.importKey(
    'raw',
    encBits,
    { name: 'AES-GCM' },
    false,
    ['encrypt', 'decrypt'],
  );
  return { authSecret, encKey };
}