import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  createSenderKey,
  buildDistribution,
  installDistribution,
  senderKeyEncrypt,
  senderKeyDecrypt,
} from '../src/sender-keys.js';
import type { GroupSession, SenderKeyState } from '../src/sender-keys.js';

const crypto = new SubtleCryptoProvider();

/**
 * Helpers: set up a group of N members who all know each other's sender keys.
 */
async function setupGroup(groupId: string, memberAddresses: string[]): Promise<{
  sessions: Map<string, GroupSession>;
}> {
  // Each member creates their own sender key state
  const senderStates = new Map<string, SenderKeyState>();
  for (const addr of memberAddresses) {
    senderStates.set(addr, await createSenderKey(crypto));
  }

  // Each member's session contains everyone's sender key
  const sessions = new Map<string, GroupSession>();
  for (const member of memberAddresses) {
    const session: GroupSession = { groupId, senderKeys: new Map() };

    // Add own sender key (with private signing key)
    session.senderKeys.set(member, senderStates.get(member)!);

    // Add public copies of everyone else's sender keys (without private key)
    for (const other of memberAddresses) {
      if (other === member) continue;
      const dist = buildDistribution(groupId, other, senderStates.get(other)!);
      installDistribution(session, dist);
    }

    sessions.set(member, session);
  }

  return { sessions };
}

describe('Sender Keys (group messaging)', () => {
  test('Alice sends to Bob + Charlie + Dave', async () => {
    const { sessions } = await setupGroup('group1', ['alice', 'bob', 'charlie', 'dave']);

    const msg = await senderKeyEncrypt(
      crypto,
      sessions.get('alice')!,
      'alice',
      new TextEncoder().encode('hello everyone'),
    );

    // All three recipients decrypt independently
    for (const recipient of ['bob', 'charlie', 'dave']) {
      const plain = await senderKeyDecrypt(crypto, sessions.get(recipient)!, 'group1', msg);
      expect(new TextDecoder().decode(plain)).toBe('hello everyone');
    }
  });

  test('multiple messages from same sender advance the chain', async () => {
    const { sessions } = await setupGroup('g', ['alice', 'bob']);

    for (let i = 0; i < 5; i++) {
      const msg = await senderKeyEncrypt(
        crypto,
        sessions.get('alice')!,
        'alice',
        new TextEncoder().encode(`msg ${i}`),
      );
      expect(msg.iteration).toBe(i);

      const plain = await senderKeyDecrypt(crypto, sessions.get('bob')!, 'g', msg);
      expect(new TextDecoder().decode(plain)).toBe(`msg ${i}`);
    }
  });

  test('messages from different senders use independent chains', async () => {
    const { sessions } = await setupGroup('g', ['alice', 'bob', 'charlie']);

    // Alice sends
    const aliceMsg = await senderKeyEncrypt(
      crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode('from alice'),
    );
    // Bob sends
    const bobMsg = await senderKeyEncrypt(
      crypto, sessions.get('bob')!, 'bob', new TextEncoder().encode('from bob'),
    );

    // Charlie decrypts both
    expect(new TextDecoder().decode(
      await senderKeyDecrypt(crypto, sessions.get('charlie')!, 'g', aliceMsg)
    )).toBe('from alice');
    expect(new TextDecoder().decode(
      await senderKeyDecrypt(crypto, sessions.get('charlie')!, 'g', bobMsg)
    )).toBe('from bob');
  });

  test('tampered ciphertext fails signature verification', async () => {
    const { sessions } = await setupGroup('g', ['alice', 'bob']);

    const msg = await senderKeyEncrypt(
      crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode('original'),
    );
    msg.ciphertext[0] ^= 0xff;

    expect(
      senderKeyDecrypt(crypto, sessions.get('bob')!, 'g', msg),
    ).rejects.toThrow(/Invalid signature|decryption/i);
  });

  test('wrong group ID fails', async () => {
    const { sessions } = await setupGroup('g', ['alice', 'bob']);

    const msg = await senderKeyEncrypt(
      crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode('hi'),
    );

    expect(
      senderKeyDecrypt(crypto, sessions.get('bob')!, 'wrong-group', msg),
    ).rejects.toThrow(/Group ID mismatch/);
  });

  test('out-of-order with small skip works', async () => {
    const { sessions } = await setupGroup('g', ['alice', 'bob']);

    // Alice sends 3 messages
    const messages = [];
    for (let i = 0; i < 3; i++) {
      messages.push(await senderKeyEncrypt(
        crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode(`msg ${i}`),
      ));
    }

    // Bob decrypts in order
    for (let i = 0; i < 3; i++) {
      const plain = await senderKeyDecrypt(crypto, sessions.get('bob')!, 'g', messages[i]!);
      expect(new TextDecoder().decode(plain)).toBe(`msg ${i}`);
    }
  });

  test('new member added later receives subsequent messages via fresh distribution', async () => {
    // Alice and Bob in a group
    const { sessions } = await setupGroup('g', ['alice', 'bob']);

    // Alice sends one message to Bob
    const msg1 = await senderKeyEncrypt(
      crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode('private msg'),
    );
    await senderKeyDecrypt(crypto, sessions.get('bob')!, 'g', msg1);

    // Charlie joins — Alice sends him her CURRENT sender key state
    const aliceStateForCharlie = sessions.get('alice')!.senderKeys.get('alice')!;
    const charlieSession: GroupSession = { groupId: 'g', senderKeys: new Map() };
    installDistribution(charlieSession, buildDistribution('g', 'alice', aliceStateForCharlie));

    // Alice sends another message
    const msg2 = await senderKeyEncrypt(
      crypto, sessions.get('alice')!, 'alice', new TextEncoder().encode('welcome charlie'),
    );

    // Charlie can decrypt (because he got the current chain state)
    const plain = await senderKeyDecrypt(crypto, charlieSession, 'g', msg2);
    expect(new TextDecoder().decode(plain)).toBe('welcome charlie');
  });
});