// AES-256-CTR range decryption for Matrix attachments (the only cipher Matrix
// has, per vendor/macltok's AttachmentEncryption). WebCrypto's AES-CTR handles
// a partial final block within one call but not a resumed one, so a block that
// isn't a multiple of 16 bytes has to wait for more ciphertext before it can be
// decrypted; see RangeDecryptor.push.

const BLOCK_SIZE = 16;

function decodeBase64(value: string): Uint8Array<ArrayBuffer> {
  const normalised = value.replace(/-/g, '+').replace(/_/g, '/');
  const padded = normalised + '='.repeat((4 - (normalised.length % 4)) % 4);
  const binary = atob(padded);
  const bytes = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i);
  return bytes;
}

export async function importAttachmentKey(k: string): Promise<CryptoKey> {
  return crypto.subtle.importKey('raw', decodeBase64(k), { name: 'AES-CTR' }, false, ['decrypt']);
}

export function decodeAttachmentIv(iv: string): Uint8Array<ArrayBuffer> {
  const bytes = decodeBase64(iv);
  if (bytes.length >= BLOCK_SIZE) return bytes.subarray(0, BLOCK_SIZE);
  const padded = new Uint8Array(BLOCK_SIZE);
  padded.set(bytes);
  return padded;
}

/** iv + blocks, added as a big-endian 64-bit value into the low 8 bytes only, per the Matrix attachment convention. */
export function addBlockOffset(iv: Uint8Array, blocks: number): Uint8Array<ArrayBuffer> {
  const counter = iv.slice();
  const view = new DataView(counter.buffer, counter.byteOffset, counter.byteLength);
  const low = view.getBigUint64(8);
  view.setBigUint64(8, (low + BigInt(blocks)) & 0xffffffffffffffffn);
  return counter;
}

function concat(a: Uint8Array<ArrayBuffer>, b: Uint8Array<ArrayBuffer>): Uint8Array<ArrayBuffer> {
  if (a.length === 0) return b;
  if (b.length === 0) return a;
  const out = new Uint8Array(a.length + b.length);
  out.set(a, 0);
  out.set(b, a.length);
  return out;
}

/**
 * Decrypts a ciphertext range that starts `startBlock` 16-byte blocks into the
 * attachment. Feed it ciphertext in arrival order via `push`, then `finish` once
 * the source is exhausted to flush the trailing partial block.
 */
export class RangeDecryptor {
  private buffered: Uint8Array<ArrayBuffer> = new Uint8Array(0);
  private blocksConsumed = 0;

  constructor(
    private readonly key: CryptoKey,
    private readonly iv: Uint8Array,
    private readonly startBlock: number
  ) {}

  async push(chunk: Uint8Array<ArrayBuffer>): Promise<Uint8Array> {
    const data = concat(this.buffered, chunk);
    const wholeLen = data.length - (data.length % BLOCK_SIZE);
    if (wholeLen === 0) {
      this.buffered = data;
      return new Uint8Array(0);
    }
    this.buffered = data.subarray(wholeLen);
    return this.decrypt(data.subarray(0, wholeLen));
  }

  async finish(): Promise<Uint8Array> {
    if (this.buffered.length === 0) return new Uint8Array(0);
    const plain = await this.decrypt(this.buffered);
    this.buffered = new Uint8Array(0);
    return plain;
  }

  private async decrypt(ciphertext: Uint8Array<ArrayBuffer>): Promise<Uint8Array> {
    const counter = addBlockOffset(this.iv, this.startBlock + this.blocksConsumed);
    this.blocksConsumed += Math.ceil(ciphertext.length / BLOCK_SIZE);
    const plain = await crypto.subtle.decrypt({ name: 'AES-CTR', counter, length: 64 }, this.key, ciphertext);
    return new Uint8Array(plain);
  }
}