// Chunked, streaming AES-256-GCM (authenticated) encryption with a
// PBKDF2-derived key, using WebCrypto (crypto.subtle) so the exact same format
// works in Bun (server) and the browser (client), and is matched byte-for-byte
// by assets/encrypt.py.
//
// The content is split into fixed-size plaintext chunks, each sealed
// independently with AES-GCM. This lets the server encrypt on upload and
// decrypt on download *while streaming* (chunk by chunk) instead of holding the
// whole file in memory. The key is still derived only once per file, so the
// expensive PBKDF2 cost is paid once regardless of size.
//
// Wire format:  salt[16] | noncePrefix[7] | encChunk_0 | encChunk_1 | ...
//   encChunk_i  = AES-GCM(key, nonce_i, plaintext_i)   (16-byte tag appended)
//   nonce_i(12) = noncePrefix[7] || uint32_be(i) || flag   (flag=1 on the final
//                 chunk, else 0)
// The per-chunk counter prevents reordering and the final-chunk flag prevents
// truncation: dropping or rearranging chunks fails authentication. Each
// non-final plaintext chunk is exactly CHUNK bytes, so the decryptor can derive
// chunk boundaries (and which chunk is last) from the total length alone —
// no per-chunk length prefixes are stored.

const SALT_LEN = 16;
const PREFIX_LEN = 7;
const TAG_LEN = 16;
const HEADER_LEN = SALT_LEN + PREFIX_LEN; // bytes before the first chunk
export const CHUNK = 64 * 1024; // plaintext bytes per chunk
const ENC_CHUNK = CHUNK + TAG_LEN; // ciphertext bytes for a full chunk

// `crypto.subtle` wants a BufferSource; Bun's lib types are stricter about
// ArrayBuffer vs SharedArrayBuffer than the values we actually pass, so cast at
// the boundary rather than littering call sites with copies.
const src = (b: Uint8Array): BufferSource => b as unknown as BufferSource;

async function deriveKey(
	password: string,
	salt: Uint8Array,
	usage: KeyUsage[],
): Promise<CryptoKey> {
	const material = await crypto.subtle.importKey(
		"raw",
		src(new TextEncoder().encode(password)),
		{ name: "PBKDF2" },
		false,
		["deriveKey"],
	);
	return crypto.subtle.deriveKey(
		{ name: "PBKDF2", salt: src(salt), iterations: 210000, hash: "SHA-512" },
		material,
		{ name: "AES-GCM", length: 256 },
		false,
		usage,
	);
}

function nonce(prefix: Uint8Array, index: number, last: boolean): Uint8Array {
	const n = new Uint8Array(12);
	n.set(prefix, 0);
	// 4-byte big-endian chunk counter at offset 7.
	n[7] = (index >>> 24) & 0xff;
	n[8] = (index >>> 16) & 0xff;
	n[9] = (index >>> 8) & 0xff;
	n[10] = index & 0xff;
	n[11] = last ? 1 : 0;
	return n;
}

function concat(parts: Uint8Array[]): Uint8Array {
	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;
}

async function sealChunk(
	key: CryptoKey,
	prefix: Uint8Array,
	index: number,
	last: boolean,
	plaintext: Uint8Array,
): Promise<Uint8Array> {
	const ct = await crypto.subtle.encrypt(
		{ name: "AES-GCM", iv: src(nonce(prefix, index, last)) },
		key,
		src(plaintext),
	);
	return new Uint8Array(ct);
}

async function openChunk(
	key: CryptoKey,
	prefix: Uint8Array,
	index: number,
	last: boolean,
	ciphertext: Uint8Array,
): Promise<Uint8Array> {
	const pt = await crypto.subtle.decrypt(
		{ name: "AES-GCM", iv: src(nonce(prefix, index, last)) },
		key,
		src(ciphertext),
	);
	return new Uint8Array(pt);
}

