import { statSync } from "node:fs";
import { readFile, writeFile } from "node:fs/promises";
import { resolve } from "node:path";
import { pathToFileURL } from "node:url";

export const INITIAL_BUNDLE_BUDGET_BYTES = 500 * 1024;

function manifestEntries(manifest) {
	if (!manifest || typeof manifest !== "object" || Array.isArray(manifest)) {
		throw new Error("Vite manifest must be an object");
	}
	return Object.entries(manifest);
}

function assetFile(entry, entryName) {
	if (!entry || typeof entry.file !== "string" || !entry.file.endsWith(".js")) {
		throw new Error(
			`Manifest entry '${entryName}' must reference a JavaScript file`,
		);
	}
	return entry.file;
}

export function getInitialAssetFiles(manifest) {
	const entries = manifestEntries(manifest);
	const entryCandidates = entries.filter(
		([, entry]) => entry?.isEntry === true,
	);
	if (entryCandidates.length !== 1) {
		throw new Error("Vite manifest must contain exactly one JavaScript entry");
	}

	const files = [];
	const visitedEntries = new Set();
	const visit = (entryName) => {
		if (visitedEntries.has(entryName)) return;
		visitedEntries.add(entryName);
		const entry = manifest[entryName];
		if (!entry)
			throw new Error(`Vite manifest is missing eager import '${entryName}'`);
		files.push(assetFile(entry, entryName));
		for (const importedEntry of entry.imports ?? []) visit(importedEntry);
	};

	visit(entryCandidates[0][0]);
	return files;
}

export function measureAssets(files, getSize) {
	return files.reduce((total, file) => {
		const size = getSize(file);
		if (!Number.isSafeInteger(size) || size < 0) {
			throw new Error(`Cannot measure bundle asset '${file}'`);
		}
		return total + size;
	}, 0);
}

export function assertBundleBudget(initialBytes) {
	if (initialBytes > INITIAL_BUNDLE_BUDGET_BYTES) {
		throw new Error(
			`Initial JavaScript payload is ${initialBytes} bytes; budget is ${INITIAL_BUNDLE_BUDGET_BYTES} bytes`,
		);
	}
}

export function getLazyRouteAssetSizes(manifest, getSize, baseline) {
	const routeEntries = manifestEntries(manifest).filter(
		([, entry]) =>
			entry?.isDynamicEntry === true && typeof entry.src === "string",
	);
	if (routeEntries.length === 0) {
		throw new Error("Vite manifest must contain lazy route chunks");
	}

	const routeSizes = Object.fromEntries(
		routeEntries
			.map(([entryName, entry]) => [
				entry.src,
				getSize(assetFile(entry, entryName)),
			])
			.sort(([left], [right]) => left.localeCompare(right)),
	);
	if (!baseline) return routeSizes;

	for (const [route, size] of Object.entries(routeSizes)) {
		if (!Number.isSafeInteger(baseline[route])) {
			throw new Error(`Missing baseline for lazy route chunk '${route}'`);
		}
		if (size > baseline[route]) {
			throw new Error(
				`Lazy route chunk '${route}' is ${size} bytes; baseline is ${baseline[route]} bytes`,
			);
		}
	}
	return routeSizes;
}

const distributionDirectory = resolve("dist");
const manifestPath = resolve(distributionDirectory, "bundle-manifest.json");
const baselinePath = resolve("scripts/lazy-route-bundle-baseline.json");

async function loadJson(path) {
	return JSON.parse(await readFile(path, "utf8"));
}

function getAssetSize(file) {
	return statSync(resolve(distributionDirectory, file)).size;
}

async function run(mode) {
	const manifest = await loadJson(manifestPath);
	const initialBytes = measureAssets(
		getInitialAssetFiles(manifest),
		getAssetSize,
	);
	assertBundleBudget(initialBytes);
	const baseline =
		mode === "baseline" ? undefined : await loadJson(baselinePath);
	const lazyRouteSizes = getLazyRouteAssetSizes(
		manifest,
		getAssetSize,
		baseline,
	);

	if (mode === "baseline") {
		await writeFile(
			baselinePath,
			`${JSON.stringify(lazyRouteSizes, null, "\t")}\n`,
		);
	}
	console.log(`Initial JavaScript payload: ${initialBytes} bytes`);
}

if (
	process.argv[1] &&
	import.meta.url === pathToFileURL(process.argv[1]).href
) {
	const mode = process.argv[2] ?? "check";
	if (mode !== "check" && mode !== "baseline") {
		throw new Error("Usage: node scripts/bundle-budget.mjs [check|baseline]");
	}
	await run(mode);
}