// End-to-end encryption for direct messages (notification module phase 2).
//
// Contract (validated against the backend): each user has an ECDH P-256 keypair. The private
// key is wrapped in the browser with an AES-GCM key derived (PBKDF2) from a message passphrase
// the server never sees; only the wrapped blob is uploaded. A conversation key is derived from
// ECDH(my private, partner public) → HKDF → AES-GCM, and each message is AES-GCM encrypted with
// a random IV. The server stores only ciphertext, so not even an admin can read a conversation.
//
// The unlocked private key is persisted per-device as a NON-EXTRACTABLE CryptoKey in IndexedDB,
// so the passphrase is needed once per device, then survives reloads (cleared with browser data).
import type { KeySetup, MyKeyBundle } from "./api";

const subtle = crypto.subtle;
const enc = new TextEncoder();
const dec = new TextDecoder();
const PBKDF2_ITERS = 600_000;
const HKDF_INFO = "siftlode-dm-v1";

const b64 = (buf: ArrayBuffer | Uint8Array): string => {
  const bytes = buf instanceof Uint8Array ? buf : new Uint8Array(buf);
  let s = "";
  for (const b of bytes) s += String.fromCharCode(b);
  return btoa(s);
};
const ub64 = (s: string): Uint8Array => Uint8Array.from(atob(s), (c) => c.charCodeAt(0));
// WebCrypto byte inputs: the DOM lib types want BufferSource backed by a plain ArrayBuffer;
// our Uint8Arrays/ArrayBuffers are functionally that, so coerce at the call sites.
const bsrc = (x: Uint8Array | ArrayBuffer): BufferSource => x as BufferSource;

// --- module state ---
let privKey: CryptoKey | null = null; // current unlocked (non-extractable) private key
const convKeys = new Map<number, CryptoKey>(); // partnerId → derived AES key

// Unlock state is shared app-wide (the Messages page and the floating chat dock both observe
// it), so expose a subscription for useSyncExternalStore.
const unlockSubs = new Set<() => void>();
function emitUnlock(): void {
  for (const cb of unlockSubs) cb();
}
export function subscribeUnlock(cb: () => void): () => void {
  unlockSubs.add(cb);
  return () => unlockSubs.delete(cb);
}

export function isUnlocked(): boolean {
  return !!privKey;
}

// --- IndexedDB (per-device persistence of the unlocked private key) ---
const DB_NAME = "siftlode-e2ee";
const STORE = "privkeys";

function idb(): Promise<IDBDatabase> {
  return new Promise((resolve, reject) => {
    const req = indexedDB.open(DB_NAME, 1);
    req.onupgradeneeded = () => req.result.createObjectStore(STORE);
    req.onsuccess = () => resolve(req.result);
    req.onerror = () => reject(req.error);
  });
}
async function idbGet(userId: number): Promise<CryptoKey | null> {
  const db = await idb();
  return new Promise((resolve) => {
    const r = db.transaction(STORE, "readonly").objectStore(STORE).get(userId);
    r.onsuccess = () => resolve((r.result as CryptoKey) ?? null);
    r.onerror = () => resolve(null);
  });
}
async function idbPut(userId: number, key: CryptoKey): Promise<void> {
  const db = await idb();
  await new Promise<void>((resolve) => {
    const tx = db.transaction(STORE, "readwrite");
    tx.objectStore(STORE).put(key, userId);
    tx.oncomplete = () => resolve();
    tx.onerror = () => resolve();
  });
}
export async function clearDevice(userId: number): Promise<void> {
  const db = await idb();
  await new Promise<void>((resolve) => {
    const tx = db.transaction(STORE, "readwrite");
    tx.objectStore(STORE).delete(userId);
    tx.oncomplete = () => resolve();
    tx.onerror = () => resolve();
  });
}

// --- key derivation ---
async function wrapKeyFrom(passphrase: string, salt: Uint8Array): Promise<CryptoKey> {
  const base = await subtle.importKey("raw", bsrc(enc.encode(passphrase)), "PBKDF2", false, ["deriveKey"]);
  return subtle.deriveKey(
    { name: "PBKDF2", salt: bsrc(salt), iterations: PBKDF2_ITERS, hash: "SHA-256" },
    base,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"]
  );
}
async function importPrivate(pkcs8: ArrayBuffer): Promise<CryptoKey> {
  // Non-extractable: usable for deriveBits but can never be exported again.
  return subtle.importKey("pkcs8", bsrc(pkcs8), { name: "ECDH", namedCurve: "P-256" }, false, ["deriveBits"]);
}

// Try to restore the unlocked key from this device's storage. Returns true if unlocked.
export async function loadFromDevice(userId: number): Promise<boolean> {
  const stored = await idbGet(userId);
  if (stored) {
    privKey = stored;
    emitUnlock();
    return true;
  }
  return false;
}