// ----------------------------------------------------------------------------
// One-shot helpers (browser uploads/downloads, tests, python parity). These
// buffer the whole payload; the server uses the streaming variants below.
// ----------------------------------------------------------------------------

export async function encrypt(
	content: Uint8Array,
	password: string,
): Promise<Uint8Array> {
	const salt = crypto.getRandomValues(new Uint8Array(SALT_LEN));
	const prefix = crypto.getRandomValues(new Uint8Array(PREFIX_LEN));
	const key = await deriveKey(password, salt, ["encrypt"]);
	const out: Uint8Array[] = [salt, prefix];
	// Always emit at least one (possibly empty) final chunk so empty input still
	// round-trips and the final-chunk flag is always present.
	const chunks = Math.max(1, Math.ceil(content.length / CHUNK));
	for (let i = 0; i < chunks; i++) {
		const start = i * CHUNK;
		const slice = content.subarray(
			start,
			Math.min(start + CHUNK, content.length),
		);
		out.push(await sealChunk(key, prefix, i, i === chunks - 1, slice));
	}
	return concat(out);
}

export async function decrypt(
	data: Uint8Array,
	password: string,
): Promise<Uint8Array> {
	if (data.length < HEADER_LEN + TAG_LEN) {
		throw new Error("ciphertext too short");
	}
	const salt = data.subarray(0, SALT_LEN);
	const prefix = data.subarray(SALT_LEN, HEADER_LEN);
	const key = await deriveKey(password, salt, ["decrypt"]);
	const body = data.subarray(HEADER_LEN);
	const out: Uint8Array[] = [];
	let offset = 0;
	let index = 0;
	while (offset < body.length) {
		const encLen = Math.min(ENC_CHUNK, body.length - offset);
		const last = offset + encLen >= body.length;
		out.push(
			await openChunk(
				key,
				prefix,
				index,
				last,
				body.subarray(offset, offset + encLen),
			),
		);
		offset += encLen;
		index++;
	}
	return concat(out);
}

// ----------------------------------------------------------------------------
// Streaming helpers (server-side encrypt on upload / decrypt on download).
// ----------------------------------------------------------------------------

// A small pull-based reader over a ReadableStream that can hand back exact byte
// counts, buffering only the unconsumed remainder.
class ByteStreamReader {
	#reader: ReadableStreamDefaultReader<Uint8Array>;
	#buf: Uint8Array = new Uint8Array(0);
	#done = false;

	constructor(stream: ReadableStream<Uint8Array>) {
		this.#reader = stream.getReader();
	}

	async #fill(): Promise<boolean> {
		if (this.#done) return false;
		const { done, value } = await this.#reader.read();
		if (done) {
			this.#done = true;
			return false;
		}
		this.#buf = this.#buf.length === 0 ? value : concat([this.#buf, value]);
		return true;
	}

