import { describe, it, expect } from "bun:test";
import { importKey, encrypt, decryptRange, ctRange, CHUNK_SIZE, CHUNK_CT_SIZE, encryptMeta, decryptMeta } from "./crypto";

async function makeKey() {
  const key = await crypto.subtle.generateKey(
    { name: "AES-GCM", length: 256 },
    true,
    ["encrypt", "decrypt"],
  );
  const raw = await crypto.subtle.exportKey("raw", key);
  const encoded = Buffer.from(raw).toString("base64url");
  return { key, encoded };
}

async function collectStream(stream: ReadableStream<Uint8Array>): Promise<Uint8Array> {
  const chunks: Uint8Array[] = [];
  const reader = stream.getReader();
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
  }
  if (chunks.length === 1) return chunks[0]!;
  const result = new Uint8Array(chunks.reduce((s, c) => s + c.length, 0));
  let offset = 0;
  for (const chunk of chunks) {
    result.set(chunk, offset);
    offset += chunk.length;
  }
  return result;
}

function toReader(data: Uint8Array): ReadableStreamDefaultReader<Uint8Array> {
  return new ReadableStream<Uint8Array>({
    start(c) { c.enqueue(data); c.close(); },
  }).getReader();
}

async function decrypt(ct: Uint8Array<ArrayBuffer>, key: CryptoKey, fileLen: number) {
  return collectStream(decryptRange(toReader(ct), key, fileLen, 0, Math.max(fileLen - 1, 0)));
}

// Simulate what the SW does: use ctRange to slice ciphertext, feed to decryptRange.
async function decryptSlice(
  ct: Uint8Array<ArrayBuffer>, key: CryptoKey, fileLen: number,
  plainStart: number, plainEnd: number,
) {
  const [ctStart, ctEnd] = ctRange(plainStart, plainEnd);
  const slice = ct.slice(ctStart, ctEnd + 1);
  return collectStream(decryptRange(toReader(slice), key, fileLen, plainStart, plainEnd));
}

describe("encryptMeta / decryptMeta round-trip", () => {
  it("recovers filename, file length, and presentation", async () => {
    const { key } = await makeKey();
    const ct = await encryptMeta("hello.txt", 42, "text", key);
    const { fileName, fileLen, presentation } = await decryptMeta(ct, key);
    expect(fileName).toBe("hello.txt");
    expect(fileLen).toBe(42);
    expect(presentation).toBe("text");
  });

  it("ciphertext is always exactly 4096 bytes", async () => {
    const { key } = await makeKey();
    const ct = await encryptMeta("hello.txt", 42, "text", key);
    expect(ct.byteLength).toBe(4096);
  });

  it("handles null presentation", async () => {
    const { key } = await makeKey();
    const ct = await encryptMeta("archive.zip", 1000, null, key);
    const { presentation } = await decryptMeta(ct, key);
    expect(presentation).toBeNull();
  });

  it("handles a unicode filename", async () => {
    const { key } = await makeKey();
    const ct = await encryptMeta("日本語ファイル.txt", 0, "text", key);
    const { fileName, fileLen } = await decryptMeta(ct, key);
    expect(fileName).toBe("日本語ファイル.txt");
    expect(fileLen).toBe(0);
  });

  it("decryption fails with wrong key", async () => {
    const { key: key1 } = await makeKey();
    const { key: key2 } = await makeKey();
    const ct = await encryptMeta("f.txt", 1, "text", key1);
    expect(decryptMeta(ct, key2)).rejects.toThrow();
  });

  it("throws when filename exceeds 2048 bytes", async () => {
    const { key } = await makeKey();
    expect(encryptMeta("a".repeat(2049), 0, null, key)).rejects.toThrow("Filename too long");
  });
});

describe("encrypt / decrypt round-trip", () => {
  it("recovers the original file data", async () => {
    const { key } = await makeKey();
    const content = new TextEncoder().encode("hello world");
    const ct = await encrypt(content.buffer, key);
    const fileData = await decrypt(ct, key, content.byteLength);
    expect(fileData).toEqual(content);
  });

  it("handles an empty file", async () => {
    const { key } = await makeKey();
    const ct = await encrypt(new ArrayBuffer(0), key);
    const fileData = await decrypt(ct, key, 0);
    expect(fileData.byteLength).toBe(0);
  });

  it("ciphertext is one 16K chunk (+ GCM tag) for small inputs", async () => {
    const { key } = await makeKey();
    const ct = await encrypt(new ArrayBuffer(1), key);
    expect(ct.byteLength).toBe(16384 + 16);
  });

  it("ciphertext grows by one chunk per 16K of input", async () => {
    const { key } = await makeKey();
    const ct1 = await encrypt(new ArrayBuffer(16384), key);
    const ct2 = await encrypt(new ArrayBuffer(16385), key);
    expect(ct1.byteLength).toBe(16384 + 16);
    expect(ct2.byteLength).toBe(2 * (16384 + 16));
  });

  it("different keys produce different ciphertexts", async () => {
    const { key: key1 } = await makeKey();
    const { key: key2 } = await makeKey();
    const content = new TextEncoder().encode("same content").buffer;

    const ct1 = await encrypt(content, key1);
    const ct2 = await encrypt(content, key2);

    expect(ct1).not.toEqual(ct2);
  });

  it("decryption fails with the wrong key", async () => {
    const { key: key1 } = await makeKey();
    const { key: key2 } = await makeKey();
    const content = new TextEncoder().encode("secret").buffer;

    const ct = await encrypt(content, key1);
    expect(decrypt(ct, key2, content.byteLength)).rejects.toThrow();
  });

  it("decryption fails with tampered ciphertext", async () => {
    const { key } = await makeKey();
    const content = new TextEncoder().encode("secret").buffer;
    const ct = await encrypt(content, key);

    ct[0] = (ct[0] ?? 0) ^ 0xff; // flip bits in first byte
    expect(decrypt(ct, key, content.byteLength)).rejects.toThrow();
  });

  it("round-trips an imported (encoded) key", async () => {
    const { encoded } = await makeKey();
    const key = await importKey(encoded);
    const content = new TextEncoder().encode("via imported key").buffer;

    const ct = await encrypt(content, key);
    const fileData = await decrypt(ct, key, content.byteLength);

    expect(new TextDecoder().decode(fileData)).toBe("via imported key");
  });

  it("body and meta use different IVs (same key, different ciphertexts for same data)", async () => {
    const { key } = await makeKey();
    // Encrypt a 13-byte payload both ways — they must differ
    const payload = new TextEncoder().encode("Hello, world!");
    const bodyCt = await encrypt(payload.buffer, key);
    const metaCt = await encryptMeta("Hello, world!", 0, null, key);
    // They must not be equal (different IVs)
    expect(bodyCt.slice(0, metaCt.byteLength)).not.toEqual(metaCt);
  });
});