// Install a freshly-imported private key as the current identity: drop any derived conversation
// keys, persist it to this device, and notify unlock subscribers. Shared by setup() and unlock().
async function installKey(userId: number, key: CryptoKey): Promise<void> {
  privKey = key;
  convKeys.clear();
  await idbPut(userId, key);
  emitUnlock();
}

// First-run setup: generate a keypair, wrap the private key with the passphrase, persist a
// non-extractable copy locally, and return the bundle to upload to the server.
export async function setup(userId: number, passphrase: string): Promise<KeySetup> {
  const kp = await subtle.generateKey({ name: "ECDH", namedCurve: "P-256" }, true, ["deriveBits"]);
  const salt = crypto.getRandomValues(new Uint8Array(16));
  const wrapIv = crypto.getRandomValues(new Uint8Array(12));
  const wrapKey = await wrapKeyFrom(passphrase, salt);
  const pkcs8 = await subtle.exportKey("pkcs8", kp.privateKey);
  const wrapped = await subtle.encrypt({ name: "AES-GCM", iv: bsrc(wrapIv) }, wrapKey, bsrc(pkcs8));
  // key_check is retained for the upload schema but MUST use its own nonce: reusing wrapIv under
  // wrapKey would be GCM nonce-reuse (leaks keystream + enables GHASH forgery). It is no longer a
  // verification oracle either — unlock() authenticates the passphrase via the wrapped-key GCM tag.
  const checkIv = crypto.getRandomValues(new Uint8Array(12));
  const keyCheck = await subtle.encrypt({ name: "AES-GCM", iv: bsrc(checkIv) }, wrapKey, bsrc(enc.encode("ok")));
  const spki = await subtle.exportKey("spki", kp.publicKey);

  await installKey(userId, await importPrivate(pkcs8));
  return {
    public_key: b64(spki),
    wrapped_private_key: b64(wrapped),
    salt: b64(salt),
    wrap_iv: b64(wrapIv),
    key_check: b64(keyCheck),
  };
}

// Unlock on a device from the server blob; throws if the passphrase is wrong.
export async function unlock(userId: number, passphrase: string, bundle: MyKeyBundle): Promise<void> {
  if (!bundle.wrapped_private_key || !bundle.salt || !bundle.wrap_iv || !bundle.key_check) {
    throw new Error("incomplete key bundle");
  }
  const wrapKey = await wrapKeyFrom(passphrase, ub64(bundle.salt));
  // Wrong passphrase → the AES-GCM auth tag fails on this decrypt. (No separate key_check oracle:
  // that was redundant with this tag and formerly shared wrapIv — see setup().)
  const pkcs8 = await subtle.decrypt(
    { name: "AES-GCM", iv: bsrc(ub64(bundle.wrap_iv)) },
    wrapKey,
    bsrc(ub64(bundle.wrapped_private_key))
  );
  await installKey(userId, await importPrivate(pkcs8));
}

// --- per-conversation encryption ---
async function convKeyFor(partnerId: number, partnerPubB64: string): Promise<CryptoKey> {
  const cached = convKeys.get(partnerId);
  if (cached) return cached;
  if (!privKey) throw new Error("locked");
  const pub = await subtle.importKey("spki", bsrc(ub64(partnerPubB64)), { name: "ECDH", namedCurve: "P-256" }, false, []);
  const bits = await subtle.deriveBits({ name: "ECDH", public: pub }, privKey, 256);
  const hk = await subtle.importKey("raw", bits, "HKDF", false, ["deriveKey"]);
  const key = await subtle.deriveKey(
    { name: "HKDF", hash: "SHA-256", salt: bsrc(new Uint8Array(0)), info: bsrc(enc.encode(HKDF_INFO)) },
    hk,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"]
  );
  convKeys.set(partnerId, key);
  return key;
}

export async function encryptFor(
  partnerId: number,
  partnerPubB64: string,
  plaintext: string
): Promise<{ ciphertext: string; iv: string }> {
  const key = await convKeyFor(partnerId, partnerPubB64);
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const ct = await subtle.encrypt({ name: "AES-GCM", iv: bsrc(iv) }, key, bsrc(enc.encode(plaintext)));
  return { ciphertext: b64(ct), iv: b64(iv) };
}

export async function decryptFrom(
  partnerId: number,
  partnerPubB64: string,
  ciphertext: string,
  iv: string
): Promise<string> {
  const key = await convKeyFor(partnerId, partnerPubB64);
  const pt = await subtle.decrypt({ name: "AES-GCM", iv: bsrc(ub64(iv)) }, key, bsrc(ub64(ciphertext)));
  return dec.decode(pt);
}