	// Reads exactly `n` bytes, or throws if the stream ends first.
	async readExact(n: number): Promise<Uint8Array> {
		while (this.#buf.length < n) {
			if (!(await this.#fill())) throw new Error("unexpected end of stream");
		}
		const out = this.#buf.subarray(0, n);
		this.#buf = this.#buf.subarray(n);
		return out;
	}

	// Reads up to `n` bytes; returns fewer only at end of stream.
	async readUpTo(n: number): Promise<Uint8Array> {
		while (this.#buf.length < n) {
			if (!(await this.#fill())) break;
		}
		const take = Math.min(n, this.#buf.length);
		const out = this.#buf.subarray(0, take);
		this.#buf = this.#buf.subarray(take);
		return out;
	}

	// True if more bytes remain, without consuming them.
	async hasMore(): Promise<boolean> {
		while (this.#buf.length === 0) {
			if (!(await this.#fill())) return false;
		}
		return true;
	}

	// Cancels the source and releases the reader lock. Used when the consuming
	// stream is cancelled (e.g. a client aborts a download) so the underlying
	// file handle isn't held until GC.
	async cancel(): Promise<void> {
		this.#done = true;
		await this.#reader.cancel().catch(() => {});
	}
}

// Encrypts a plaintext stream into the chunked wire format. `onHead`, if given,
// is invoked once with up to `headBytes` of leading plaintext (used to sniff a
// media type for previews) before the stream completes.
export async function encryptStream(
	input: ReadableStream<Uint8Array>,
	password: string,
	onHead?: (head: Uint8Array) => void,
	headBytes = 4100,
): Promise<ReadableStream<Uint8Array>> {
	const salt = crypto.getRandomValues(new Uint8Array(SALT_LEN));
	const prefix = crypto.getRandomValues(new Uint8Array(PREFIX_LEN));
	const key = await deriveKey(password, salt, ["encrypt"]);
	const reader = new ByteStreamReader(input);
	let index = 0;
	let headDone = onHead === undefined;
	const headParts: Uint8Array[] = [];
	let headLen = 0;

	function recordHead(chunk: Uint8Array) {
		if (headDone) return;
		const need = headBytes - headLen;
		if (need > 0) {
			const slice = chunk.subarray(0, need);
			headParts.push(slice);
			headLen += slice.length;
		}
	}
	function flushHead() {
		if (!headDone) {
			headDone = true;
			onHead?.(concat(headParts));
		}
	}

	return new ReadableStream<Uint8Array>({
		start(controller) {
			// salt + noncePrefix come first, before any chunk.
			controller.enqueue(concat([salt, prefix]));
		},
		async pull(controller) {
			const current = await reader.readUpTo(CHUNK);
			recordHead(current);
			// Peek (without consuming) whether more plaintext follows so the final
			// chunk's flag is set correctly. An empty input still yields one final
			// (empty) chunk on the first pull.
			const last = !(await reader.hasMore());
			controller.enqueue(await sealChunk(key, prefix, index, last, current));
			index++;
			if (last) {
				flushHead();
				controller.close();
			}
		},
		cancel() {
			return reader.cancel();
		},
	});
}

// Decrypts a stream in the chunked wire format. `totalLen` is the full byte
// length of the source (e.g. the on-disk file size) so chunk boundaries and the
// final chunk can be derived. Returns the first plaintext chunk eagerly (for
// MIME sniffing) plus a `body` stream that re-emits it and then the rest, all
// from a single key derivation.
export async function decryptToStream(
	source: ReadableStream<Uint8Array>,
	password: string,
	totalLen: number,
): Promise<{ firstChunk: Uint8Array; body: ReadableStream<Uint8Array> }> {
	const dataLen = totalLen - HEADER_LEN;
	if (dataLen < TAG_LEN) throw new Error("ciphertext too short");
	const reader = new ByteStreamReader(source);
	const header = await reader.readExact(HEADER_LEN);
	const salt = header.subarray(0, SALT_LEN);
	const prefix = header.subarray(SALT_LEN, HEADER_LEN);
	const key = await deriveKey(password, salt, ["decrypt"]);

	let consumed = 0;
	let index = 0;
	async function next(): Promise<Uint8Array | null> {
		if (consumed >= dataLen) return null;
		const encLen = Math.min(ENC_CHUNK, dataLen - consumed);
		const last = consumed + encLen >= dataLen;
		const ct = await reader.readExact(encLen);
		consumed += encLen;
		const pt = await openChunk(key, prefix, index, last, ct);
		index++;
		return pt;
	}

	// Decrypt the first chunk now (throws on a wrong password / tamper).
	const firstChunk = (await next()) ?? new Uint8Array(0);
	let firstEmitted = false;
	const body = new ReadableStream<Uint8Array>({
		async pull(controller) {
			if (!firstEmitted) {
				firstEmitted = true;
				controller.enqueue(firstChunk);
				return;
			}
			const chunk = await next();
			if (chunk === null) controller.close();
			else controller.enqueue(chunk);
		},
		cancel() {
			return reader.cancel();
		},
	});
	return { firstChunk, body };
}