// The browser side of core/crypt.js: the same formats, byte for byte, on Web
// Crypto (so this file also runs under Node's test runner). See core/crypt.js
// for the design; only what differs is explained here.
//
// Web Crypto has no AES-ECB, CMAC or AES-SIV. One-block ECB and CMAC come from
// AES-CBC with a zero IV (Web Crypto adds a PKCS#7 block at the end, which is
// dropped); SIV's counter mode is AES-CTR over the full 128-bit block, as
// Node's aes-ctr does.

const subtle = globalThis.crypto.subtle;
const enc = new TextEncoder();
const dec = new TextDecoder('utf-8', { fatal: true });

export const HEADER_LEN = 24;
export const TAG_LEN = 16;
const MAGIC = [0x50, 0x54, 0x45, 0x31]; // "PTE1"
const LOG2_CHUNK = 20; // 1 MiB
const ZERO = new Uint8Array(16);

const u8 = (x) => (x instanceof Uint8Array ? x : new Uint8Array(x));
export const utf8 = (s) => enc.encode(s);
export function concat(...parts) {
  const out = new Uint8Array(parts.reduce((n, p) => n + p.length, 0));
  let o = 0;
  for (const p of parts) { out.set(p, o); o += p.length; }
  return out;
}

// ---- keys as CryptoKeys, cached per raw key ----------------------------------

const keyCache = new Map();
async function importKey(raw, alg, usages) {
  const id = `${alg}:${b64u(raw)}`;
  let k = keyCache.get(id);
  if (!k) {
    k = await subtle.importKey('raw', raw, alg, false, usages);
    if (keyCache.size > 256) keyCache.clear();
    keyCache.set(id, k);
  }
  return k;
}

// ---- AES-SIV (RFC 5297), AES-CMAC (RFC 4493) ---------------------------------

async function cbc(key, data) {
  const out = u8(await subtle.encrypt({ name: 'AES-CBC', iv: ZERO }, await importKey(key, 'AES-CBC', ['encrypt']), data));
  return out.subarray(0, data.length); // drop the PKCS#7 block
}

function dbl(b) {
  const out = new Uint8Array(16);
  let carry = 0;
  for (let i = 15; i >= 0; i--) {
    out[i] = ((b[i] << 1) | carry) & 0xff;
    carry = b[i] >> 7;
  }
  if (b[0] & 0x80) out[15] ^= 0x87;
  return out;
}

function xor(a, b) {
  const out = new Uint8Array(a.length);
  for (let i = 0; i < a.length; i++) out[i] = a[i] ^ b[i];
  return out;
}

export async function cmac(key, msg) {
  const k1 = dbl(await cbc(key, ZERO));
  const k2 = dbl(k1);
  const n = Math.max(1, Math.ceil(msg.length / 16));
  const complete = msg.length > 0 && msg.length % 16 === 0;
  const last = new Uint8Array(16);
  last.set(msg.subarray((n - 1) * 16));
  if (!complete) last[msg.length - (n - 1) * 16] = 0x80;
  const m = concat(msg.subarray(0, (n - 1) * 16), xor(last, complete ? k1 : k2));
  const out = await cbc(key, m);
  return out.subarray(out.length - 16);
}

async function s2v(key, ads, plain) {
  let d = await cmac(key, ZERO);
  for (const ad of ads) d = xor(dbl(d), await cmac(key, ad));
  let t;
  if (plain.length >= 16) {
    t = Uint8Array.from(plain);
    const tail = t.length - 16;
    for (let i = 0; i < 16; i++) t[tail + i] ^= d[i];
  } else {
    const padded = new Uint8Array(16);
    padded.set(plain);
    padded[plain.length] = 0x80;
    t = xor(dbl(d), padded);
  }
  return cmac(key, t);
}

async function sivCtr(key, v, data) {
  const q = Uint8Array.from(v);
  q[8] &= 0x7f;
  q[12] &= 0x7f;
  return u8(await subtle.encrypt({ name: 'AES-CTR', counter: q, length: 128 }, await importKey(key, 'AES-CTR', ['encrypt']), data));
}

function equal(a, b) {
  if (a.length !== b.length) return false;
  let d = 0;
  for (let i = 0; i < a.length; i++) d |= a[i] ^ b[i];
  return d === 0;
}

// key: 32, 48 or 64 bytes (two AES keys). Returns tag || ciphertext.
export async function sivEncrypt(key, ads, plain) {
  const half = key.length / 2;
  const v = await s2v(key.subarray(0, half), ads, plain);
  return concat(v, await sivCtr(key.subarray(half), v, plain));
}

