import { createHash } from "node:crypto";
import { availableParallelism } from "node:os";
import { getHeapStatistics } from "node:v8";
import type { Worker } from "node:worker_threads";
import {
	FAILED_RENDER_TTL_MS,
	MAX_RENDER_WORKERS,
	RENDER_CACHE_CHARS,
	RENDER_CACHE_MAX_HEAP_SHARE,
	RENDER_QUEUE_WAIT_MS,
	RENDER_WORKER_HEAP_MB,
	type RenderTask
} from "#lib/ssr/render.ts";
import type { WorkerMessage } from "#lib/ssr/worker.ts";

// leaves a core for the server
const POOL_SIZE = Math.max(1, Math.min(MAX_RENDER_WORKERS, availableParallelism() - 1));

interface Job {
	id: number;
	key: string;
	task: RenderTask;
	owner: object;
	wanted: (() => boolean)[];
	done: Promise<string | null>;
	resolve: (html: string | null) => void;
	timer: ReturnType<typeof setTimeout>;
	queueTimer: ReturnType<typeof setTimeout>;
	// 0 while queued
	started: number;
}

interface Slot {
	worker: Worker;
	job: Job | null;
}

let nextId = 0;
// one lane per page, served round robin. a big page doesn't starve a small one
const lanes = new Map<object, Job[]>();
const pending = new Map<string, Job>();
// map order is lru order. null means the render failed
const cache = new Map<string, { html: string | null; until: number }>();
let cacheChars = 0;
const slots: Slot[] = [];
// without workers the browser renders everything.
// a render on the main thread can't be stopped if it hangs
let unavailable = false;
let starting: Promise<void> | null = null;

const entryChars = (key: string, html: string | null) => key.length + (html?.length ?? 0);

const recall = (key: string): string | null | undefined => {
	const entry = cache.get(key);
	if (!entry) return undefined;
	cache.delete(key);
	if (entry.until < Date.now()) {
		cacheChars -= entryChars(key, entry.html);
		return undefined;
	}
	cache.set(key, entry);
	return entry.html;
};

// bun's getHeapStatistics walks the whole heap and its limit grows with the heap.
// there only the cache size cap bounds it
const HEAP_CHECKED = process.versions.bun === undefined;
const PRESSURE_CHECK_MS = 1_000;
let pressure = { at: -Infinity, under: false };

const underPressure = () => {
	if (!HEAP_CHECKED) return false;
	const now = Date.now();
	if (now - pressure.at < PRESSURE_CHECK_MS) return pressure.under;
	const { used_heap_size, heap_size_limit } = getHeapStatistics();
	pressure = { at: now, under: used_heap_size > heap_size_limit * RENDER_CACHE_MAX_HEAP_SHARE };
	return pressure.under;
};

const remember = (key: string, html: string | null) => {
	if (underPressure()) {
		cache.clear();
		cacheChars = 0;
		return;
	}
	const until = html === null ? Date.now() + FAILED_RENDER_TTL_MS : Infinity;
	cache.set(key, { html, until });
	cacheChars += entryChars(key, html);
	for (const [oldest, entry] of cache) {
		if (cacheChars <= RENDER_CACHE_CHARS) break;
		cache.delete(oldest);
		cacheChars -= entryChars(oldest, entry.html);
	}
};

// only cache what a worker actually ran, not queue drops
const settle = (job: Job, html: string | null, kept: boolean) => {
	clearTimeout(job.timer);
	clearTimeout(job.queueTimer);
	pending.delete(job.key);
	if (kept) remember(job.key, html);
	job.resolve(html);
};

