import { randomUUID } from "node:crypto";
import { mkdirSync } from "node:fs";
import { unlink } from "node:fs/promises";
import { Readable } from "node:stream";
import { Database } from "bun:sqlite";
import cron from "@elysiajs/cron";
import { html } from "@elysiajs/html";
import staticPlugin from "@elysiajs/static";
import busboy from "busboy";
import { Elysia, StatusMap, t } from "elysia";
import { fileTypeFromBuffer } from "file-type";
import {
	filetypes,
	Index,
	NotFound,
	SetCookie,
	ShowFile,
	textPreviewHtml,
	WrongPassword,
	type Preview,
} from "./components";
import { config } from "./config";
import { decryptToStream, encryptStream } from "./crypto";

const BLOB_DIR = "./db/blobs";
mkdirSync(BLOB_DIR, { recursive: true });
const blobPath = (uuid: string) => `${BLOB_DIR}/${uuid}`;
const safeUnlink = (path: string) => unlink(path).catch(() => {});

const filetypeSet = new Set(filetypes);
const SNIFF_BYTES = 4100; // enough for file-type's magic-number detection

const db = new Database("./db/db.sqlite");
db.run("PRAGMA foreign_keys = ON");
db.run("PRAGMA journal_mode = WAL");
// Content is stored on disk at ./db/blobs/<uuid>; the row keeps only metadata.
// `size` is the on-disk byte count (post-encryption) and drives the total-bytes
// cap; `media_mime` is the MIME sniffed from the plaintext head at upload time,
// letting /show preview media without decrypting.
db.run(
	"CREATE TABLE IF NOT EXISTS files (uuid TEXT PRIMARY KEY, filename TEXT NOT NULL, filetype TEXT NOT NULL, encrypted INTEGER NOT NULL, size INTEGER NOT NULL, media_mime TEXT, delete_at INTEGER) STRICT",
);
db.run("PRAGMA optimize");

// Running total of stored content bytes, initialized once from the DB and then
// maintained in memory (incremented on upload, decremented when files expire).
let totalBytes = (
	db.prepare("SELECT COALESCE(SUM(size), 0) AS t FROM files").get() as {
		t: number;
	}
).t;

// uuid route params are constrained to this shape so they can't be used to
// inject CRLF/extra directives into the Set-Cookie Path or content-disposition,
// or to escape the blob directory.
const UUID_PATTERN = "^[0-9a-fA-F-]{36}$";

// Per-IP timestamp of the last accepted upload / last server-side decryption
// attempt, used for the respective cooldowns. Pruned by the cron below so they
// can't grow without bound.
const lastUpload = new Map<string, number>();
const lastDecrypt = new Map<string, number>();

type MinimalServer = {
	requestIP(req: Request): { address: string } | null;
} | null;

function clientIp(
	server: MinimalServer,
	request: Request,
	headers: Record<string, string | undefined>,
): string {
	if (config.behindProxy) {
		const xff = headers["x-forwarded-for"]?.split(",")[0]?.trim();
		if (xff) return xff;
		const real = headers["x-real-ip"];
		if (real) return real;
	}
	return server?.requestIP(request)?.address ?? "unknown";
}

// An upload failure that maps to a specific HTTP status + message.
class UploadError extends Error {
	constructor(
		readonly status: number,
		message: string,
	) {
		super(message);
	}
}

async function pump(
	stream: ReadableStream<Uint8Array>,
	onChunk: (chunk: Uint8Array) => void | Promise<void>,
): Promise<void> {
	const reader = stream.getReader();
	while (true) {
		const { done, value } = await reader.read();
		if (done) break;
		// Await the callback so a slow sink applies backpressure (Bun's FileSink
		// .write returns a Promise when the write is still pending) instead of
		// buffering the whole upload in memory.
		await onChunk(value);
	}
}

async function sniffMime(bytes: Uint8Array): Promise<string | null> {
	if (bytes.length === 0) return null;
	return (await fileTypeFromBuffer(bytes))?.mime ?? null;
}

