// High-level E2E crypto built on the vendored Signal Protocol library.
// This module is the ONLY place message plaintext exists in the app besides the UI.

import {
  KeyHelper,
  SignalProtocolAddress,
  SessionBuilder,
  SessionCipher,
} from "./vendor/libsignal.js";
import { SignalStore } from "./store.js";

export const store = new SignalStore();

// ---- encoding helpers ----
export function abToB64(ab) {
  const bytes = new Uint8Array(ab);
  let bin = "";
  for (let i = 0; i < bytes.length; i++) bin += String.fromCharCode(bytes[i]);
  return btoa(bin);
}
export function b64ToAb(b64) {
  const bin = atob(b64);
  const bytes = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
  return bytes.buffer;
}
// Signal message bodies are binary strings; base64 them for JSON transport.
function binStrToB64(s) {
  return btoa(s);
}
function b64ToBinStr(b64) {
  return atob(b64);
}

const NUM_ONE_TIME_PREKEYS = 100;
const SIGNED_PREKEY_ID = 1;

// Generate a fresh identity + prekeys and persist private halves locally.
// Returns the PUBLIC registration payload to send to the server.
export async function generateRegistration() {
  const identityKeyPair = await KeyHelper.generateIdentityKeyPair();
  const registrationId = KeyHelper.generateRegistrationId();
  await store.setIdentity(identityKeyPair, registrationId);

  const signedPreKey = await KeyHelper.generateSignedPreKey(identityKeyPair, SIGNED_PREKEY_ID);
  await store.storeSignedPreKey(SIGNED_PREKEY_ID, signedPreKey.keyPair);

  const oneTimePreKeys = [];
  const startId = 1;
  for (let i = 0; i < NUM_ONE_TIME_PREKEYS; i++) {
    const id = startId + i;
    const pk = await KeyHelper.generatePreKey(id);
    await store.storePreKey(id, pk.keyPair);
    oneTimePreKeys.push({ id, pub: abToB64(pk.keyPair.pubKey) });
  }

  return {
    registrationId,
    identityKey: abToB64(identityKeyPair.pubKey),
    signedPreKey: {
      id: SIGNED_PREKEY_ID,
      pub: abToB64(signedPreKey.keyPair.pubKey),
      sig: abToB64(signedPreKey.signature),
    },
    oneTimePreKeys,
  };
}

// Generate a batch of fresh one-time prekeys to replenish the server's supply.
export async function generateMoreOneTimePreKeys(count = 100) {
  // Continue numbering past whatever we've used; store a counter in meta-less way:
  // derive from a random high base to avoid collisions (server ignores dupes).
  const base = 1000 + Math.floor(Math.random() * 1_000_000);
  const oneTimePreKeys = [];
  for (let i = 0; i < count; i++) {
    const id = base + i;
    const pk = await KeyHelper.generatePreKey(id);
    await store.storePreKey(id, pk.keyPair);
    oneTimePreKeys.push({ id, pub: abToB64(pk.keyPair.pubKey) });
  }
  return oneTimePreKeys;
}

// ---- multi-device key sync (recovery backup + QR linking) ----

// Serialize the on-device identity so it can be restored on another device.
// This is the SECRET material — it is only ever emitted encrypted.
export async function exportIdentityBlob(account) {
  const idkp = await store.getIdentityKeyPair();
  const registrationId = await store.getLocalRegistrationId();
  if (!idkp) throw new Error("no identity on this device");
  return {
    v: 1,
    userId: account.userId,
    deviceId: account.deviceId,
    username: account.username,
    displayName: account.displayName,
    registrationId,
    identityKey: { pub: abToB64(idkp.pubKey), priv: abToB64(idkp.privKey) },
  };
}

// Restore an identity blob into the (already namespaced) store, then generate a
// fresh set of prekeys to publish. Caller must have called useAccount() first.
// Old sessions/ratchets don't survive — new conversations re-handshake cleanly.
export async function importIdentityBlob(blob) {
  const identityKeyPair = {
    pubKey: b64ToAb(blob.identityKey.pub),
    privKey: b64ToAb(blob.identityKey.priv),
  };
  await store.setIdentity(identityKeyPair, blob.registrationId);

  const signedPreKey = await KeyHelper.generateSignedPreKey(identityKeyPair, SIGNED_PREKEY_ID);
  await store.storeSignedPreKey(SIGNED_PREKEY_ID, signedPreKey.keyPair);
  const oneTimePreKeys = [];
  for (let i = 1; i <= NUM_ONE_TIME_PREKEYS; i++) {
    const pk = await KeyHelper.generatePreKey(i);
    await store.storePreKey(i, pk.keyPair);
    oneTimePreKeys.push({ id: i, pub: abToB64(pk.keyPair.pubKey) });
  }
  return {
    signedPreKey: {
      id: SIGNED_PREKEY_ID,
      pub: abToB64(signedPreKey.keyPair.pubKey),
      sig: abToB64(signedPreKey.signature),
    },
    oneTimePreKeys,
  };
}

