import { describe, test, expect, beforeEach } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import {
  ShadeSessionManager,
  generateOneTimePreKeys,
  createPreKeyBundle,
} from '../src/index.js';

const crypto = new SubtleCryptoProvider();

describe('ShadeSessionManager', () => {
  let alice: ShadeSessionManager;
  let bob: ShadeSessionManager;
  let aliceStorage: MemoryStorage;
  let bobStorage: MemoryStorage;

  beforeEach(async () => {
    aliceStorage = new MemoryStorage();
    bobStorage = new MemoryStorage();
    alice = new ShadeSessionManager(crypto, aliceStorage);
    bob = new ShadeSessionManager(crypto, bobStorage);
    await alice.initialize();
    await bob.initialize();
  });

  /** Helper: establish a session from Alice to Bob */
  async function establishSession() {
    // Bob generates one-time prekeys
    const otpks = await bob.generateOneTimePreKeys(10);

    // Bob creates a bundle (with one-time prekey)
    const bobBundle = await bob.createPreKeyBundle();
    // Add a one-time prekey to the bundle (in real life the prekey server does this)
    const otpk = otpks[0];
    bobBundle.oneTimePreKey = { keyId: otpk.keyId, publicKey: otpk.keyPair.publicKey };

    // Alice initiates session
    await alice.initSessionFromBundle('bob', bobBundle);
  }

  // ─── Initialization ────────────────────────────────────────

  describe('initialization', () => {
    test('generates identity keys on first init', async () => {
      const pub = alice.getPublicIdentity();
      expect(pub.signingKey.length).toBe(32);
      expect(pub.dhKey.length).toBe(32);
    });

    test('reuses identity keys on second init', async () => {
      const pub1 = alice.getPublicIdentity();
      const alice2 = new ShadeSessionManager(crypto, aliceStorage);
      await alice2.initialize();
      const pub2 = alice2.getPublicIdentity();
      expect(pub1.signingKey).toEqual(pub2.signingKey);
      expect(pub1.dhKey).toEqual(pub2.dhKey);
    });

    test('creates a prekey bundle', async () => {
      const bundle = await alice.createPreKeyBundle();
      expect(bundle.identitySigningKey.length).toBe(32);
      expect(bundle.identityDHKey.length).toBe(32);
      expect(bundle.signedPreKey.keyId).toBe(1);
      expect(bundle.signedPreKey.publicKey.length).toBe(32);
      expect(bundle.signedPreKey.signature.length).toBe(64);
    });

    test('generates one-time prekeys', async () => {
      const keys = await alice.generateOneTimePreKeys(5);
      expect(keys.length).toBe(5);
      expect(await aliceStorage.getOneTimePreKeyCount()).toBe(5);
    });
  });

  // ─── Full Conversation ─────────────────────────────────────

  describe('full conversation via managers', () => {
    test('Alice sends to Bob, Bob replies', async () => {
      await establishSession();

      // Alice → Bob (first message = PreKeyMessage)
      const env1 = await alice.encrypt('bob', 'Hello Bob!');
      expect(env1.type).toBe('prekey');

      const plain1 = await bob.decrypt('alice', env1);
      expect(plain1).toBe('Hello Bob!');

      // Alice → Bob (second message = RatchetMessage)
      const env2 = await alice.encrypt('bob', 'Still me');
      expect(env2.type).toBe('ratchet');

      const plain2 = await bob.decrypt('alice', env2);
      expect(plain2).toBe('Still me');

      // Bob → Alice (reply)
      const env3 = await bob.encrypt('alice', 'Hi Alice!');
      expect(env3.type).toBe('ratchet');

      const plain3 = await alice.decrypt('bob', env3);
      expect(plain3).toBe('Hi Alice!');
    });

    test('extended conversation with many turns', async () => {
      await establishSession();

      // Alice sends first message to establish Bob's session
      const first = await alice.encrypt('bob', 'init');
      await bob.decrypt('alice', first);

      for (let i = 0; i < 20; i++) {
        const senderMgr = i % 2 === 0 ? alice : bob;
        const receiverMgr = i % 2 === 0 ? bob : alice;
        const senderAddr = i % 2 === 0 ? 'bob' : 'alice';
        const receiverAddr = i % 2 === 0 ? 'alice' : 'bob';

        const text = `Turn ${i}`;
        const env = await senderMgr.encrypt(senderAddr, text);
        const plain = await receiverMgr.decrypt(receiverAddr, env);
        expect(plain).toBe(text);
      }
    });

    test('burst messages then reply', async () => {
      await establishSession();

      // Alice sends 5 messages
      const envelopes = [];
      for (let i = 0; i < 5; i++) {
        envelopes.push(await alice.encrypt('bob', `msg-${i}`));
      }

      // Bob decrypts all 5
      for (let i = 0; i < 5; i++) {
        expect(await bob.decrypt('alice', envelopes[i])).toBe(`msg-${i}`);
      }

      // Bob replies
      const reply = await bob.encrypt('alice', 'Got all 5!');
      expect(await alice.decrypt('bob', reply)).toBe('Got all 5!');
    });
  });

  // ─── Prekey Rotation ──────────────────────────────────────

  describe('prekey rotation', () => {
    test('rotated signed prekey works for new sessions', async () => {
      // Bob rotates his signed prekey
      await bob.rotateSignedPreKey();

      // New bundle uses the new signed prekey
      const bundle = await bob.createPreKeyBundle();
      expect(bundle.signedPreKey.keyId).toBe(2);

      // Alice can still establish a session with the new bundle
      await alice.initSessionFromBundle('bob', bundle);
      const env = await alice.encrypt('bob', 'After rotation');
      const plain = await bob.decrypt('alice', env);
      expect(plain).toBe('After rotation');
    });

    test('old sessions continue working after rotation', async () => {
      await establishSession();

      // Establish session and exchange messages
      const env1 = await alice.encrypt('bob', 'Before rotation');
      await bob.decrypt('alice', env1);

      // Bob rotates
      await bob.rotateSignedPreKey();

      // Existing session still works (uses ratchet keys, not prekeys)
      const env2 = await bob.encrypt('alice', 'After rotation, same session');
      expect(await alice.decrypt('bob', env2)).toBe('After rotation, same session');
    });
  });

  // ─── Multi-Peer ───────────────────────────────────────────

  describe('multi-peer sessions', () => {
    test('Alice talks to Bob and Charlie simultaneously', async () => {
      const charlieStorage = new MemoryStorage();
      const charlie = new ShadeSessionManager(crypto, charlieStorage);
      await charlie.initialize();
      await charlie.generateOneTimePreKeys(5);

      // Establish Alice → Bob
      await establishSession();
      const envB = await alice.encrypt('bob', 'Hi Bob');
      await bob.decrypt('alice', envB);

      // Establish Alice → Charlie
      const charlieBundle = await charlie.createPreKeyBundle();
      const otpks = await charlie.generateOneTimePreKeys(1);
      charlieBundle.oneTimePreKey = { keyId: otpks[0].keyId, publicKey: otpks[0].keyPair.publicKey };
      await alice.initSessionFromBundle('charlie', charlieBundle);

      const envC = await alice.encrypt('charlie', 'Hi Charlie');
      expect(await charlie.decrypt('alice', envC)).toBe('Hi Charlie');

      // Both sessions work independently
      const envB2 = await alice.encrypt('bob', 'Still talking to you Bob');
      expect(await bob.decrypt('alice', envB2)).toBe('Still talking to you Bob');

      const envC2 = await alice.encrypt('charlie', 'And you Charlie');
      expect(await charlie.decrypt('alice', envC2)).toBe('And you Charlie');
    });
  });

  // ─── Error Cases ──────────────────────────────────────────

  describe('error cases', () => {
    test('encrypt to unknown peer throws NoSessionError', async () => {
      expect(alice.encrypt('nobody', 'test')).rejects.toThrow('No session');
    });

    test('decrypt from unknown peer throws NoSessionError', async () => {
      const fakeEnvelope = {
        type: 'ratchet' as const,
        content: {
          dhPublicKey: crypto.randomBytes(32),
          previousCounter: 0,
          counter: 0,
          ciphertext: crypto.randomBytes(48),
          nonce: crypto.randomBytes(12),
        },
        timestamp: Date.now(),
        senderAddress: 'nobody',
      };
      expect(alice.decrypt('nobody', fakeEnvelope)).rejects.toThrow('No session');
    });
  });
});