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

const crypto = new SubtleCryptoProvider();

async function main() {
  console.log('=== Shade Prekey Server Demo ===\n');

  // Start the prekey server
  const store = new MemoryPrekeyStore();
  const server = createPrekeyServer({ crypto, store });
  const port = 19850;
  const handle = Bun.serve({ port, fetch: server.fetch });
  console.log(`Prekey server listening on http://localhost:${port}\n`);

  try {
    // ─── Bob registers ────────────────────────────────────
    const bobMgr = new ShadeSessionManager(crypto, new MemoryStorage());
    await bobMgr.initialize();

    const bobIdentity = await new MemoryStorage().getIdentityKeyPair();
    // Note: in real usage, get the storage instance you passed in. Simplified here:
    const bobStorageRef = new MemoryStorage();
    const bobMgr2 = new ShadeSessionManager(crypto, bobStorageRef);
    await bobMgr2.initialize();
    const bobKp = await bobStorageRef.getIdentityKeyPair();

    const bobTransport = new ShadeFetchTransport({
      baseUrl: `http://localhost:${port}`,
      crypto,
      signingPrivateKey: bobKp!.signingPrivateKey,
    });

    const bobOTPKs = await bobMgr2.generateOneTimePreKeys(10);
    const bobBundle = await bobMgr2.createPreKeyBundle();
    await bobTransport.register('bob', bobMgr2.getPublicIdentity(), bobBundle.signedPreKey, bobOTPKs);
    console.log('Bob registered with prekey server (10 one-time prekeys uploaded)');
    console.log('Bob fingerprint:', await bobMgr2.getIdentityFingerprint());
    console.log();

    // ─── Alice fetches Bob's bundle anonymously ───────────
    const aliceMgr = new ShadeSessionManager(crypto, new MemoryStorage());
    await aliceMgr.initialize();
    console.log('Alice fingerprint:', await aliceMgr.getIdentityFingerprint());

    // Alice doesn't need a signing key to fetch
    const aliceTransport = new ShadeFetchTransport({
      baseUrl: `http://localhost:${port}`,
      crypto,
    });

    const fetchedBundle = await aliceTransport.fetchBundle('bob');
    console.log('\nAlice fetched Bob\'s bundle');
    console.log('Remaining one-time prekeys:', await aliceTransport.getKeyCount('bob'));

    await aliceMgr.initSessionFromBundle('bob', fetchedBundle);

    // ─── Encrypted conversation ───────────────────────────
    console.log('\n=== Encrypted conversation ===');
    const env1 = await aliceMgr.encrypt('bob', 'Hello via prekey server!');
    console.log('→ Alice sends encrypted message');
    const plain1 = await bobMgr2.decrypt('alice', env1);
    console.log(`← Bob decrypts: "${plain1}"`);

    const env2 = await bobMgr2.encrypt('alice', 'Got your message!');
    console.log('\n→ Bob replies (DH ratchet step)');
    const plain2 = await aliceMgr.decrypt('bob', env2);
    console.log(`← Alice decrypts: "${plain2}"`);

    console.log('\n=== Done ===');
  } finally {
    handle.stop();
  }
}

main().catch(console.error);