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

const crypto = new SubtleCryptoProvider();
const enc = new TextEncoder();
const dec = new TextDecoder();

/**
 * Establish a *bidirectional* parent session: Alice→Bob X3DH, then one
 * Alice→Bob message Bob decrypts so Bob also has a session for 'alice'.
 * Both sides then hold the peer's pinned identity DH key — the input the
 * stream handshake derives from.
 */
async function bidirectionalPair() {
  const aliceStorage = new MemoryStorage();
  const bobStorage = new MemoryStorage();
  const alice = new ShadeSessionManager(crypto, aliceStorage);
  const bob = new ShadeSessionManager(crypto, bobStorage);
  await alice.initialize();
  await bob.initialize();

  const otpks = await bob.generateOneTimePreKeys(4);
  const bundle = await bob.createPreKeyBundle();
  bundle.oneTimePreKey = { keyId: otpks[0]!.keyId, publicKey: otpks[0]!.keyPair.publicKey };
  await alice.initSessionFromBundle('bob', bundle);

  const hello = await alice.encrypt('bob', 'parent-hello');
  expect(await bob.decrypt('alice', hello)).toBe('parent-hello');

  return { alice, bob, aliceStorage, bobStorage };
}

/** Run the full STREAM_OPEN / STREAM_OPEN_ACK handshake between managers. */
async function openStreamPair(alice: ShadeSessionManager, bob: ShadeSessionManager) {
  const begun = await alice.beginStream('bob'); // initiator
  const accepted = await bob.acceptStream('alice', begun.streamId, begun.ephemeralPublicKey);
  const aliceStream = await begun.complete(accepted.ephemeralPublicKey);
  return { aliceStream, bobStream: accepted.stream, streamId: begun.streamId };
}

