/**
 * Shade E2EE Demo — shows exactly how a frontend and backend use Shade.
 *
 * Run: bun demo.ts
 */
import { SubtleCryptoProvider, MemoryStorage } from './packages/shade-crypto-web/src/index.js';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  processPreKeyMessage,
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from './packages/shade-core/src/index.js';

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

function log(who: string, msg: string) {
  const color = who === 'FRONTEND' ? '\x1b[36m' : who === 'BACKEND' ? '\x1b[33m' : '\x1b[90m';
  console.log(`${color}[${who}]\x1b[0m ${msg}`);
}

async function main() {
  console.log('\n\x1b[1m=== SHADE E2EE DEMO ===\x1b[0m\n');

  // ─── BACKEND: One-time setup (happens once on deploy) ──────
  log('BACKEND', 'Generating identity keys...');
  const backendStorage = new MemoryStorage();
  const backendIdentity = await generateIdentityKeyPair(crypto);
  await backendStorage.saveIdentityKeyPair(backendIdentity);
  await backendStorage.saveLocalRegistrationId(1);

  const signedPreKey = await generateSignedPreKey(crypto, backendIdentity, 1);
  await backendStorage.saveSignedPreKey(signedPreKey);

  const oneTimePreKeys = await generateOneTimePreKeys(crypto, 100, 10);
  for (const otpk of oneTimePreKeys) await backendStorage.saveOneTimePreKey(otpk);

  log('BACKEND', `Published prekey bundle (10 one-time prekeys ready)`);

  // This bundle would be stored on the Shade Prekey Server
  const bundle = createPreKeyBundle(1, backendIdentity, signedPreKey, oneTimePreKeys[0]);

  // ─── FRONTEND: Connect to backend ─────────────────────────
  console.log('\n\x1b[90m--- Frontend wants to talk to Backend ---\x1b[0m\n');

  log('FRONTEND', 'Generating identity keys...');
  const frontendStorage = new MemoryStorage();
  const frontendIdentity = await generateIdentityKeyPair(crypto);
  await frontendStorage.saveIdentityKeyPair(frontendIdentity);

  log('FRONTEND', 'Fetching backend prekey bundle from Shade server...');
  // In reality: const bundle = await fetch('https://shade-server/v1/keys/bundle/backend')

  log('FRONTEND', 'Processing bundle (X3DH key agreement)...');
  const x3dhResult = await processPreKeyBundle(crypto, frontendStorage, bundle);
  log('FRONTEND', `Shared secret derived (${x3dhResult.rootKey.length * 8}-bit key)`);

  // Frontend creates its ratchet session
  const frontendSession = await initSenderSession(
    crypto,
    x3dhResult.rootKey,
    x3dhResult.remoteIdentityKey,
    x3dhResult.remoteSignedPreKey,
  );

  // ─── FRONTEND: Send first encrypted message ───────────────
  console.log('\n\x1b[90m--- Encrypted conversation begins ---\x1b[0m\n');

  const firstPlaintext = 'Hello Backend! This message is end-to-end encrypted.';
  log('FRONTEND', `Encrypting: "${firstPlaintext}"`);
  const firstEncrypted = await ratchetEncrypt(crypto, frontendSession, enc.encode(firstPlaintext));

  // Bundle it as a PreKeyMessage (first message includes X3DH info)
  const preKeyMessage = {
    registrationId: 1,
    preKeyId: x3dhResult.preKeyId,
    signedPreKeyId: x3dhResult.signedPreKeyId,
    ephemeralKey: x3dhResult.ephemeralPublicKey,
    identityDHKey: frontendIdentity.dhPublicKey,
    message: firstEncrypted,
  };

  log('FRONTEND', `Sending PreKeyMessage (${firstEncrypted.ciphertext.length} bytes ciphertext)`);
  // In reality: await fetch('https://backend/api/messages', { body: preKeyMessage })

  // ─── BACKEND: Receive and decrypt ─────────────────────────
  log('BACKEND', 'Received PreKeyMessage, processing X3DH...');
  const backendX3dh = await processPreKeyMessage(crypto, backendStorage, preKeyMessage);
  log('BACKEND', `Shared secret derived (matches frontend: ${arrEq(backendX3dh.rootKey, x3dhResult.rootKey)})`);

  const backendSession = initReceiverSession(
    backendX3dh.rootKey,
    backendX3dh.remoteIdentityKey,
    signedPreKey.keyPair,
  );

  const decrypted1 = await ratchetDecrypt(crypto, backendSession, backendX3dh.initialMessage);
  log('BACKEND', `Decrypted: "${dec.decode(decrypted1)}"`);

  // ─── BACKEND: Reply (triggers DH ratchet) ──────────────────
  const replyText = 'Got it! Every message from now on uses a unique key.';
  log('BACKEND', `Encrypting reply: "${replyText}"`);
  const replyEncrypted = await ratchetEncrypt(crypto, backendSession, enc.encode(replyText));
  log('BACKEND', `Sending RatchetMessage (new DH key generated)`);

  // ─── FRONTEND: Receive reply ──────────────────────────────
  const decrypted2 = await ratchetDecrypt(crypto, frontendSession, replyEncrypted);
  log('FRONTEND', `Decrypted: "${dec.decode(decrypted2)}"`);

  // ─── A few more messages ──────────────────────────────────
  console.log('\n\x1b[90m--- Continued conversation (each message = new key) ---\x1b[0m\n');

  const conversation = [
    { from: 'FRONTEND', text: 'If someone steals this key, they cannot read previous messages.' },
    { from: 'BACKEND', text: 'Correct — that is forward secrecy via the Double Ratchet.' },
    { from: 'FRONTEND', text: 'And if they steal my device, future messages re-secure automatically?' },
    { from: 'BACKEND', text: 'Yes — the next DH ratchet step generates fresh keys. Post-compromise recovery.' },
  ];

  for (const { from, text } of conversation) {
    const sender = from === 'FRONTEND' ? frontendSession : backendSession;
    const receiver = from === 'FRONTEND' ? backendSession : frontendSession;
    const receiverName = from === 'FRONTEND' ? 'BACKEND' : 'FRONTEND';

    const encrypted = await ratchetEncrypt(crypto, sender, enc.encode(text));
    const plain = dec.decode(await ratchetDecrypt(crypto, receiver, encrypted));

    log(from, `>> "${text}"`);
    log(receiverName, `<< decrypted OK`);
  }

  // ─── Summary ──────────────────────────────────────────────
  console.log('\n\x1b[1m=== SUMMARY ===\x1b[0m');
  console.log(`  Messages exchanged:  6`);
  console.log(`  DH ratchet steps:    ${4} (new keys generated)`);
  console.log(`  One-time prekeys remaining: ${await backendStorage.getOneTimePreKeyCount()}`);
  console.log(`  Forward secrecy:     YES`);
  console.log(`  Post-compromise:     YES`);
  console.log(`  Plaintext leaked:    ZERO\n`);
}

function arrEq(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
  return true;
}

main().catch(console.error);