const app = new Elysia({
	serve: {
		maxRequestBodySize: config.maxUploadBytes,
	},
})
	.use(staticPlugin({ assets: "./assets", prefix: "/" }))
	.use(html())
	.use(
		cron({
			name: "delete",
			pattern: "*/5 * * * * *",
			async run() {
				const expired = db
					.prepare(
						"SELECT uuid, size FROM files WHERE delete_at < strftime('%s', 'now')",
					)
					.all() as { uuid: string; size: number }[];
				if (expired.length > 0) {
					// Delete by the exact uuids selected above, not a second
					// strftime('now') comparison (which evaluates at a later instant and
					// could delete a row we didn't account for here — leaking the counter
					// and orphaning its blob).
					const placeholders = expired.map(() => "?").join(",");
					db.exec(
						`DELETE FROM files WHERE uuid IN (${placeholders})`,
						expired.map((e) => e.uuid),
					);
					for (const { uuid, size } of expired) {
						await safeUnlink(blobPath(uuid));
						totalBytes -= size;
					}
					if (totalBytes < 0) totalBytes = 0;
				}
				const now = Date.now();
				if (config.uploadCooldownSeconds > 0) {
					const cutoff = now - config.uploadCooldownSeconds * 1000;
					for (const [ip, ts] of lastUpload) {
						if (ts < cutoff) lastUpload.delete(ip);
					}
				}
				if (config.decryptCooldownSeconds > 0) {
					const cutoff = now - config.decryptCooldownSeconds * 1000;
					for (const [ip, ts] of lastDecrypt) {
						if (ts < cutoff) lastDecrypt.delete(ip);
					}
				}
			},
		}),
	)
	.use(
		cron({
			// Run SQLite's optimizer periodically (not just at startup), per its docs.
			name: "optimize",
			pattern: "0 0 * * * *",
			run() {
				db.exec("PRAGMA optimize");
			},
		}),
	)
	.get("/", ({ server }) => Index(server?.url.toString() ?? ""))
	.post(
		"/upload",
		async ({ set, server, request, headers }) => {
			const ip = clientIp(server, request, headers);
			const now = Date.now();
			if (config.uploadCooldownSeconds > 0) {
				const last = lastUpload.get(ip) ?? 0;
				if (now - last < config.uploadCooldownSeconds * 1000) {
					set.status = 429; // Too Many Requests
					return "Upload cooldown active, please wait before uploading again";
				}
			}

			const contentType = request.headers.get("content-type") ?? "";
			if (!contentType.includes("multipart/form-data") || !request.body) {
				set.status = 400;
				return "Expected a multipart/form-data upload";
			}

			const uuid = randomUUID();
			const path = blobPath(uuid);

			// The in-flight blob write. A rejection can settle the upload while this
			// is still draining to disk; the catch awaits it before unlinking so a
			// late write can't recreate the blob after cleanup (orphan).
			let writing: Promise<unknown> | undefined;

			try {
				// Parse the multipart body as it streams in. The file part must come
				// LAST so the other fields (password, filetype, ...) are known before
				// the bytes flow and can drive on-the-fly encryption to disk.
				const result = await new Promise<{
					filename: string;
					filetype: string;
					encrypted: boolean;
					size: number;
					mediaMime: string | null;
					deleteAt: number | null;
				}>((resolve, reject) => {
					const bb = busboy({
						headers: { "content-type": contentType },
						limits: { files: 1, fileSize: config.maxUploadBytes },
					});
					const fields: Record<string, string> = {};
					let fileSeen = false;
					// A form field arriving after the file part means the file wasn't
					// sent last. We don't act on it here (the file is still streaming) —
					// we record it and report it once the whole body is parsed, so the
					// "file must be last" error always wins over an incidental symptom
					// like a not-yet-seen filetype.
					let fieldAfterFile = false;

					bb.on("field", (name, value) => {
						if (fileSeen) fieldAfterFile = true;
						else fields[name] = value;
					});

					bb.on("file", (name, stream, info) => {
						if (name !== "file") {
							stream.resume();
							return;
						}
						fileSeen = true;
						writing = streamToBlob(stream, info);
					});

					bb.on("error", reject);

					// Settle once the entire body is parsed: by now the field set and the
					// file's position relative to other fields are both fully known.
					bb.on("close", async () => {
						try {
							if (!fileSeen) throw new UploadError(400, "No file provided");
							if (fieldAfterFile) {
								throw new UploadError(
									400,
									"the file field must be the last form field",
								);
							}
							// Set because a file part was seen; await it to surface any
							// streaming error and read the stored size/type.
							const file = await (writing as ReturnType<typeof streamToBlob>);

							const filetype = fields.filetype ?? "";
							if (!filetypeSet.has(filetype)) {
								throw new UploadError(400, "Invalid or missing filetype");
							}
							const dim = fields.delete_in_minutes;
							if (dim && !/^[0-9]+$/.test(dim)) {
								throw new UploadError(400, "Invalid delete_in_minutes");
							}
							let minutes: number | null =
								dim && Number(dim) > 0 ? Number(dim) : null;
							if (config.maxAgeMinutes !== null) {
								minutes = Math.min(
									minutes ?? config.maxAgeMinutes,
									config.maxAgeMinutes,
								);
							}
							const deleteAt =
								minutes !== null ? Math.floor(now / 1000) + minutes * 60 : null;

							resolve({
								filename: fields.filename || file.infoFilename || "file",
								filetype,
								encrypted: file.encrypted,
								size: file.size,
								mediaMime: file.mediaMime,
								deleteAt,
							});
						} catch (e) {
							reject(e);
						}
					});

					// Streams the file part to ./db/blobs/<uuid>, encrypting on the fly
					// when a password was provided (server-side) or storing opaque bytes
					// for an already-encrypted upload. Uses the fields seen so far; if the
					// file wasn't last that set is incomplete, but the close handler
					// rejects such uploads before anything is persisted.
					async function streamToBlob(stream: Readable, info: busboy.FileInfo) {
						let limitExceeded = false;
						stream.on("limit", () => {
							limitExceeded = true;
						});
						const webIn = Readable.toWeb(
							stream,
						) as unknown as ReadableStream<Uint8Array>;

						const sink = Bun.file(path).writer();
						let size = 0;
						let mediaMime: string | null = null;
						let encrypted = false;

						if (fields.encrypted === "on") {
							// Client-side encrypted: opaque bytes, store as-is, no sniffing.
							encrypted = true;
							await pump(webIn, async (c) => {
								size += c.length;
								await sink.write(c);
							});
						} else if (fields.password) {
							// Server-side encryption: encrypt the stream to disk, sniffing
							// the plaintext head for a preview MIME type.
							encrypted = true;
							let head: Uint8Array | undefined;
							const cipher = await encryptStream(
								webIn,
								fields.password,
								(h) => {
									head = h;
								},
								SNIFF_BYTES,
							);
							await pump(cipher, async (c) => {
								size += c.length;
								await sink.write(c);
							});
							if (head) mediaMime = await sniffMime(head);
						} else {
							// Plaintext: stream to disk, collecting the head for sniffing.
							const headParts: Uint8Array[] = [];
							let headLen = 0;
							await pump(webIn, async (c) => {
								size += c.length;
								await sink.write(c);
								if (headLen < SNIFF_BYTES) {
									const slice = c.subarray(0, SNIFF_BYTES - headLen);
									headParts.push(slice);
									headLen += slice.length;
								}
							});
							mediaMime = await sniffMime(Buffer.concat(headParts));
						}

						await sink.end();
						if (limitExceeded) {
							throw new UploadError(
								413,
								`File exceeds the maximum upload size of ${config.maxUploadBytes} bytes`,
							);
						}
						return {
							encrypted,
							size,
							mediaMime,
							infoFilename: info.filename ?? "",
						};
					}

					Readable.fromWeb(
						request.body as unknown as import("node:stream/web").ReadableStream<Uint8Array>,
					).pipe(bb);
				});

				// Enforce the total-bytes cap against the in-memory counter. Concurrent
				// uploads can transiently overshoot by up to (concurrency * per-file)
				// before this check; acceptable at the expected scale.
				if (
					config.maxTotalBytes !== null &&
					totalBytes + result.size > config.maxTotalBytes
				) {
					await safeUnlink(path);
					set.status = 507; // Insufficient Storage
					return "Server storage is full, try again later";
				}

				db.exec(
					"INSERT INTO files (uuid, filename, filetype, encrypted, size, media_mime, delete_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
					[
						uuid,
						result.filename,
						result.filetype,
						result.encrypted,
						result.size,
						result.mediaMime,
						result.deleteAt,
					],
				);
				totalBytes += result.size;
				if (config.uploadCooldownSeconds > 0) lastUpload.set(ip, now);
				set.status = StatusMap["See Other"];
				set.headers.location = `/show/${uuid}`;
				return `Created with id: ${uuid}`;
			} catch (e) {
				// Let any in-flight write finish so it can't recreate the blob after
				// we unlink it below.
				if (writing) await writing.catch(() => {});
				await safeUnlink(path);
				if (e instanceof UploadError) {
					set.status = e.status;
					return e.message;
				}
				throw e;
			}
		},
		{
			// Body parsing is handled manually from the raw stream, so disable
			// Elysia's parser (which would otherwise buffer the whole upload).
			parse: "none",
		},
	)
	.get(
		"/show/:uuid",
		async ({ set, params, cookie, server, request, headers }) => {
			const row =
				(db
					.prepare(
						"SELECT filename, filetype, encrypted, size, media_mime, delete_at FROM files WHERE uuid = ?",
					)
					.get(params.uuid) as {
					filename: string;
					filetype: string;
					encrypted: number;
					size: number;
					media_mime: string | null;
					delete_at: number | null;
				}) || null;
			if (!row) {
				set.status = StatusMap["Not Found"];
				return NotFound();
			}

			const path = blobPath(params.uuid);
			let preview: Preview;
			let shownSize: number | null = row.size;

			if (row.encrypted) {
				const password = cookie.password.value;
				if (!password) {
					// No password yet: let the client-side flow handle decryption.
					preview = { kind: "await" };
					shownSize = null;
				} else if (row.filetype === "blob") {
					// Binary/media: don't decrypt here — preview points at /raw, which
					// performs the single decryption for this view.
					preview = row.media_mime
						? { kind: "media", mime: row.media_mime }
						: { kind: "none" };
				} else {
					// Text: this is the one server-side decryption for the view.
					const ip = clientIp(server, request, headers);
					if (decryptBlocked(ip)) {
						set.status = 429;
						return "Decryption cooldown active, please wait and reload";
					}
					try {
						const content = await readDecrypted(path, password, row.size);
						recordDecrypt(ip);
						const out = new Uint8Array(content);
						shownSize = out.byteLength;
						const text = textPreviewHtml(out, row.filetype);
						preview =
							text !== null ? { kind: "text", html: text } : { kind: "none" };
					} catch (_e) {
						recordDecrypt(ip);
						const secure = config.behindProxy ? "; Secure" : "";
						set.status = StatusMap.Forbidden;
						set.headers["set-cookie"] = [
							`password=; Path=/show/${params.uuid}; SameSite=lax; HttpOnly${secure}; Expires=Thu, 01 Jan 1970 00:00:00 GMT`,
							`password=; Path=/raw/${params.uuid}; SameSite=lax; HttpOnly${secure}; Expires=Thu, 01 Jan 1970 00:00:00 GMT`,
						];
						return WrongPassword();
					}
				}
			} else if (row.filetype !== "blob") {
				// Plaintext text: read from disk and render inline.
				const content = new Uint8Array(await Bun.file(path).bytes());
				shownSize = content.byteLength;
				const text = textPreviewHtml(content, row.filetype);
				preview =
					text !== null
						? { kind: "text", html: text }
						: row.media_mime
							? { kind: "media", mime: row.media_mime }
							: { kind: "none" };
			} else {
				// Plaintext binary/media: preview via /raw using the sniffed MIME.
				preview = row.media_mime
					? { kind: "media", mime: row.media_mime }
					: { kind: "none" };
			}

			return ShowFile({
				filename: row.filename,
				uuid: params.uuid,
				filetype: row.filetype,
				deleteAt: row.delete_at,
				size: shownSize,
				preview,
			});
		},
		{
			params: t.Object({
				uuid: t.String({ format: "regex", pattern: UUID_PATTERN }),
			}),
			cookie: t.Object({ password: t.Optional(t.String()) }),
		},
	)
	.post(
		"/set-cookie/:uuid",
		({ set, body, params }) => {
			// encodeURIComponent keeps ';', CR/LF and other separators out of the
			// cookie value; Elysia URL-decodes the value again when it reads it back.
			const value = encodeURIComponent(body.password);
			const secure = config.behindProxy ? "; Secure" : "";
			set.headers["set-cookie"] = [
				`password=${value}; Path=/show/${params.uuid}; SameSite=lax; HttpOnly${secure}`,
				`password=${value}; Path=/raw/${params.uuid}; SameSite=lax; HttpOnly${secure}`,
			];
			set.headers.location = `/show/${params.uuid}`;
			set.status = StatusMap["See Other"];
			return SetCookie(params.uuid);
		},
		{
			body: t.Object({
				password: t.String(),
			}),
			params: t.Object({
				uuid: t.String({ format: "regex", pattern: UUID_PATTERN }),
			}),
		},
	)
	.get(
		"/raw/:uuid",
		async ({ set, params, cookie, query, server, request, headers }) => {
			const row =
				(db
					.prepare(
						"SELECT filename, encrypted, filetype, size FROM files WHERE uuid = ?",
					)
					.get(params.uuid) as {
					filename: string;
					encrypted: number;
					filetype: string;
					size: number;
				}) || null;
			if (!row) {
				set.status = StatusMap["Not Found"];
				return "File not found";
			}

			const path = blobPath(params.uuid);
			const servingEncrypted =
				row.encrypted && query.ignore_password === "true";

			let body: ReadableStream<Uint8Array> | ReturnType<typeof Bun.file>;
			let sniff: Uint8Array | null = null;

			if (row.encrypted && !servingEncrypted) {
				const password = cookie.password.value;
				if (!password) {
					set.status = StatusMap.Unauthorized;
					return 'This file is encrypted, set the cookie "password" with the correct password to allow the server to decrypt it';
				}
				const ip = clientIp(server, request, headers);
				if (decryptBlocked(ip)) {
					set.status = 429;
					return "Decryption cooldown active, please wait and retry";
				}
				try {
					const { firstChunk, stream } = await openDecrypted(
						path,
						password,
						row.size,
					);
					recordDecrypt(ip);
					body = stream;
					sniff = firstChunk.subarray(0, SNIFF_BYTES);
				} catch (_e) {
					recordDecrypt(ip);
					set.status = StatusMap.Forbidden;
					return "Incorrect password";
				}
			} else {
				// Serve the file as-is: plaintext, or (with ignore_password) the still
				// encrypted bytes. BunFile streams and supports range requests.
				body = Bun.file(path);
				if (!servingEncrypted) {
					sniff = new Uint8Array(
						await Bun.file(path).slice(0, SNIFF_BYTES).arrayBuffer(),
					);
				}
			}

			// Never let the browser sniff stored content into an executable type
			// (e.g. HTML/SVG running as same-origin script). Only whitelisted,
			// non-scriptable media is served inline with its real type; everything
			// else (incl. still-encrypted bytes) is an octet-stream attachment.
			let mime = "application/octet-stream";
			let disposition = "attachment";
			if (sniff) {
				const detected = await fileTypeFromBuffer(sniff);
				if (
					detected &&
					detected.mime !== "image/svg+xml" &&
					(detected.mime.startsWith("image/") ||
						detected.mime.startsWith("audio/") ||
						detected.mime.startsWith("video/"))
				) {
					mime = detected.mime;
					disposition = "inline";
				}
			}

			const safeName = encodeURIComponent(row.filename);
			set.headers["x-content-type-options"] = "nosniff";
			set.headers["content-type"] = mime;
			set.headers.encrypted = row.encrypted ? "true" : "false";
			set.headers.filetype = row.filetype;
			set.headers.filename = safeName;
			set.headers["content-disposition"] =
				`${disposition}; filename*=UTF-8''${safeName}`;
			return body;
		},
		{
			params: t.Object({
				uuid: t.String({ format: "regex", pattern: UUID_PATTERN }),
			}),
			cookie: t.Object({ password: t.Optional(t.String()) }),
			query: t.Object({ ignore_password: t.Optional(t.String()) }),
		},
	)
	.listen(3000);