describe('streaming sub-ratchet (V4.11)', () => {
  let alice: ShadeSessionManager;
  let bob: ShadeSessionManager;
  let aliceStorage: MemoryStorage;

  beforeEach(async () => {
    ({ alice, bob, aliceStorage } = await bidirectionalPair());
  });

  test('both sides derive the same stream root (round-trips frames)', async () => {
    const { aliceStream, bobStream } = await openStreamPair(alice, bob);

    // Initiator → responder (first frame triggers responder DH step).
    const f1 = await aliceStream.seal(enc.encode('log line 1'));
    expect(dec.decode(await bobStream.open(f1))).toBe('log line 1');

    // Responder → initiator (now responder may seal).
    const r1 = await bobStream.seal(enc.encode('command-response 1'));
    expect(dec.decode(await aliceStream.open(r1))).toBe('command-response 1');
  });

  test('two streams to the same peer get independent roots', async () => {
    const s1 = await openStreamPair(alice, bob);
    const s2 = await openStreamPair(alice, bob);
    expect(s1.streamId).not.toEqual(s2.streamId);

    const a = await s1.aliceStream.seal(enc.encode('on stream 1'));
    // A frame from stream 1 must not decrypt on stream 2's ratchet.
    await expect(s2.bobStream.open(a)).rejects.toBeInstanceOf(DecryptionError);
    // …but does on its own.
    expect(dec.decode(await s1.bobStream.open(a))).toBe('on stream 1');
  });

  test('R2/R3: long one-directional burst stays correct and memory-bounded', async () => {
    const { aliceStream, bobStream } = await openStreamPair(alice, bob);
    const N = 5000;

    // Capture a live receive-chain key buffer to prove forward secrecy:
    // ratchetDecrypt zeroizes the previous chain key in place.
    await bobStream.open(await aliceStream.seal(enc.encode('frame-0')));
    const bobSession = (bobStream as unknown as { session: { receiveChain: { chainKey: Uint8Array }; skippedKeys: Map<string, Uint8Array> } }).session;
    const staleChainKey = bobSession.receiveChain.chainKey;
    const staleCopy = staleChainKey.slice();
    expect(staleCopy.some((b) => b !== 0)).toBe(true);

    for (let i = 1; i < N; i++) {
      const wire = await aliceStream.seal(enc.encode(`frame-${i}`));
      expect(dec.decode(await bobStream.open(wire))).toBe(`frame-${i}`);
    }

    // In-order delivery ⇒ zero skipped keys retained across 5k frames.
    expect(bobSession.skippedKeys.size).toBe(0);
    // The chain key in use at frame 0 was overwritten (forward secrecy).
    expect(staleChainKey.every((b) => b === 0)).toBe(true);
  });

  test('R5: opening/using/closing a stream never touches the parent session', async () => {
    const before = await aliceStorage.getSession('bob');
    const snapshot = JSON.stringify({
      root: Array.from(before!.rootKey),
      sendCtr: before!.sendChain.counter,
      prevCtr: before!.previousSendCounter,
    });

    const { aliceStream, bobStream } = await openStreamPair(alice, bob);
    for (let i = 0; i < 200; i++) {
      await bobStream.open(await aliceStream.seal(enc.encode(`x${i}`)));
    }
    await aliceStream.close();
    await bobStream.close();

    const after = await aliceStorage.getSession('bob');
    expect(
      JSON.stringify({
        root: Array.from(after!.rootKey),
        sendCtr: after!.sendChain.counter,
        prevCtr: after!.previousSendCounter,
      }),
    ).toBe(snapshot);

    // Parent HTTP path still works after the stream lifecycle.
    const env = await alice.encrypt('bob', 'after-stream');
    expect(await bob.decrypt('alice', env)).toBe('after-stream');
  });

  test('R1: replayed / rewound frame is rejected', async () => {
    const { aliceStream, bobStream } = await openStreamPair(alice, bob);
    const f1 = await aliceStream.seal(enc.encode('once'));
    expect(dec.decode(await bobStream.open(f1))).toBe('once');
    // Re-delivering the exact same sealed frame must fail.
    await expect(bobStream.open(f1)).rejects.toBeInstanceOf(DecryptionError);
  });

  test('close() zeroizes and blocks further use; idempotent', async () => {
    const { aliceStream, bobStream } = await openStreamPair(alice, bob);
    await aliceStream.close();
    await aliceStream.close(); // idempotent
    expect(aliceStream.closed).toBe(true);
    await expect(aliceStream.seal(enc.encode('nope'))).rejects.toBeInstanceOf(
      StreamClosedError,
    );
    // The peer end is unaffected by our local close.
    expect(bobStream.closed).toBe(false);
  });

  test('handshake is mutually authenticated against pinned identities', async () => {
    // A third party (mallory) with its own identity cannot stand in for
    // bob: alice derives against bob's pinned identity key, so a
    // handshake completed with mallory's ephemeral yields a different
    // root and frames fail to open.
    const mStorage = new MemoryStorage();
    const mallory = new ShadeSessionManager(crypto, mStorage);
    await mallory.initialize();
    // Give mallory a parent session label so acceptStream has identity
    // material, but pinned to the WRONG (alice) identity vs what alice
    // pinned for 'bob'.
    const otpks = await mallory.generateOneTimePreKeys(2);
    const mb = await mallory.createPreKeyBundle();
    mb.oneTimePreKey = { keyId: otpks[0]!.keyId, publicKey: otpks[0]!.keyPair.publicKey };
    await alice.initSessionFromBundle('mallory', mb);
    const helo = await alice.encrypt('mallory', 'hi');
    await mallory.decrypt('alice', helo);

    const begun = await alice.beginStream('bob');
    const mAccept = await mallory.acceptStream('alice', begun.streamId, begun.ephemeralPublicKey);
    const aliceStream = await begun.complete(mAccept.ephemeralPublicKey);
    const frame = await aliceStream.seal(enc.encode('secret'));
    await expect(mAccept.stream.open(frame)).rejects.toBeInstanceOf(DecryptionError);
  });
});