describe("decryptRange", () => {
  // 3 chunks: chunk 0 = 16K, chunk 1 = 16K, chunk 2 = 5 bytes
  const MULTI_CHUNK_LEN = CHUNK_SIZE * 2 + 5;
  let key: CryptoKey;
  let plain: Uint8Array;
  let ct: Uint8Array<ArrayBuffer>;

  // Fill plaintext with sequential bytes so we can verify any slice
  async function setup() {
    ({ key } = await makeKey());
    plain = new Uint8Array(MULTI_CHUNK_LEN);
    for (let i = 0; i < plain.length; i++) plain[i] = i & 0xff;
    ct = await encrypt(plain.buffer as ArrayBuffer, key);
  }

  it("full range recovers all data", async () => {
    await setup();
    const result = await decryptSlice(ct, key, plain.length, 0, plain.length - 1);
    expect(result).toEqual(plain);
  });

  it("first byte only", async () => {
    await setup();
    const result = await decryptSlice(ct, key, plain.length, 0, 0);
    expect(result).toEqual(plain.slice(0, 1));
  });

  it("last byte only", async () => {
    await setup();
    const last = plain.length - 1;
    const result = await decryptSlice(ct, key, plain.length, last, last);
    expect(result).toEqual(plain.slice(last));
  });

  it("range within a single chunk", async () => {
    await setup();
    const result = await decryptSlice(ct, key, plain.length, 10, 100);
    expect(result).toEqual(plain.slice(10, 101));
  });

  it("range spanning two chunks", async () => {
    await setup();
    const start = CHUNK_SIZE - 5;
    const end = CHUNK_SIZE + 5;
    const result = await decryptSlice(ct, key, plain.length, start, end);
    expect(result).toEqual(plain.slice(start, end + 1));
  });

  it("range spanning all three chunks", async () => {
    await setup();
    const start = 100;
    const end = CHUNK_SIZE * 2 + 3;
    const result = await decryptSlice(ct, key, plain.length, start, end);
    expect(result).toEqual(plain.slice(start, end + 1));
  });

  it("exact chunk boundary as start", async () => {
    await setup();
    const result = await decryptSlice(ct, key, plain.length, CHUNK_SIZE, CHUNK_SIZE + 10);
    expect(result).toEqual(plain.slice(CHUNK_SIZE, CHUNK_SIZE + 11));
  });

  it("exact chunk boundary as end", async () => {
    await setup();
    const result = await decryptSlice(ct, key, plain.length, CHUNK_SIZE - 10, CHUNK_SIZE - 1);
    expect(result).toEqual(plain.slice(CHUNK_SIZE - 10, CHUNK_SIZE));
  });

  it("entire final (partial) chunk", async () => {
    await setup();
    const start = CHUNK_SIZE * 2;
    const end = plain.length - 1;
    const result = await decryptSlice(ct, key, plain.length, start, end);
    expect(result).toEqual(plain.slice(start));
  });

  it("single-chunk file with sub-range", async () => {
    const { key: k } = await makeKey();
    const small = new TextEncoder().encode("hello world");
    const smallCt = await encrypt(small.buffer, k);
    const result = await decryptSlice(smallCt, k, small.length, 2, 7);
    expect(new TextDecoder().decode(result)).toBe("llo wo");
  });

  it("ctRange maps plaintext range to correct ciphertext offsets", () => {
    // Range within first chunk
    expect(ctRange(0, 100)).toEqual([0, CHUNK_CT_SIZE - 1]);
    // Range spanning chunks 0-1
    expect(ctRange(0, CHUNK_SIZE)).toEqual([0, 2 * CHUNK_CT_SIZE - 1]);
    // Range entirely in chunk 2
    expect(ctRange(CHUNK_SIZE * 2, CHUNK_SIZE * 2 + 5)).toEqual([
      2 * CHUNK_CT_SIZE,
      3 * CHUNK_CT_SIZE - 1,
    ]);
  });
});