/**
 * 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);
}

// ---------------------------------------------------------------------------
// Per-profile DEK wrap/unwrap (Phase A2).
//
// Each profile owns a random AES-256-GCM DEK that encrypts all of that
// profile's data. The profile DEK is wrapped under the *account* DEK (not a
// KEK derived from a secret) and stored on the server. At login/unlock the
// client unwraps each profile DEK using the account DEK. See
// docs/adr/multi-person-sharing.md §1.
// ---------------------------------------------------------------------------

/** Generate a fresh random profile DEK (a random AES-256-GCM key). */
export async function generateProfileDek(): Promise<CryptoKey> {
  return generateDek();
}

/** Wrap a profile DEK under the account DEK. The account DEK is a CryptoKey
 * usable for AES-GCM encrypt/decrypt, so it serves directly as the wrap key. */
export async function wrapProfileDek(
  profileDek: CryptoKey,
  accountDek: CryptoKey,
): Promise<CipherPayload> {
  return wrapDek(profileDek, accountDek);
}

/** Unwrap a profile DEK from its account-wrapped form. Inverse of
 * wrapProfileDek. Throws on tamper / wrong account DEK. */
export async function unwrapProfileDek(
  payload: CipherPayload,
  accountDek: CryptoKey,
): Promise<CryptoKey> {
  return unwrapDek(payload, accountDek);
}

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

// ---------------------------------------------------------------------------
// Account X25519 identity keypair (Phase A1, multi-person sharing prep).
//
// Each account gets an X25519 keypair at registration. The PUBLIC half is
// stored in plaintext (public keys are not secret); other accounts will wrap
// profile DEKs to it in Phase B. The PRIVATE half is wrapped (AES-GCM
// encrypted) under the account DEK and stored on the server as opaque
// ciphertext. The server never sees the private key.
//
// See docs/adr/multi-person-sharing.md §1.
// ---------------------------------------------------------------------------

/** Generate an X25519 keypair for the account. Returns the public key as a
 * base64 raw string (for the server to store) and the private key as a
 * CryptoKey (kept in memory, wrapped before upload). */
export async function generateIdentityKeyPair(): Promise<{
  publicKey: string;
  privateKey: CryptoKey;
}> {
  const { publicKey, privateKey } = (await crypto.subtle.generateKey(
    { name: 'ECDH', namedCurve: 'X25519' },
    true, // extractable so we can export the private key for wrapping
    ['deriveBits'],
  )) as CryptoKeyPair;
  const rawPublic = await crypto.subtle.exportKey('raw', publicKey);
  return { publicKey: base64(new Uint8Array(rawPublic)), privateKey };
}

/** Wrap (encrypt) the X25519 private key under the account DEK. The DEK is an
 * AES-GCM key, so we reuse the string-cipher from cipher.ts: export the private
 * key to raw bytes, base64-encode, and encrypt. */
export async function wrapIdentityPrivateKey(
  privateKey: CryptoKey,
  dek: CryptoKey,
): Promise<CipherPayload> {
  const rawPrivate = await crypto.subtle.exportKey('raw', privateKey);
  return encrypt(base64(new Uint8Array(rawPrivate)), dek);
}

/** Unwrap (decrypt) the X25519 private key from its DEK-wrapped form. Inverse
 * of wrapIdentityPrivateKey. The result is importable for ECDH deriveBits. */
export async function unwrapIdentityPrivateKey(
  payload: CipherPayload,
  dek: CryptoKey,
): Promise<CryptoKey> {
  const rawPrivateB64 = await decrypt(payload, dek);
  const rawPrivate = Uint8Array.from(atob(rawPrivateB64), (c) => c.charCodeAt(0));
  return crypto.subtle.importKey(
    'raw',
    rawPrivate as BufferSource,
    { name: 'ECDH', namedCurve: 'X25519' },
    true, // extractable so it can be re-wrapped if ever needed
    ['deriveBits'],
  );
}

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

// ---------------------------------------------------------------------------
// In-memory identity-private-key store (Phase A1). Same lifecycle as the DEK:
// held only for the authenticated session, cleared on logout, re-derived from
// the wrapped form (under the account DEK) on unlock. Not consumed until Phase B.
// ---------------------------------------------------------------------------

let currentIdentityPrivate: CryptoKey | null = null;

export function setIdentityPrivate(key: CryptoKey): void {
  currentIdentityPrivate = key;
}

export function getIdentityPrivate(): CryptoKey | null {
  return currentIdentityPrivate;
}

export function clearIdentityPrivate(): void {
  currentIdentityPrivate = null;
}

export function hasIdentityPrivate(): boolean {
  return currentIdentityPrivate !== null;
}

// ---------------------------------------------------------------------------
// In-memory per-profile DEK store (Phase A2). One DEK per profile, plus a
// notion of the "active" profile whose DEK encrypts/decrypts the data the UI
// is currently showing. Same lifecycle as the account DEK: rebuilt from the
// server-stored wrapped forms on unlock, cleared on logout.
// ---------------------------------------------------------------------------

const profileDeks: Map<string, CryptoKey> = new Map();
let activeProfileId: string | null = null;

export function setProfileDek(profileId: string, dek: CryptoKey): void {
  profileDeks.set(profileId, dek);
}

export function getProfileDek(profileId: string): CryptoKey | null {
  return profileDeks.get(profileId) ?? null;
}

export function setActiveProfileId(profileId: string): void {
  activeProfileId = profileId;
}

export function getActiveProfileId(): string | null {
  return activeProfileId;
}

/** The active profile's DEK — the key all encrypt/decrypt call sites use. */
export function getActiveProfileDek(): CryptoKey | null {
  if (activeProfileId === null) return null;
  return profileDeks.get(activeProfileId) ?? null;
}

export function clearProfileDeks(): void {
  profileDeks.clear();
  activeProfileId = 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 };
}