// --- server-side decryption cooldown ----------------------------------------

function decryptBlocked(ip: string): boolean {
	if (config.decryptCooldownSeconds <= 0) return false;
	const last = lastDecrypt.get(ip) ?? 0;
	return Date.now() - last < config.decryptCooldownSeconds * 1000;
}
function recordDecrypt(ip: string): void {
	if (config.decryptCooldownSeconds > 0) lastDecrypt.set(ip, Date.now());
}

// Fully decrypt an on-disk blob to memory (used for text previews in /show).
async function readDecrypted(
	path: string,
	password: string,
	size: number,
): Promise<Uint8Array> {
	const { stream } = await openDecrypted(path, password, size);
	const parts: Uint8Array[] = [];
	await pump(stream, (c) => {
		parts.push(c);
	});
	return Buffer.concat(parts);
}

// Open an on-disk blob for streaming decryption, exposing the first plaintext
// chunk (for MIME sniffing) and the full plaintext stream from one derivation.
async function openDecrypted(
	path: string,
	password: string,
	size: number,
): Promise<{ firstChunk: Uint8Array; stream: ReadableStream<Uint8Array> }> {
	const { firstChunk, body } = await decryptToStream(
		Bun.file(path).stream(),
		password,
		size,
	);
	return { firstChunk, stream: body };
}

console.log(
	`⚡ ZBin is running at ${app.server?.hostname}:${app.server?.port} ⚡`,
);