// The plaintext, or null when the tag does not match.
export async function sivDecrypt(key, ads, data) {
  if (data.length < 16) return null;
  const half = key.length / 2;
  const v = data.subarray(0, 16);
  const plain = await sivCtr(key.subarray(half), v, data.subarray(16));
  return equal(await s2v(key.subarray(0, half), ads, plain), v) ? plain : null;
}

// ---- keys ------------------------------------------------------------------------

export async function hkdf(ikm, salt, info, len) {
  const k = await subtle.importKey('raw', ikm, 'HKDF', false, ['deriveBits']);
  return u8(await subtle.deriveBits({ name: 'HKDF', hash: 'SHA-256', salt, info: utf8(info) }, k, len * 8));
}

export async function driveKeys(master, bucket) {
  if (!(master instanceof Uint8Array) || master.length !== 32) throw new Error('the encryption key must be 32 bytes');
  const okm = await hkdf(master, new Uint8Array(0), `pithos/drive/v1/${bucket}`, 96);
  return { bucket, name: okm.subarray(0, 64), body: okm.subarray(64) };
}

// ---- names -----------------------------------------------------------------------

export function b64u(bytes) {
  let s = '';
  for (let i = 0; i < bytes.length; i += 0x8000) s += String.fromCharCode(...bytes.subarray(i, i + 0x8000));
  return btoa(s).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
}

export function unb64u(s) {
  if (!/^[A-Za-z0-9_-]*$/.test(s)) return null;
  const pad = s.length % 4 === 2 ? '==' : s.length % 4 === 3 ? '=' : s.length % 4 === 1 ? null : '';
  if (pad == null) return null;
  try { return Uint8Array.from(atob(s.replace(/-/g, '+').replace(/_/g, '/') + pad), (c) => c.charCodeAt(0)); } catch { return null; }
}

const segmentAds = (dk, parent) => [utf8(dk.bucket), utf8(parent)];

export async function encryptSegment(dk, parent, seg) {
  return '~' + b64u(await sivEncrypt(dk.name, segmentAds(dk, parent), utf8(seg)));
}

// null when the segment is not one of ours (a plain name).
export async function decryptSegment(dk, parent, seg) {
  if (!seg.startsWith('~') || seg.length < 23) return null;
  const raw = unb64u(seg.slice(1));
  if (!raw || b64u(raw) !== seg.slice(1)) return null;
  const p = await sivDecrypt(dk.name, segmentAds(dk, parent), raw);
  if (!p) return null;
  try { return dec.decode(p); } catch { return null; }
}

export async function encryptPath(dk, plain) {
  if (!plain) return '';
  const folder = plain.endsWith('/');
  const segs = (folder ? plain.slice(0, -1) : plain).split('/');
  const out = [];
  let parent = '';
  for (const s of segs) {
    if (!s) throw new Error('empty folder names are not supported in encrypted drives');
    out.push(await encryptSegment(dk, parent, s));
    parent += s + '/';
  }
  const key = out.join('/') + (folder ? '/' : '');
  if (utf8(key).length > 1024) throw new Error('this name is too long once encrypted (S3 allows 1024 bytes)');
  return key;
}

// { path, encrypted }: encrypted only if every segment was ours.
export async function decryptPath(dk, key) {
  if (!key) return { path: '', encrypted: false };
  const folder = key.endsWith('/');
  const segs = (folder ? key.slice(0, -1) : key).split('/');
  const out = [];
  let parent = '';
  let all = true;
  for (const s of segs) {
    const d = await decryptSegment(dk, parent, s);
    if (d == null) all = false;
    const plain = d ?? s;
    out.push(plain);
    parent += plain + '/';
  }
  return { path: out.join('/') + (folder ? '/' : ''), encrypted: all };
}

// ---- bodies ------------------------------------------------------------------------

export const chunkSize = (log2 = LOG2_CHUNK) => 2 ** log2;

export function cipherSize(plainSize, cs = chunkSize()) {
  return HEADER_LEN + plainSize + Math.max(1, Math.ceil(plainSize / cs)) * TAG_LEN;
}

export function plainSize(cipherSizeBytes, cs = chunkSize()) {
  const c = cipherSizeBytes - HEADER_LEN;
  if (c < TAG_LEN) return null;
  const chunks = Math.ceil(c / (cs + TAG_LEN));
  const p = c - chunks * TAG_LEN;
  return p >= 0 && cipherSize(p, cs) === cipherSizeBytes ? p : null;
}