const spawn = async (): Promise<Slot> => {
	const [{ Worker }, { workerUrl, workerInit }] = await Promise.all([
		import("node:worker_threads"),
		import("#lib/ssr/worker.ts")
	]);
	const worker = new Worker(new URL(workerUrl), {
		workerData: workerInit(),
		resourceLimits: { maxOldGenerationSizeMb: RENDER_WORKER_HEAP_MB }
	});
	worker.unref();
	const slot: Slot = { worker, job: null };
	await new Promise<void>((resolve, reject) => {
		const onMessage = (message: WorkerMessage) => {
			if (!("ready" in message)) return;
			worker.off("error", reject);
			resolve();
		};
		worker.on("message", onMessage);
		worker.once("error", reject);
	});
	worker.on("message", (message: WorkerMessage) => {
		const job = slot.job;
		if ("ready" in message || job?.id !== message.id) return;
		slot.job = null;
		settle(job, "html" in message ? message.html : null, true);
		dispatch();
	});
	// without an error listener a worker oom takes the server down. the exit handler replaces it
	worker.on("error", (cause) => console.warn("[ssr] render worker failed:", cause));
	worker.on("exit", () => replace(slot));
	return slot;
};

const replace = (slot: Slot, kept = true) => {
	const index = slots.indexOf(slot);
	if (index === -1) return;
	slots.splice(index, 1);
	if (slot.job) settle(slot.job, null, kept);
	slot.job = null;
	void slot.worker.terminate();
	void spawn().then(
		(fresh) => {
			if (unavailable) return void fresh.worker.terminate();
			slots.push(fresh);
			dispatch();
		},
		() => undefined
	);
};

// moving the lane to the end of the map makes this round robin
const take = (): Job | undefined => {
	for (const [owner, lane] of lanes) {
		const job = lane.shift();
		lanes.delete(owner);
		if (lane.length > 0) lanes.set(owner, lane);
		return job;
	}
	return undefined;
};

const dispatch = () => {
	for (const slot of slots) {
		while (!slot.job) {
			const job = take();
			if (!job) return;
			if (!job.wanted.some((wanted) => wanted())) {
				settle(job, null, false);
				continue;
			}
			slot.job = job;
			job.started = Date.now();
			clearTimeout(job.queueTimer);
			slot.worker.postMessage({ id: job.id, task: job.task });
		}
	}
};

const start = (): Promise<void> =>
	(starting ??= Promise.all(Array.from({ length: POOL_SIZE }, spawn)).then(
		(started) => {
			if (unavailable) for (const slot of started) void slot.worker.terminate();
			else slots.push(...started);
		},
		(cause: unknown) => {
			for (const slot of slots) void slot.worker.terminate();
			console.warn("[ssr] render workers unavailable, the browser renders instead:", cause);
			unavailable = true;
		}
	));

// live workers hold the event loop open and the server would never exit
process.once("sveltekit:shutdown", () => {
	unavailable = true;
	for (const lane of lanes.values()) for (const job of lane) settle(job, null, false);
	lanes.clear();
	for (const slot of slots.splice(0)) {
		if (slot.job) settle(slot.job, null, false);
		void slot.worker.terminate();
	}
});

const drop = (job: Job) => {
	const waiting = lanes.get(job.owner) ?? [];
	waiting.splice(waiting.indexOf(job), 1);
	if (waiting.length === 0) lanes.delete(job.owner);
	settle(job, null, false);
};

export const render = async (
	task: RenderTask,
	timeoutMs: number,
	owner: object = {},
	wanted: () => boolean = () => true
): Promise<string | null> => {
	await start();
	const key = createHash("sha256").update(JSON.stringify(task)).digest("hex");
	const known = recall(key);
	if (known !== undefined) return known;
	if (unavailable) return null;
	const queued = pending.get(key);
	if (queued) {
		queued.wanted.push(wanted);
		return queued.done;
	}
	let resolve!: (html: string | null) => void;
	const done = new Promise<string | null>((settled) => (resolve = settled));
	const job: Job = {
		id: nextId++,
		key,
		task,
		owner,
		wanted: [wanted],
		done,
		resolve,
		started: 0,
		timer: setTimeout(() => {
			const slot = slots.find((candidate) => candidate.job === job);
			// a timeout mostly spent queued isn't the document's fault. don't cache it
			if (slot) return replace(slot, Date.now() - job.started >= timeoutMs / 2);
			drop(job);
		}, timeoutMs),
		queueTimer: setTimeout(() => drop(job), RENDER_QUEUE_WAIT_MS)
	};
	pending.set(key, job);
	const lane = lanes.get(owner);
	if (lane) lane.push(job);
	else lanes.set(owner, [job]);
	dispatch();
	return done;
};