import { describe, it, expect } from 'vitest';
import {
  setupEncryption,
  unlockWithPassword,
  unlockWithRecovery,
  rewrapDek,
  encryptJson,
  decryptJson,
  encrypt,
  decrypt,
  setEncKey,
  clearEncKey,
  hasEncKey,
  generateIdentityKeyPair,
  wrapIdentityPrivateKey,
  unwrapIdentityPrivateKey,
  setIdentityPrivate,
  clearIdentityPrivate,
  hasIdentityPrivate,
  generateProfileDek,
  wrapProfileDek,
  unwrapProfileDek,
  setProfileDek,
  getProfileDek,
  clearProfileDeks,
} from './index';

/** X25519 in Web Crypto is a recent addition (Node 16+, modern browsers). Some
 * test runtimes (older jsdom/Node) may not implement it; detect once and skip
 * the identity-keypair tests rather than fail spuriously. */
async function x25519Supported(): Promise<boolean> {
  try {
    await crypto.subtle.generateKey(
      { name: 'ECDH', namedCurve: 'X25519' },
      true,
      ['deriveBits'],
    );
    return true;
  } catch {
    return false;
  }
}

describe('zero-knowledge crypto lifecycle', () => {
  const password = 'my-secret-password';
  const recoveryPhrase = 'my-recovery-phrase';
  const sampleData = { name: 'Aspirin', dosage: '100mg', frequency: 'daily' };

  it('full lifecycle: setup → encrypt → unlock → decrypt → recover → rewrap', async () => {
    // 1. Setup (register): generate DEK, wrap under password + recovery.
    const setup = await setupEncryption(password, recoveryPhrase);
    expect(setup.dek).toBeDefined();
    expect(setup.passwordWrappedDek.data).not.toBe('');
    expect(setup.recoveryWrappedDek).toBeDefined();
    expect(setup.recoveryKekHash).toBeDefined();

    // 2. Encrypt data with the DEK.
    const encrypted = await encryptJson(sampleData, setup.dek);
    // Verify it's ciphertext, not plaintext.
    expect(encrypted.data).not.toContain('Aspirin');
    expect(encrypted.iv).not.toBe('');

    // 3. Decrypt with the same DEK — round-trip.
    const decrypted = await decryptJson<typeof sampleData>(encrypted, setup.dek);
    expect(decrypted.name).toBe('Aspirin');
    expect(decrypted.dosage).toBe('100mg');

    // 4. Simulate page reload: DEK is lost.
    clearEncKey();
    expect(hasEncKey()).toBe(false);

    // 5. Unlock with password (the unlock screen flow).
    const unlockResult = await unlockWithPassword(password, setup.passwordWrappedDek);
    setEncKey(unlockResult.dek);
    expect(hasEncKey()).toBe(true);

    // 6. The unlocked DEK can decrypt the same data.
    const decryptedAfterUnlock = await decryptJson<typeof sampleData>(
      encrypted,
      unlockResult.dek,
    );
    expect(decryptedAfterUnlock.name).toBe('Aspirin');

    // 7. Recover via recovery phrase (forgot password flow).
    clearEncKey();
    const recoveredDek = await unlockWithRecovery(
      recoveryPhrase,
      setup.recoveryWrappedDek!,
    );
    const decryptedAfterRecovery = await decryptJson<typeof sampleData>(
      encrypted,
      recoveredDek,
    );
    expect(decryptedAfterRecovery.name).toBe('Aspirin');

    // 8. Rewrap under a new password (post-recovery).
    const newPassword = 'my-new-password';
    const rewrapped = await rewrapDek(recoveredDek, newPassword);
    const unlockWithNew = await unlockWithPassword(newPassword, rewrapped);
    const decryptedAfterRewrap = await decryptJson<typeof sampleData>(
      encrypted,
      unlockWithNew.dek,
    );
    expect(decryptedAfterRewrap.name).toBe('Aspirin');

    clearEncKey();
  });

  it('wrong password fails to unlock', async () => {
    const setup = await setupEncryption(password);
    await expect(
      unlockWithPassword('wrong-password', setup.passwordWrappedDek),
    ).rejects.toThrow();
  });

  it('wrong recovery phrase fails to unlock', async () => {
    const setup = await setupEncryption(password, recoveryPhrase);
    await expect(
      unlockWithRecovery('wrong-phrase', setup.recoveryWrappedDek!),
    ).rejects.toThrow();
  });

  it('data encrypted under one key cannot be decrypted by another', async () => {
    const setupA = await setupEncryption('user-a-password');
    const setupB = await setupEncryption('user-b-password');
    const encrypted = await encryptJson(sampleData, setupA.dek);

    // User B's key cannot decrypt User A's data.
    await expect(decryptJson(encrypted, setupB.dek)).rejects.toThrow();
  });
});