export function isEncryptedHeader(buf) {
  return buf.length >= HEADER_LEN && MAGIC.every((b, i) => buf[i] === b);
}

export function parseHeader(header) {
  if (!isEncryptedHeader(header)) throw new Error('not an encrypted Pithos file');
  const log2 = header[20];
  if (log2 < 12 || log2 > 24) throw new Error('unsupported chunk size in an encrypted file');
  return { salt: header.subarray(4, 20), cs: 2 ** log2 };
}

export function newHeader() {
  const h = new Uint8Array(HEADER_LEN);
  h.set(MAGIC);
  h.set(globalThis.crypto.getRandomValues(new Uint8Array(16)), 4);
  h[20] = LOG2_CHUNK;
  return h;
}

// What sealing or opening one file's chunks needs.
export async function bodyContext(dk, plainPath, header) {
  const { salt, cs } = parseHeader(header);
  const raw = await hkdf(dk.body, salt, 'pithos/body/v1', 32);
  const key = await subtle.importKey('raw', raw, 'AES-GCM', false, ['encrypt', 'decrypt']);
  return { key, cs, ad: concat(header, utf8(dk.bucket + '\0' + plainPath)) };
}

function nonce(i, last) {
  const n = new Uint8Array(12);
  new DataView(n.buffer).setBigUint64(0, BigInt(i));
  n[8] = last ? 1 : 0;
  return n;
}

export async function sealChunk(ctx, i, last, piece) {
  return u8(await subtle.encrypt({ name: 'AES-GCM', iv: nonce(i, last), additionalData: ctx.ad, tagLength: 128 }, ctx.key, piece));
}

export async function openChunk(ctx, i, last, sealed) {
  try {
    return u8(await subtle.decrypt({ name: 'AES-GCM', iv: nonce(i, last), additionalData: ctx.ad, tagLength: 128 }, ctx.key, sealed));
  } catch {
    throw new Error('this file was changed on the server or belongs to another drive or key');
  }
}

// Encrypts a Blob/File (or bytes) chunk by chunk without holding it whole:
// yields the header, then each sealed chunk. Total = cipherSize(size).
export async function* encryptBlob(dk, plainPath, blob) {
  const src = blob instanceof Uint8Array ? new Blob([blob]) : blob;
  const header = newHeader();
  const ctx = await bodyContext(dk, plainPath, header);
  const chunks = Math.max(1, Math.ceil(src.size / ctx.cs));
  yield header;
  for (let i = 0; i < chunks; i++) {
    const piece = u8(await src.slice(i * ctx.cs, (i + 1) * ctx.cs).arrayBuffer());
    yield await sealChunk(ctx, i, i === chunks - 1, piece);
  }
}

// What to fetch for a plaintext range [start, end] of a cipherLen-byte file.
export function cipherRange(cipherLen, range, cs = chunkSize()) {
  const plain = plainSize(cipherLen, cs);
  if (plain == null) throw new Error('not an encrypted Pithos file (size does not match)');
  const chunks = Math.max(1, Math.ceil(plain / cs));
  let start = 0;
  let end = plain - 1;
  if (range) {
    start = range.start;
    end = Math.min(range.end ?? plain - 1, plain - 1);
    if (start > end || start >= plain) return { plain, unsatisfiable: true };
  }
  const first = plain ? Math.floor(start / cs) : 0;
  const lastChunk = plain ? Math.floor(end / cs) : 0;
  return {
    plain, start, end, chunks, first, lastChunk,
    from: HEADER_LEN + first * (cs + TAG_LEN),
    to: Math.min(cipherLen - 1, HEADER_LEN + (lastChunk + 1) * (cs + TAG_LEN) - 1),
  };
}

// Opens chunks r.first..r.lastChunk read from `readChunk(i)` (which returns the
// sealed bytes of chunk i) and yields only plaintext bytes r.start..r.end.
export async function* decryptChunks(ctx, r, readChunk) {
  for (let i = r.first; i <= r.lastChunk; i++) {
    const plain = await openChunk(ctx, i, i === r.chunks - 1, await readChunk(i));
    const base = i * ctx.cs;
    const a = Math.max(0, r.start - base);
    const b = Math.min(plain.length, r.end - base + 1);
    if (b > a) yield plain.subarray(a, b);
  }
}