import assert from 'node:assert/strict';
import { test } from 'node:test';
import { addBlockOffset, importAttachmentKey, RangeDecryptor } from '../src/lib/media-decrypt.ts';
import { attachmentContentType, imageContentType } from '../src/lib/media-sniff.ts';

function randomBytes(length: number): Uint8Array<ArrayBuffer> {
  const bytes = new Uint8Array(length);
  crypto.getRandomValues(bytes);
  return bytes;
}

function textBytes(text: string): Uint8Array<ArrayBuffer> {
  return new TextEncoder().encode(text);
}

function concatBytes(parts: Uint8Array[]): Uint8Array<ArrayBuffer> {
  const total = parts.reduce((sum, part) => sum + part.length, 0);
  const out = new Uint8Array(total);
  let offset = 0;
  for (const part of parts) {
    out.set(part, offset);
    offset += part.length;
  }
  return out;
}

function base64url(bytes: Uint8Array): string {
  let binary = '';
  for (const byte of bytes) binary += String.fromCharCode(byte);
  return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
}

async function encryptWhole(
  rawKey: Uint8Array<ArrayBuffer>,
  iv: Uint8Array,
  plaintext: Uint8Array<ArrayBuffer>
): Promise<Uint8Array<ArrayBuffer>> {
  const encryptKey = await crypto.subtle.importKey('raw', rawKey, { name: 'AES-CTR' }, false, ['encrypt']);
  const counter = addBlockOffset(iv, 0);
  const cipher = await crypto.subtle.encrypt({ name: 'AES-CTR', counter, length: 64 }, encryptKey, plaintext);
  return new Uint8Array(cipher);
}

async function decryptRange(
  key: CryptoKey,
  iv: Uint8Array,
  ciphertext: Uint8Array<ArrayBuffer>,
  start: number,
  end: number,
  chunkSize: number
): Promise<Uint8Array> {
  const alignedStart = start - (start % 16);
  const slice = ciphertext.subarray(alignedStart, end + 1);
  const decryptor = new RangeDecryptor(key, iv, alignedStart / 16);
  const parts: Uint8Array[] = [];
  for (let offset = 0; offset < slice.length; offset += chunkSize) {
    const chunk = slice.subarray(offset, Math.min(offset + chunkSize, slice.length));
    const plain = await decryptor.push(chunk);
    if (plain.length > 0) parts.push(plain);
  }
  const last = await decryptor.finish();
  if (last.length > 0) parts.push(last);
  return concatBytes(parts).subarray(start - alignedStart);
}

async function setup(plaintextLength: number) {
  const rawKey = randomBytes(32);
  const iv = randomBytes(16);
  const key = await importAttachmentKey(base64url(rawKey));
  const plaintext = randomBytes(plaintextLength);
  const ciphertext = await encryptWhole(rawKey, iv, plaintext);
  return { key, iv, plaintext, ciphertext };
}

test('range decryptor: whole read matches the plaintext', async () => {
  const { key, iv, plaintext, ciphertext } = await setup(1000);
  const decryptor = new RangeDecryptor(key, iv, 0);
  const decrypted = concatBytes([await decryptor.push(ciphertext), await decryptor.finish()]);
  assert.deepEqual(decrypted, plaintext);
});

test('range decryptor: unaligned start and end', async () => {
  const { key, iv, plaintext, ciphertext } = await setup(1000);
  const start = 37;
  const end = 601;
  const got = await decryptRange(key, iv, ciphertext, start, end, 512);
  assert.deepEqual(got, plaintext.subarray(start, end + 1));
});

test('range decryptor: chunk sizes that split blocks', async () => {
  const { key, iv, plaintext, ciphertext } = await setup(1000);
  const start = 100;
  const end = 900;
  const got = await decryptRange(key, iv, ciphertext, start, end, 7);
  assert.deepEqual(got, plaintext.subarray(start, end + 1));
});

test('range decryptor: range starting past the first block', async () => {
  const { key, iv, plaintext, ciphertext } = await setup(2000);
  const start = 1040;
  const end = 1500;
  const got = await decryptRange(key, iv, ciphertext, start, end, 33);
  assert.deepEqual(got, plaintext.subarray(start, end + 1));
});

test('sniff: recognises a few magic numbers', () => {
  assert.equal(imageContentType(Uint8Array.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])), 'image/png');
  assert.equal(imageContentType(Uint8Array.from([0xff, 0xd8, 0xff, 0xe0])), 'image/jpeg');
  assert.equal(attachmentContentType(textBytes('OggS')), 'audio/ogg');
  assert.equal(attachmentContentType(Uint8Array.from([0x1a, 0x45, 0xdf, 0xa3])), 'video/webm');
  const mp4 = concatBytes([new Uint8Array(4), textBytes('ftyp'), textBytes('isom')]);
  assert.equal(attachmentContentType(mp4), 'video/mp4');
  assert.equal(attachmentContentType(textBytes('not a media file')), null);
});