describe('account X25519 identity keypair (Phase A1)', () => {
  const password = 'my-secret-password';

  it('full lifecycle: generate → wrap → unlock DEK → unwrap → deriveBits', async () => {
    if (!(await x25519Supported())) {
      console.log('[crypto] skipped (X25519 unsupported in this runtime)');
      return;
    }

    // 1. Register: set up the account DEK and generate the identity keypair.
    const setup = await setupEncryption(password);
    const { publicKey, privateKey } = await generateIdentityKeyPair();
    expect(publicKey).not.toBe('');
    setIdentityPrivate(privateKey);
    expect(hasIdentityPrivate()).toBe(true);

    // 2. Wrap the private key under the account DEK and "store" it.
    const wrapped = await wrapIdentityPrivateKey(privateKey, setup.dek);
    expect(wrapped.data).not.toBe('');

    // 3. Simulate page reload: both keys are lost from memory.
    clearIdentityPrivate();
    expect(hasIdentityPrivate()).toBe(false);

    // 4. Unlock the DEK with the password (unlock screen flow).
    const { dek } = await unlockWithPassword(password, setup.passwordWrappedDek);

    // 5. Unwrap the identity private key from the stored wrapped form.
    const recovered = await unwrapIdentityPrivateKey(wrapped, dek);
    setIdentityPrivate(recovered);
    expect(hasIdentityPrivate()).toBe(true);

    // 6. The recovered private key is a USABLE X25519 key: it can derive a
    // shared secret against the original public key. Import the public key
    // and ECDH-derive — this proves the round-trip preserved a working key.
    const pubBytes = Uint8Array.from(atob(publicKey), (c) => c.charCodeAt(0));
    const pubKey = await crypto.subtle.importKey(
      'raw',
      pubBytes as BufferSource,
      { name: 'ECDH', namedCurve: 'X25519' },
      false,
      [],
    );
    const sharedSecret = await crypto.subtle.deriveBits(
      { name: 'ECDH', public: pubKey },
      recovered,
      256,
    );
    expect(sharedSecret.byteLength).toBe(32);

    clearIdentityPrivate();
  });

  it('a wrapped private key cannot be unwrapped with the wrong DEK', async () => {
    if (!(await x25519Supported())) {
      console.log('[crypto] skipped (X25519 unsupported in this runtime)');
      return;
    }
    // DEK for account A; a different DEK for account B.
    const setupA = await setupEncryption('user-a-password');
    const setupB = await setupEncryption('user-b-password');
    const { privateKey } = await generateIdentityKeyPair();
    const wrapped = await wrapIdentityPrivateKey(privateKey, setupA.dek);

    // User B's DEK cannot unwrap user A's identity private key.
    await expect(unwrapIdentityPrivateKey(wrapped, setupB.dek)).rejects.toThrow();
  });

  it('two identity keypairs derive different shared secrets with the same peer', async () => {
    if (!(await x25519Supported())) {
      console.log('[crypto] skipped (X25519 unsupported in this runtime)');
      return;
    }
    // Sanity: distinct keypairs must produce distinct ECDH outputs.
    const peer = await generateIdentityKeyPair();
    const peerPubBytes = Uint8Array.from(atob(peer.publicKey), (c) => c.charCodeAt(0));
    const peerPub = await crypto.subtle.importKey(
      'raw',
      peerPubBytes as BufferSource,
      { name: 'ECDH', namedCurve: 'X25519' },
      false,
      [],
    );

    const a = await generateIdentityKeyPair();
    const b = await generateIdentityKeyPair();
    const secretA = await crypto.subtle.deriveBits(
      { name: 'ECDH', public: peerPub },
      a.privateKey,
      256,
    );
    const secretB = await crypto.subtle.deriveBits(
      { name: 'ECDH', public: peerPub },
      b.privateKey,
      256,
    );

    const aHex = Array.from(new Uint8Array(secretA)).map((x) => x.toString(16)).join('');
    const bHex = Array.from(new Uint8Array(secretB)).map((x) => x.toString(16)).join('');
    expect(aHex).not.toBe(bHex);
  });
});

describe('per-profile DEK isolation (Phase A2)', () => {
  const password = 'my-secret-password';

  it('two profiles get distinct DEKs, each wrapped under the account DEK', async () => {
    // One account DEK (unlocked from the password).
    const setup = await setupEncryption(password);
    const accountDek = setup.dek;

    // Two profiles, each with its own random DEK.
    const childDek = await generateProfileDek();
    const petDek = await generateProfileDek();

    // Each profile DEK is wrapped under the account DEK for storage.
    const childWrapped = await wrapProfileDek(childDek, accountDek);
    const petWrapped = await wrapProfileDek(petDek, accountDek);

    // Both wrapped blobs decrypt cleanly under the account DEK.
    const childUnwrapped = await unwrapProfileDek(childWrapped, accountDek);
    const petUnwrapped = await unwrapProfileDek(petWrapped, accountDek);
    expect(childUnwrapped).toBeDefined();
    expect(petUnwrapped).toBeDefined();

    // Data encrypted under the child profile's DEK cannot be decrypted by the
    // pet profile's DEK (and vice versa) — this is the per-profile isolation
    // guarantee that makes the parent/child and pet cases safe.
    const childData = await encrypt('child-medication', childDek);
    await expect(decrypt(childData, petDek)).rejects.toThrow();
    const petData = await encrypt('pet-weight', petDek);
    await expect(decrypt(petData, childDek)).rejects.toThrow();

    // ...but each decrypts with its own key.
    expect(await decrypt(childData, childDek)).toBe('child-medication');
    expect(await decrypt(petData, petDek)).toBe('pet-weight');
  });

  it('the in-memory profile-DEK store maps profile ids to DEKs', async () => {
    const childDek = await generateProfileDek();
    const petDek = await generateProfileDek();
    setProfileDek('profile_child', childDek);
    setProfileDek('profile_pet', petDek);
    // Each id resolves to the DEK that was stored under it; unknown ids resolve to null.
    expect(getProfileDek('profile_child')).toBe(childDek);
    expect(getProfileDek('profile_pet')).toBe(petDek);
    expect(getProfileDek('profile_unknown')).toBeNull();
    // Clearing drops everything.
    clearProfileDeks();
    expect(getProfileDek('profile_child')).toBeNull();
  });
});