// A human-writable recovery code: 24 Crockford-base32 chars in 6 groups of 4.
// ~120 bits of entropy — infeasible to brute force even with the blob in hand.
export function generateRecoveryCode() {
  const alphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"; // no I L O U
  const bytes = crypto.getRandomValues(new Uint8Array(15));
  let bits = 0, value = 0, out = "";
  for (const b of bytes) {
    value = (value << 8) | b;
    bits += 8;
    while (bits >= 5) {
      out += alphabet[(value >>> (bits - 5)) & 31];
      bits -= 5;
    }
  }
  return out.match(/.{1,4}/g).join("-"); // e.g. XK4P-92MB-...
}

function normalizeCode(code) {
  return code.toUpperCase().replace(/[^0-9A-Z]/g, "");
}

async function deriveCodeKey(code, salt) {
  const base = await crypto.subtle.importKey(
    "raw",
    new TextEncoder().encode(normalizeCode(code)),
    "PBKDF2",
    false,
    ["deriveKey"]
  );
  // Client-side (browser) has no 100k PBKDF2 cap, unlike the Worker.
  return crypto.subtle.deriveKey(
    { name: "PBKDF2", hash: "SHA-256", salt, iterations: 250000 },
    base,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"]
  );
}

// Encrypt an object under a recovery code -> {salt, iv, ct} (all base64).
export async function encryptWithCode(obj, code) {
  const salt = crypto.getRandomValues(new Uint8Array(16));
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const key = await deriveCodeKey(code, salt);
  const ct = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv },
    key,
    new TextEncoder().encode(JSON.stringify(obj))
  );
  return { salt: abToB64(salt.buffer), iv: abToB64(iv.buffer), ct: abToB64(ct) };
}

export async function decryptWithCode(payload, code) {
  const key = await deriveCodeKey(code, new Uint8Array(b64ToAb(payload.salt)));
  const pt = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv: new Uint8Array(b64ToAb(payload.iv)) },
    key,
    b64ToAb(payload.ct)
  );
  return JSON.parse(new TextDecoder().decode(new Uint8Array(pt)));
}

// ---- QR linking: ECDH to the new device's ephemeral public key ----

// New device generates an ephemeral P-256 keypair; the raw public key goes in
// the QR. The private key never leaves this device.
export async function generateEphemeralKeyPair() {
  const kp = await crypto.subtle.generateKey({ name: "ECDH", namedCurve: "P-256" }, true, [
    "deriveKey",
  ]);
  const rawPub = await crypto.subtle.exportKey("raw", kp.publicKey);
  return { keyPair: kp, pubB64: abToB64(rawPub) };
}

// Authorizing device: encrypt a blob to the new device's ephemeral public key.
// Uses an ephemeral-static ECDH (fresh sender key each time) -> AES-GCM.
export async function encryptToEphemeral(obj, peerPubB64) {
  const peerPub = await crypto.subtle.importKey(
    "raw",
    b64ToAb(peerPubB64),
    { name: "ECDH", namedCurve: "P-256" },
    false,
    []
  );
  const my = await crypto.subtle.generateKey({ name: "ECDH", namedCurve: "P-256" }, true, [
    "deriveKey",
  ]);
  const key = await crypto.subtle.deriveKey(
    { name: "ECDH", public: peerPub },
    my.privateKey,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt"]
  );
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const ct = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv },
    key,
    new TextEncoder().encode(JSON.stringify(obj))
  );
  const myRawPub = await crypto.subtle.exportKey("raw", my.publicKey);
  return { epk: abToB64(myRawPub), iv: abToB64(iv.buffer), ct: abToB64(ct) };
}

