import { describe, expect, test } from "bun:test";
import {
	CHUNK,
	decrypt,
	decryptToStream,
	encrypt,
	encryptStream,
} from "./crypto";

function streamOf(data: Uint8Array, pieces = 1): ReadableStream<Uint8Array> {
	const size = Math.ceil(data.length / pieces) || 1;
	let off = 0;
	return new ReadableStream<Uint8Array>({
		pull(c) {
			if (off >= data.length) {
				c.close();
				return;
			}
			c.enqueue(data.subarray(off, off + size));
			off += size;
		},
	});
}

async function collect(
	stream: ReadableStream<Uint8Array>,
): Promise<Uint8Array> {
	const parts: Uint8Array[] = [];
	const reader = stream.getReader();
	while (true) {
		const { done, value } = await reader.read();
		if (done) break;
		parts.push(value);
	}
	let total = 0;
	for (const p of parts) total += p.length;
	const out = new Uint8Array(total);
	let off = 0;
	for (const p of parts) {
		out.set(p, off);
		off += p.length;
	}
	return out;
}

describe("crypto (chunked AES-256-GCM)", () => {
	test("round-trips content with the correct password", async () => {
		const data = new TextEncoder().encode("hello zbin 🔐 multi-byte");
		const enc = await encrypt(data, "correct horse battery staple");
		const dec = await decrypt(enc, "correct horse battery staple");
		expect(new TextDecoder().decode(dec)).toBe("hello zbin 🔐 multi-byte");
	});

	test("wire format is salt[16] | prefix[7] | ct+tag", async () => {
		const enc = await encrypt(new Uint8Array([1, 2, 3]), "pw");
		// 16 (salt) + 7 (prefix) + 3 (plaintext) + 16 (GCM tag)
		expect(enc.length).toBe(16 + 7 + 3 + 16);
	});

	test("round-trips empty input", async () => {
		const enc = await encrypt(new Uint8Array(0), "pw");
		expect(enc.length).toBe(16 + 7 + 16); // single empty final chunk (tag only)
		const dec = await decrypt(enc, "pw");
		expect(dec.length).toBe(0);
	});

	test("round-trips multi-chunk content (> CHUNK)", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK * 2 + 1234));
		const enc = await encrypt(data, "pw");
		const dec = await decrypt(enc, "pw");
		expect(dec).toEqual(data);
	});

	test("round-trips content of exactly CHUNK bytes", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK));
		const dec = await decrypt(await encrypt(data, "pw"), "pw");
		expect(dec).toEqual(data);
	});

	test("rejects a wrong password (authenticated)", async () => {
		const enc = await encrypt(new Uint8Array([1, 2, 3]), "right");
		await expect(decrypt(enc, "wrong")).rejects.toThrow();
	});

	test("rejects tampered ciphertext", async () => {
		const enc = await encrypt(new Uint8Array([9, 9, 9]), "pw");
		enc[enc.length - 1] ^= 0xff; // flip a tag byte
		await expect(decrypt(enc, "pw")).rejects.toThrow();
	});

	test("rejects truncation (dropping the final chunk)", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK * 2));
		const enc = await encrypt(data, "pw");
		// Drop the final (full) chunk; the now-last chunk was sealed with flag=0
		// but will be opened with flag=1, so authentication must fail.
		const truncated = enc.subarray(0, enc.length - (CHUNK + 16));
		await expect(decrypt(truncated, "pw")).rejects.toThrow();
	});
});

describe("crypto streaming", () => {
	test("encryptStream output decrypts via one-shot decrypt", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK * 3 + 7));
		const enc = await collect(await encryptStream(streamOf(data, 5), "pw"));
		expect(await decrypt(enc, "pw")).toEqual(data);
	});

	test("encryptStream invokes onHead with leading plaintext", async () => {
		const data = new TextEncoder().encode("GIF89a-ish header then more data");
		let head: Uint8Array | undefined;
		await collect(
			await encryptStream(streamOf(data, 3), "pw", (h) => {
				head = h;
			}),
		);
		expect(head).toBeDefined();
		expect(new TextDecoder().decode((head as Uint8Array).subarray(0, 6))).toBe(
			"GIF89a",
		);
	});

	test("decryptToStream round-trips and exposes the first chunk", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK + 500));
		const enc = await encrypt(data, "pw");
		const { firstChunk, body } = await decryptToStream(
			streamOf(enc, 4),
			"pw",
			enc.length,
		);
		expect(firstChunk).toEqual(data.subarray(0, CHUNK));
		expect(await collect(body)).toEqual(data);
	});

	test("decryptToStream rejects a wrong password", async () => {
		const enc = await encrypt(new Uint8Array([1, 2, 3]), "right");
		await expect(
			decryptToStream(streamOf(enc), "wrong", enc.length),
		).rejects.toThrow();
	});

	test("stream encrypt -> stream decrypt round-trips", async () => {
		const data = crypto.getRandomValues(new Uint8Array(CHUNK * 2 + 42));
		const enc = await collect(await encryptStream(streamOf(data, 7), "pw"));
		const { body } = await decryptToStream(streamOf(enc, 3), "pw", enc.length);
		expect(await collect(body)).toEqual(data);
	});
});