import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { kdfRootKey, kdfChainKey, deriveInitialRootKey } from '../src/keys.js';

const crypto = new SubtleCryptoProvider();

describe('KDF Chain Functions', () => {
  describe('kdfRootKey', () => {
    test('produces 32-byte root key and 32-byte chain key', async () => {
      const rootKey = crypto.randomBytes(32);
      const dhOutput = crypto.randomBytes(32);

      const { newRootKey, chainKey } = await kdfRootKey(crypto, rootKey, dhOutput);
      expect(newRootKey.length).toBe(32);
      expect(chainKey.length).toBe(32);
    });

    test('new root key differs from input root key', async () => {
      const rootKey = crypto.randomBytes(32);
      const dhOutput = crypto.randomBytes(32);

      const { newRootKey } = await kdfRootKey(crypto, rootKey, dhOutput);
      expect(newRootKey).not.toEqual(rootKey);
    });

    test('root key and chain key differ from each other', async () => {
      const rootKey = crypto.randomBytes(32);
      const dhOutput = crypto.randomBytes(32);

      const { newRootKey, chainKey } = await kdfRootKey(crypto, rootKey, dhOutput);
      expect(newRootKey).not.toEqual(chainKey);
    });

    test('deterministic: same inputs produce same outputs', async () => {
      const rootKey = new Uint8Array(32).fill(0x11);
      const dhOutput = new Uint8Array(32).fill(0x22);

      const a = await kdfRootKey(crypto, rootKey, dhOutput);
      const b = await kdfRootKey(crypto, rootKey, dhOutput);
      expect(a.newRootKey).toEqual(b.newRootKey);
      expect(a.chainKey).toEqual(b.chainKey);
    });

    test('different DH output produces different keys', async () => {
      const rootKey = crypto.randomBytes(32);
      const dh1 = crypto.randomBytes(32);
      const dh2 = crypto.randomBytes(32);

      const a = await kdfRootKey(crypto, rootKey, dh1);
      const b = await kdfRootKey(crypto, rootKey, dh2);
      expect(a.newRootKey).not.toEqual(b.newRootKey);
      expect(a.chainKey).not.toEqual(b.chainKey);
    });
  });

  describe('kdfChainKey', () => {
    test('produces 32-byte chain key and 32-byte message key', async () => {
      const chainKey = crypto.randomBytes(32);

      const { newChainKey, messageKey } = await kdfChainKey(crypto, chainKey);
      expect(newChainKey.length).toBe(32);
      expect(messageKey.length).toBe(32);
    });

    test('chain key and message key differ', async () => {
      const chainKey = crypto.randomBytes(32);

      const { newChainKey, messageKey } = await kdfChainKey(crypto, chainKey);
      expect(newChainKey).not.toEqual(messageKey);
    });

    test('chain ratchet is one-way: cannot derive previous chain key', async () => {
      const ck0 = crypto.randomBytes(32);
      const { newChainKey: ck1 } = await kdfChainKey(crypto, ck0);
      const { newChainKey: ck2 } = await kdfChainKey(crypto, ck1);

      // All three are different
      expect(ck0).not.toEqual(ck1);
      expect(ck1).not.toEqual(ck2);
      expect(ck0).not.toEqual(ck2);
    });

    test('deterministic: same input produces same output', async () => {
      const chainKey = new Uint8Array(32).fill(0x33);

      const a = await kdfChainKey(crypto, chainKey);
      const b = await kdfChainKey(crypto, chainKey);
      expect(a.newChainKey).toEqual(b.newChainKey);
      expect(a.messageKey).toEqual(b.messageKey);
    });

    test('sequential chain steps produce unique message keys', async () => {
      let ck = crypto.randomBytes(32);
      const messageKeys: Uint8Array[] = [];

      for (let i = 0; i < 10; i++) {
        const { newChainKey, messageKey } = await kdfChainKey(crypto, ck);
        messageKeys.push(messageKey);
        ck = newChainKey;
      }

      // All message keys should be unique
      for (let i = 0; i < messageKeys.length; i++) {
        for (let j = i + 1; j < messageKeys.length; j++) {
          expect(messageKeys[i]).not.toEqual(messageKeys[j]);
        }
      }
    });
  });

  describe('deriveInitialRootKey', () => {
    test('produces 32-byte root key from multiple DH outputs', async () => {
      const secrets = [
        crypto.randomBytes(32),
        crypto.randomBytes(32),
        crypto.randomBytes(32),
      ];

      const rootKey = await deriveInitialRootKey(crypto, secrets);
      expect(rootKey.length).toBe(32);
    });

    test('works with 3 secrets (no one-time prekey)', async () => {
      const secrets = [
        crypto.randomBytes(32),
        crypto.randomBytes(32),
        crypto.randomBytes(32),
      ];

      const rootKey = await deriveInitialRootKey(crypto, secrets);
      expect(rootKey.length).toBe(32);
    });

    test('works with 4 secrets (with one-time prekey)', async () => {
      const secrets = [
        crypto.randomBytes(32),
        crypto.randomBytes(32),
        crypto.randomBytes(32),
        crypto.randomBytes(32),
      ];

      const rootKey = await deriveInitialRootKey(crypto, secrets);
      expect(rootKey.length).toBe(32);
    });

    test('deterministic: same secrets produce same root key', async () => {
      const secrets = [
        new Uint8Array(32).fill(0xaa),
        new Uint8Array(32).fill(0xbb),
        new Uint8Array(32).fill(0xcc),
      ];

      const a = await deriveInitialRootKey(crypto, secrets);
      const b = await deriveInitialRootKey(crypto, secrets);
      expect(a).toEqual(b);
    });

    test('different secrets produce different root keys', async () => {
      const secretsA = [crypto.randomBytes(32), crypto.randomBytes(32), crypto.randomBytes(32)];
      const secretsB = [crypto.randomBytes(32), crypto.randomBytes(32), crypto.randomBytes(32)];

      const a = await deriveInitialRootKey(crypto, secretsA);
      const b = await deriveInitialRootKey(crypto, secretsB);
      expect(a).not.toEqual(b);
    });

    test('adding a 4th secret changes the root key', async () => {
      const base = [
        new Uint8Array(32).fill(0x11),
        new Uint8Array(32).fill(0x22),
        new Uint8Array(32).fill(0x33),
      ];

      const without = await deriveInitialRootKey(crypto, base);
      const withExtra = await deriveInitialRootKey(crypto, [...base, new Uint8Array(32).fill(0x44)]);
      expect(without).not.toEqual(withExtra);
    });
  });
});