// New device: decrypt the relayed blob using its ephemeral private key.
export async function decryptFromEphemeral(payload, ephemeralKeyPair) {
  const senderPub = await crypto.subtle.importKey(
    "raw",
    b64ToAb(payload.epk),
    { name: "ECDH", namedCurve: "P-256" },
    false,
    []
  );
  const key = await crypto.subtle.deriveKey(
    { name: "ECDH", public: senderPub },
    ephemeralKeyPair.privateKey,
    { name: "AES-GCM", length: 256 },
    false,
    ["decrypt"]
  );
  const pt = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv: new Uint8Array(b64ToAb(payload.iv)) },
    key,
    b64ToAb(payload.ct)
  );
  return JSON.parse(new TextDecoder().decode(new Uint8Array(pt)));
}

// Establish outbound sessions with every device in a fetched prekey bundle.
export async function ensureSessions(bundlesResponse) {
  for (const b of bundlesResponse.bundles) {
    const address = new SignalProtocolAddress(bundlesResponse.userId, deviceIndex(b.deviceId));
    const existing = await store.loadSession(address.toString());
    if (existing) continue;
    const builder = new SessionBuilder(store, address);
    await builder.processPreKey({
      registrationId: b.registrationId,
      identityKey: b64ToAb(b.identityKey),
      signedPreKey: {
        keyId: b.signedPreKey.id,
        publicKey: b64ToAb(b.signedPreKey.pub),
        signature: b64ToAb(b.signedPreKey.sig),
      },
      preKey: b.oneTimePreKey
        ? { keyId: b.oneTimePreKey.id, publicKey: b64ToAb(b.oneTimePreKey.pub) }
        : undefined,
    });
  }
}

// Signal addresses need an integer deviceId. Map our uuid deviceIds to a stable
// small int by keeping a per-user registry in memory (and the session key string
// already encodes the full identity, so uniqueness within a user is all we need).
const deviceIndexMap = new Map();
let deviceCounter = 1;
function deviceIndex(deviceId) {
  if (!deviceIndexMap.has(deviceId)) deviceIndexMap.set(deviceId, deviceCounter++);
  return deviceIndexMap.get(deviceId);
}

// Encrypt a plaintext string for every device of a recipient. Returns the array
// of per-device messages to POST to /api/messages/send.
export async function encryptForUser(userId, deviceIds, plaintext) {
  const buffer = new TextEncoder().encode(plaintext).buffer;
  const messages = [];
  for (const deviceId of deviceIds) {
    const address = new SignalProtocolAddress(userId, deviceIndex(deviceId));
    const cipher = new SessionCipher(store, address);
    const ct = await cipher.encrypt(buffer);
    messages.push({ toDeviceId: deviceId, type: ct.type, body: binStrToB64(ct.body) });
  }
  return messages;
}

// Decrypt an inbound envelope. Handles both prekey (type 3) and whisper (type 1).
export async function decryptEnvelope(envelope, senderUserId) {
  const address = new SignalProtocolAddress(senderUserId, deviceIndex(envelope.fromDeviceId));
  const cipher = new SessionCipher(store, address);
  const body = b64ToBinStr(envelope.body);
  let plaintextBuf;
  if (envelope.type === 3) {
    plaintextBuf = await cipher.decryptPreKeyWhisperMessage(body, "binary");
  } else {
    plaintextBuf = await cipher.decryptWhisperMessage(body, "binary");
  }
  return new TextDecoder().decode(new Uint8Array(plaintextBuf));
}

// ---- media encryption (AES-GCM, key travels inside the E2E message) ----
export async function encryptMedia(fileArrayBuffer) {
  const key = await crypto.subtle.generateKey({ name: "AES-GCM", length: 256 }, true, [
    "encrypt",
    "decrypt",
  ]);
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const ciphertext = await crypto.subtle.encrypt({ name: "AES-GCM", iv }, key, fileArrayBuffer);
  const rawKey = await crypto.subtle.exportKey("raw", key);
  return {
    ciphertext, // ArrayBuffer to upload to R2
    keyB64: abToB64(rawKey),
    ivB64: abToB64(iv.buffer),
  };
}

export async function decryptMedia(ciphertext, keyB64, ivB64) {
  const key = await crypto.subtle.importKey(
    "raw",
    b64ToAb(keyB64),
    { name: "AES-GCM" },
    false,
    ["decrypt"]
  );
  const iv = new Uint8Array(b64ToAb(ivB64));
  return crypto.subtle.decrypt({ name: "AES-GCM", iv }, key, ciphertext);
}