import { ShadeSessionManager } from '../../packages/shade-core/src/index.js';
import { SubtleCryptoProvider, MemoryStorage } from '../../packages/shade-crypto-web/src/index.js';
import { ShadeWebSocket } from '../../packages/shade-transport/src/index.js';

const crypto = new SubtleCryptoProvider();

async function main() {
  console.log('=== Shade WebSocket Tunnel ===\n');

  // Set up Alice and Bob
  const alice = new ShadeSessionManager(crypto, new MemoryStorage());
  const bob = new ShadeSessionManager(crypto, new MemoryStorage());
  await alice.initialize();
  await bob.initialize();

  // Establish session (skipping prekey server for brevity)
  const otpks = await bob.generateOneTimePreKeys(5);
  const bundle = await bob.createPreKeyBundle();
  bundle.oneTimePreKey = { keyId: otpks[0].keyId, publicKey: otpks[0].keyPair.publicKey };
  await alice.initSessionFromBundle('bob', bundle);

  // Pair of WebSockets connected to each other
  // Use Bun's built-in WebSocket server + client for the demo
  const port = 19851;
  const messages: string[] = [];

  const server = Bun.serve({
    port,
    fetch(req, srv) {
      if (srv.upgrade(req)) return;
      return new Response('upgrade failed', { status: 500 });
    },
    websocket: {
      async message(ws, message) {
        // Server side decrypts (acts as Bob)
        const bytes = message instanceof Uint8Array ? message : new Uint8Array(Buffer.from(message as string));
        // Need to manually use ShadeWebSocket or do raw decoding
        console.log(`[Bob received ${bytes.length} encrypted bytes]`);
        // Forward to Bob's session
        const { decodeEnvelope } = await import('../../packages/shade-proto/src/index.js');
        const envelope = decodeEnvelope(bytes);
        const plain = await bob.decrypt('alice', envelope);
        messages.push(plain);
        console.log(`[Bob decrypted]: "${plain}"`);
      },
    },
  });

  // Alice opens a WebSocket and uses ShadeWebSocket wrapper
  const aliceWs = new WebSocket(`ws://localhost:${port}/`);
  await new Promise((r) => aliceWs.addEventListener('open', r));

  const aliceShade = new ShadeWebSocket(aliceWs, alice, 'bob');

  console.log('Alice → Bob: "Hello over WebSocket"');
  await aliceShade.send('Hello over WebSocket');

  console.log('Alice → Bob: "Second message"');
  await aliceShade.send('Second message');

  // Wait for messages to arrive
  await new Promise((r) => setTimeout(r, 200));

  console.log('\nMessages Bob received:', messages);
  console.log('\n=== Done ===');

  aliceShade.close();
  server.stop();
}

main().catch(console.error);