import { b64ToBuf, bufToB64, toArrayBuffer } from "./storage";
import {
  E2eeEnvelope,
  parseEnvelope,
  serializeEnvelope,
} from "./envelope";

const ECDH_PARAMS: EcKeyGenParams = { name: "ECDH", namedCurve: "P-256" };
const AES_BITS = 256;

export async function generateIdentityKeyPair(): Promise<CryptoKeyPair> {
  return crypto.subtle.generateKey(ECDH_PARAMS, true, ["deriveBits"]);
}

export async function exportPublicKeySpkiB64(publicKey: CryptoKey): Promise<string> {
  const spki = await crypto.subtle.exportKey("spki", publicKey);
  return bufToB64(spki);
}

export async function importPublicKeySpkiB64(b64: string): Promise<CryptoKey> {
  return crypto.subtle.importKey(
    "spki",
    b64ToBuf(b64),
    ECDH_PARAMS,
    true,
    []
  );
}

export async function exportPrivateKeyJwk(privateKey: CryptoKey): Promise<JsonWebKey> {
  return crypto.subtle.exportKey("jwk", privateKey);
}

export async function importPrivateKeyJwk(jwk: JsonWebKey): Promise<CryptoKey> {
  return crypto.subtle.importKey("jwk", jwk, ECDH_PARAMS, true, ["deriveBits"]);
}

async function deriveAesKey(
  privateKey: CryptoKey,
  publicKey: CryptoKey,
  infoLabel: string
): Promise<CryptoKey> {
  const bits = await crypto.subtle.deriveBits(
    { name: "ECDH", public: publicKey },
    privateKey,
    256
  );

  const hkdfKey = await crypto.subtle.importKey("raw", bits, "HKDF", false, [
    "deriveKey",
  ]);

  const info = new TextEncoder().encode(infoLabel);
  const salt = new Uint8Array(16);

  return crypto.subtle.deriveKey(
    { name: "HKDF", hash: "SHA-256", salt, info },
    hkdfKey,
    { name: "AES-GCM", length: AES_BITS },
    false,
    ["encrypt", "decrypt"]
  );
}

export async function deriveDmAesKey(
  myPrivateKey: CryptoKey,
  peerPublicKey: CryptoKey,
  userIdA: string,
  userIdB: string
): Promise<CryptoKey> {
  const [a, b] = [String(userIdA), String(userIdB)].sort();
  return deriveAesKey(myPrivateKey, peerPublicKey, `modstagram-dm-v1:${a}:${b}`);
}

export async function encryptWithAesGcm(
  key: CryptoKey,
  plaintext: string
): Promise<E2eeEnvelope> {
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const encoded = new TextEncoder().encode(plaintext);
  const ct = await crypto.subtle.encrypt({ name: "AES-GCM", iv }, key, encoded);
  return {
    v: 1,
    e2ee: true,
    iv: bufToB64(iv),
    alg: "AES-GCM",
    ciphertext: bufToB64(ct),
  };
}

export async function decryptWithAesGcm(
  key: CryptoKey,
  envelope: E2eeEnvelope
): Promise<string> {
  const iv = b64ToBuf(envelope.iv);
  const ct = b64ToBuf(envelope.ciphertext);
  const pt = await crypto.subtle.decrypt({ name: "AES-GCM", iv }, key, ct);
  return new TextDecoder().decode(pt);
}

export async function encryptText(
  key: CryptoKey,
  plaintext: string
): Promise<string> {
  return serializeEnvelope(await encryptWithAesGcm(key, plaintext));
}

export async function decryptText(
  key: CryptoKey,
  content: string
): Promise<string | null> {
  const env = parseEnvelope(content);
  if (!env) return null;
  try {
    return await decryptWithAesGcm(key, env);
  } catch {
    return null;
  }
}

export async function generateRoomKeyRaw(): Promise<Uint8Array> {
  return crypto.getRandomValues(new Uint8Array(32));
}

export async function importRoomAesKey(raw: Uint8Array): Promise<CryptoKey> {
  return crypto.subtle.importKey(
    "raw",
    toArrayBuffer(raw),
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"]
  );
}

export async function wrapRoomKeyForRecipient(
  roomKeyRaw: Uint8Array,
  recipientPublicKey: CryptoKey
): Promise<string> {
  const ephemeral = await crypto.subtle.generateKey(ECDH_PARAMS, true, [
    "deriveBits",
  ]);
  const aesKey = await deriveAesKey(
    ephemeral.privateKey,
    recipientPublicKey,
    "modstagram-room-wrap-v1"
  );
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const ct = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv },
    aesKey,
    toArrayBuffer(roomKeyRaw)
  );
  const epk = await exportPublicKeySpkiB64(ephemeral.publicKey);
  return bufToB64(
    new TextEncoder().encode(
      JSON.stringify({
        epk,
        iv: bufToB64(iv),
        ct: bufToB64(ct),
      })
    )
  );
}

export async function unwrapRoomKey(
  wrappedB64: string,
  myPrivateKey: CryptoKey
): Promise<Uint8Array> {
  const parsed = JSON.parse(new TextDecoder().decode(b64ToBuf(wrappedB64))) as {
    epk: string;
    iv: string;
    ct: string;
  };
  const epk = await importPublicKeySpkiB64(parsed.epk);
  const aesKey = await deriveAesKey(
    myPrivateKey,
    epk,
    "modstagram-room-wrap-v1"
  );
  const pt = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv: b64ToBuf(parsed.iv) },
    aesKey,
    b64ToBuf(parsed.ct)
  );
  return new Uint8Array(pt);
}
