// 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;
}

// 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);
}