factorio-clusterio/packages/host/src/export.ts

// Export of item icons and locale
import fs from "node:fs/promises";
import path from "path";
import { createJimp } from "@jimp/core";
import png from "@jimp/js-png";
import * as crop from "@jimp/plugin-crop";
import * as resize from "@jimp/plugin-resize";
import type {} from "@jimp/types"; // needed to name jimp types in export.d.ts
import yazl from "yazl";

import * as lib from "@clusterio/lib";
import * as libBuildMod from "@clusterio/lib/dist/node/build_mod.js";
import { findVersion, isHeadless, type FactorioServer } from "./server.js";
import { definesPrototypesFallback } from "./export_fallback.js";

interface SimpleIconSpecification {
	icons?: never;
	icon: string;
	icon_size?: number;
	icon_mipmaps?: number;
}
interface IconLayer {
	icon: string;
	icon_size?: number;
	tint?: lib.FactorioColor,
	shift?: [number, number],
	scale?: number,
	icon_mipmaps?: number,
}
interface LayeredIconSpecification {
	icons: IconLayer[];
	icon_size?: number;
	icon_mipmaps?: number;
}
type IconSpecification = LayeredIconSpecification | SimpleIconSpecification;
type PrototypeWithIcon = lib.FactorioPrototypeBase & IconSpecification;


/**
 * Generate the export mod needed for exportData
 *
 * Packs the lua export mod located in lua/export into the mods folder for
 * the server given, with dependencies generated for all the mods present in
 * the server's mods folder.
 *
 * @param server -
 *     The server to generate the export mod for.
 */
async function generateExportMod(server: FactorioServer) {
	let dependencies = [];
	let splitter = /^(.*)_(\d+\.\d+\.\d+)(\.zip)?$/;
	for (let entry of await fs.readdir(server.writePath("mods"))) {
		let match = splitter.exec(entry);
		if (match && match[1] !== "export") {
			dependencies.push(`? ${match[1]}`);
		}
	}

	await libBuildMod.build({
		clean: false,
		build: true,
		pack: true,
		sourceDir: path.join(import.meta.dirname, "..", "..", "..", "lua", "export"),
		outputDir: server.writePath("mods"),
		bumpPatch: false,
		factorioVersion: server.version?.replace(/\.\d+$/, ""),
		dependencies,
	});
}

interface ModZip {
	zip: lib.ZipArchive;
	root: string;
}
type ZipCache = Map<string, ModZip>;
async function loadZip(server: FactorioServer, modVersions: Map<string, string>, zipCache: ZipCache, mod: string) {
	let modVersion = modVersions.get(mod);
	if (!modVersion) {
		throw new Error(`Got path for unknown mod ${mod}`);
	}

	let zipPath = server.writePath("mods", `${mod}_${modVersion}.zip`);
	let modZip = zipCache.get(zipPath);
	if (!modZip) {
		const zip = await lib.ZipArchive.fromFile(zipPath);
		try {
			modZip = { zip, root: lib.findRoot(zip) };
		} catch (err) {
			zip.close();
			throw err;
		}
		zipCache.set(zipPath, modZip);
	}

	return modZip;
}

/** Where icon files are read from during an export. */
interface ExportFiles {
	/** Mod zips opened during the export, closed when it ends. */
	zipCache: ZipCache;
	/** Data dir of a full install to fall back to for builtin mod files. */
	graphicsDataDir?: string;
	/** Why no full install was found for a headless server. */
	graphicsError?: string;
}

function closeZipCache(zipCache: ZipCache) {
	for (const { zip } of zipCache.values()) {
		zip.close();
	}
	zipCache.clear();
}

/**
 * Load the given Factorio file path into a Buffer
 *
 * @param server - The server to load the file from.
 * @param modVersions - Mapping of mod to version used.
 * @param files - Where to read files from.
 * @param modPath - Factorio style path to the file to load.
 * @returns The content of the file or null if not found.
 * @internal
 */
async function loadFile(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	modPath: string,
) {
	let match = /^__([^\/]+)__\/(.*)$/.exec(modPath);
	if (!match) {
		server._logger.warn(`Skipping icon with bad mod path: ${JSON.stringify(modPath)}`);
		return null;
	}

	let [, mod, filePath] = match;

	const builtinModNames = ["core", ...lib.ModPack.getBuiltinModNames(server.version!)];
	if (builtinModNames.includes(mod)) {
		const dataDirs = [server.dataPath()];
		if (files.graphicsDataDir) {
			dataDirs.push(files.graphicsDataDir);
		}
		for (const dataDir of dataDirs) {
			try {
				return await fs.readFile(path.join(dataDir, mod, filePath));
			} catch (err: any) {
				if (err.code !== "ENOENT") {
					throw err;
				}
			}
		}
		return null;
	}

	let modZip;
	try {
		modZip = await loadZip(server, modVersions, files.zipCache, mod);
	} catch (err: any) {
		if (err.code === "ENOENT") {
			return null;
		}
		throw err;
	}

	return await modZip.zip.readFile(path.posix.join(modZip.root, filePath));
}

// Only the parts of jimp the export uses
const Jimp = createJimp({ formats: [png], plugins: [crop.methods, resize.methods] });
type Image = InstanceType<typeof Jimp>;
type IconCache = Map<string, Image | null>;
async function loadIcon(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	iconPath: string,
	iconSize: number,
	iconMipmaps: number,
	iconCache: IconCache,
) {
	let icon = iconCache.get(iconPath);
	if (icon === undefined) {
		let fileContent = await loadFile(server, modVersions, files, iconPath);
		if (fileContent) {
			icon = await Jimp.fromBuffer(fileContent) as Image;
			icon.crop({ x: 0, y: 0, w: iconSize, h: iconSize });
			iconCache.set(iconPath, icon);
		} else {
			icon = null;
		}
		iconCache.set(iconPath, icon);
	}
	return icon;
}

async function loadSimpleIcon(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	item: SimpleIconSpecification,
	size: number,
	iconCache: IconCache,
): Promise<Image | null> {
	let icon = await loadIcon(server, modVersions, files, item.icon, item.icon_size ?? 64, 0, iconCache);
	if (icon) {
		let iconScale = size / (item.icon_size ?? 64);
		if (iconScale !== 1) {
			icon = icon.clone();
			icon.scale(iconScale);
		}
	}
	return icon;
}

async function loadLayeredIcon(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	item: LayeredIconSpecification,
	size: number,
	iconCache: IconCache,
): Promise<Image | null> {
	let baseLayerSize = (item.icons[0].icon_size || item.icon_size) ?? 64;
	let icon = new Jimp({ width: size, height: size });

	// The scaling factor of the base layer
	let baseLayerScale = item.icons[0].scale || 32 / baseLayerSize;

	// The size in pixels of one unit
	let baseUnit = size / (baseLayerSize * baseLayerScale);

	for (let layer of item.icons) {
		let layerSize = (layer.icon_size || item.icon_size) ?? 64;
		let iconLayer = await loadIcon(
			server, modVersions, files, layer.icon, layerSize, layer.icon_mipmaps || 0, iconCache,
		);

		// A partial icon would be wrong, skip it like a missing simple icon
		if (!iconLayer) {
			return null;
		}

		iconLayer = iconLayer.clone();

		let tint: lib.Color;
		if (layer.tint) {
			tint = lib.normalizeColor(layer.tint);
		} else {
			tint = { r: 1, b: 1, g: 1, a: 1 };
		}

		let layerScale = layer.scale || 32 / layerSize;
		let [xs, ys] = layer.shift || [0, 0];
		let realScale = layerScale * baseUnit;
		let sizeShift = (size - layerSize * realScale) / 2;

		xs = xs * baseUnit + sizeShift;
		ys = ys * baseUnit + sizeShift;

		if (realScale !== 1) {
			iconLayer.scale(realScale);
		}


		iconLayer.scan(0, 0, iconLayer.bitmap.width, iconLayer.bitmap.height, (x, y, sidx) => {
			x += xs;
			y += ys;
			if (x < 0 || x >= size || y < 0 || y >= size) {
				return;
			}
			let sdata = iconLayer!.bitmap.data;
			let ddata = icon.bitmap.data;
			let didx = icon.getPixelIndex(x, y);

			let sa = sdata[sidx + 3] / 255;
			let da = ddata[didx + 3] / 255;
			let db = da * (1 - sa * tint.a);
			let cb = sa + da * (1 - sa);
			ddata[didx + 0] = Math.min(255, (sdata[sidx + 0] * sa * tint.r + ddata[didx + 0] * db) / cb);
			ddata[didx + 1] = Math.min(255, (sdata[sidx + 1] * sa * tint.g + ddata[didx + 1] * db) / cb);
			ddata[didx + 2] = Math.min(255, (sdata[sidx + 2] * sa * tint.b + ddata[didx + 2] * db) / cb);
			ddata[didx + 3] = Math.min(255, sdata[sidx + 3] + ddata[didx + 3] * (1 - sa));
		});
	}

	return icon;
}

function fixIcons(item: PrototypeWithIcon) {
	const icons = item.icons;
	if (typeof icons === "object" && !(icons instanceof Array) && icons !== null) {
		// It's possible to specify icons as an object with arbitrary keys
		// and the game will still accept it, cast the value to array if
		// this is the case.
		item = { ...item };
		item.icons = Object.values(icons);
	}
	return item;
}

const prototypesWithIcon = [
	"item",
	"fluid",
	"recipe",
	"virtual-signal",
	"technology",
	"space-location",
	"quality",
	"entity",
] as const;

function filterAndFixPrototypesWithIcons(prototypes: lib.FactorioPrototypes, allowedTypes: Set<string>) {
	return Object.entries(prototypes)
		.filter(([type]) => allowedTypes.has(type))
		.flatMap(([_, typePrototypes]) => Object.values(typePrototypes) as PrototypeWithIcon[])
		.map(fixIcons);
}

/** Shared state for building a unified spritesheet across all categories. */
interface SheetState {
	sheet: Image;
	metadata: Record<string, lib.ExportMetadataEntry[]>;
	iconCache: IconCache;
	simpleIcons: Map<string, number>;
	pos: number;
}

const SHEET_ICON_SIZE = 32;
const SHEET_WIDTH = 1024;

async function createSheetState(): Promise<SheetState> {
	return {
		sheet: new Jimp({ width: SHEET_WIDTH, height: SHEET_ICON_SIZE }),
		metadata: Object.create(null),
		iconCache: new Map(),
		simpleIcons: new Map(),
		pos: 0,
	};
}

/** Grow the spritesheet vertically if needed to fit another row. */
async function growSheet(state: SheetState) {
	const row = Math.floor(state.pos / (SHEET_WIDTH / SHEET_ICON_SIZE));
	const neededHeight = (row + 1) * SHEET_ICON_SIZE;
	if (neededHeight > state.sheet.bitmap.height) {
		const extended = new Jimp({ width: SHEET_WIDTH, height: neededHeight });
		extended.composite(state.sheet, 0, 0);
		(state.sheet as any).bitmap = (extended as any).bitmap;
	}
}

async function exportIcon(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	state: SheetState,
	prototype: PrototypeWithIcon,
) {
	let icon: Image | null = null;
	let iconPos: number | undefined;
	let iconData: lib.ExportMetadataEntry["icon"];
	let iconPath: string | undefined;

	if (prototype.icons) {
		icon = await loadLayeredIcon(
			server, modVersions, files, prototype, SHEET_ICON_SIZE, state.iconCache,
		);
		if (icon) {
			iconPos = state.pos;
		}
		iconPath = prototype.icons[0].icon;

	} else if (prototype.icon) {
		iconPos = state.simpleIcons.get(prototype.icon);
		iconPath = prototype.icon;
		if (iconPos === undefined) {
			icon = await loadSimpleIcon(
				server, modVersions, files, prototype, SHEET_ICON_SIZE, state.iconCache,
			);
			if (icon) {
				iconPos = state.pos;
				state.simpleIcons.set(prototype.icon, state.pos);
			}
		}
	}

	if (iconPos !== undefined) {
		iconData = {
			x: iconPos * SHEET_ICON_SIZE % SHEET_WIDTH,
			y: Math.floor(iconPos / (SHEET_WIDTH / SHEET_ICON_SIZE)) * SHEET_ICON_SIZE,
			size: SHEET_ICON_SIZE,
			path: iconPath,
		};
	}

	if (icon) {
		await growSheet(state);
		state.sheet.composite(
			icon,
			state.pos * SHEET_ICON_SIZE % SHEET_WIDTH,
			Math.floor(state.pos / (SHEET_WIDTH / SHEET_ICON_SIZE)) * SHEET_ICON_SIZE,
		);
		state.pos += 1;
	}

	return iconData;
}

/**
 * Export prototype icons for a single base prototype onto the shared spritesheet.
 *
 * @param server - The server to load icons from.
 * @param modVersions - Mapping of mod name to version.
 * @param files - Where to read icon files from.
 * @param prototypes - Array of prototypes with icon to export for this base type.
 * @param baseName - name of base type (used in metadata).
 * @param state - Shared spritesheet state (mutated in place).
 * @returns Number of icons exported for this base type.
 * @internal
 */
async function exportIcons(
	server: FactorioServer,
	modVersions: Map<string, string>,
	files: ExportFiles,
	prototypes: PrototypeWithIcon[],
	baseName: string,
	state: SheetState,
) {
	const metadata = state.metadata[baseName] ?? [];
	state.metadata[baseName] ??= metadata;
	let count = 0;
	for (let prototype of prototypes) {
		let icon: lib.ExportMetadataEntry["icon"];
		try {
			icon = await exportIcon(server, modVersions, files, state, prototype);
		} catch (err: any) {
			// One broken icon should not stop the rest of the export
			server._logger.warn(`Skipping icon of ${prototype.type} ${prototype.name}: ${err.message}`);
		}
		if (icon) {
			count += 1;
		}
		metadata.push({
			name: prototype.name,
			type: prototype.type,
			localised_name: prototype.localised_name,
			icon,
		});
	}

	return count;
}

/**
 * Exports the locale files for the base game and the given mods
 *
 * Parses and merges all the locales for the all the mods given through
 * `modVersions` and `modOrder`.
 *
 * @param server - The server to export the locale from.
 * @param modVersions - Mapping of mod name to version to export locale from.
 * @param zipCache - Mod zips opened by this export.
 * @param modOrder - Load order of the mods.
 * @param languageCode - Language to export locale for.
 * @returns merged locale information
 * @internal
 */
async function exportLocale(
	server: FactorioServer,
	modVersions: Map<string, string>,
	zipCache: ZipCache,
	modOrder: string[],
	languageCode: string
) {
	let mergedLocales = new Map<string, string>();

	function mergeLocale(locale: Record<string, string | Record<string, string>>) {
		for (let [category, entries] of Object.entries(locale)) {
			if (typeof entries === "string") {
				mergedLocales.set(category, entries);

			} else {
				for (let [key, value] of Object.entries(entries)) {
					mergedLocales.set(`${category}.${key}`, value);
				}
			}
		}
	}

	const builtinModNames = [
		// Filter if on builtin mods so "core" does not need to be specified in modOrder
		"core", ...lib.ModPack.getBuiltinModNames(server.version!).filter(mod => modOrder.includes(mod)),
	];
	for (const builtinModName of builtinModNames) {
		const localeFilePath = server.dataPath(builtinModName, "locale", languageCode, `${builtinModName}.cfg`);
		mergeLocale(lib.parse(await fs.readFile(localeFilePath, "utf8")));
	}

	const builtinModNamesExport = ["export", ...builtinModNames];
	for (let mod of modOrder) {
		if (builtinModNamesExport.includes(mod)) {
			continue;
		}

		let modZip;
		try {
			modZip = await loadZip(server, modVersions, zipCache, mod);
		} catch (err: any) {
			if (err.code === "ENOENT") {
				server._logger.warn(`Skipping locale of ${mod}, its mod zip was not found`);
				continue;
			}
			throw err;
		}
		const localePath = `${path.posix.join(modZip.root, "locale", languageCode)}/`;
		for (const entry of modZip.zip.entries.values()) {
			if (entry.fileName.startsWith(localePath) && entry.fileName.endsWith(".cfg")) {
				let content = await modZip.zip.readEntry(entry);
				mergeLocale(lib.parse(content.toString("utf8")));
			}
		}
	}

	return mergedLocales;
}

/**
 * Export locale, icons and metadata for the given factorio server
 *
 * If the server runs the headless build, builtin mod graphics are read
 * from a full install of the same major and minor version found next to
 * it, if there is one.
 *
 * @param server - The server to export the data from.
 * @returns zip file with exported data.
 */
export async function exportData(server: FactorioServer) {
	const files: ExportFiles = { zipCache: new Map() };
	if (await isHeadless(server.dataPath())) {
		const majorMinor = server.version!.replace(/\.\d+$/, "") as lib.MajorMinorVersion;
		try {
			[files.graphicsDataDir] = await findVersion(server._factorioDir, majorMinor, false);
		} catch (err: any) {
			// Without it the export still has everything but the base game icons
			files.graphicsError = err.message;
		}
	}
	try {
		return await buildExport(server, files);
	} finally {
		// Free up the zip files loaded during the export.
		closeZipCache(files.zipCache);
	}
}

/**
 * Warn if any builtin mod icon files could not be found
 *
 * @param server - The server the export ran on.
 * @param files - Where icon files were read from.
 * @param iconCache - Icons loaded by the export, null for missing files.
 * @internal
 */
function checkMissingBuiltinIcons(server: FactorioServer, files: ExportFiles, iconCache: IconCache) {
	const builtinModNames = ["core", ...lib.ModPack.getBuiltinModNames(server.version!)];
	let missing = 0;
	for (const [iconPath, icon] of iconCache) {
		const mod = /^__([^\/]+)__\//.exec(iconPath)?.[1];
		if (icon === null && mod && builtinModNames.includes(mod)) {
			missing += 1;
		}
	}
	if (!missing) {
		return;
	}

	const dirs = [server.dataPath(), files.graphicsDataDir].filter(Boolean).join(" or ");
	let message =
		`${missing} icon files of the base game were not found in ${dirs}, icons using them were not exported.`;
	if (files.graphicsError !== undefined) {
		const majorMinor = server.version!.replace(/\.\d+$/, "");
		message +=
			` The headless build of Factorio has no graphics and no full install was found (${files.graphicsError}).` +
			` Add a full install of Factorio ${majorMinor} to host.factorio_directory to read them from there,` +
			" or export the data on a host with a full install.";
	}
	server._logger.warn(message);
}

async function buildExport(server: FactorioServer, files: ExportFiles) {
	await generateExportMod(server);

	let settings: lib.ExportSettings = {};
	let prototypes: lib.ExportPrototypes = {};
	let modVersions = new Map();
	let modOrder: string[] = [];
	let defines: lib.FactorioDefines | undefined;

	function add(obj: lib.FactorioPrototypes, prototype: lib.FactorioPrototypeBase) {
		if (!Object.prototype.hasOwnProperty.call(obj, prototype.type)) {
			obj[prototype.type] = {};
		}
		obj[prototype.type][prototype.name] = prototype;
	}

	server.on("ipc-prototype_export", data => add(prototypes, data));
	server.on("ipc-settings_export", data => add(settings, data));
	server.on("ipc-mod_setting_mod", ({ name, mod }) => {
		for (let type of Object.values(settings)) {
			if (Object.prototype.hasOwnProperty.call(type, name)) {
				type[name].mod = mod;
				return;
			}
		}
		server._logger.error(`Unable to find ${name} in settings prototypes`);
	});
	server.on("ipc-mod_list", data => { modVersions = new Map(Object.entries(data)); });
	server.on("ipc-defines", data => { defines = data; });
	server.on("output", parsed => {
		if (parsed.format === "seconds" && parsed.type === "generic") {
			let match = /^Checksum of (.*): \d+$/.exec(parsed.message);
			if (match) {
				modOrder.push(match[1]);
			}
		}
	});

	try {
		await server.startScenario("base/freeplay");
		await server.kill();
	} finally {
		await fs.unlink(server.writePath("mods", "export_0.0.0.zip"));
	}

	if (!Object.keys(prototypes).length) {
		throw new Error("No prototypes got exported");
	}

	if (!defines) {
		throw new Error("No defines got exported");
	}

	// compat with versions prior to 1.1.31
	if (!defines.prototypes) {
		defines.prototypes = definesPrototypesFallback;
	}

	// Some mod authors put leading zeros into the versions of their zip files.
	let splitter = /^(.*)_(\d+)\.(\d+)\.(\d+)\.zip?$/;
	for (let entry of await fs.readdir(server.writePath("mods"))) {
		let match = splitter.exec(entry);
		if (!match) {
			continue;
		}

		let modVersion = `${match[2]}.${match[3]}.${match[4]}`;
		let normalizedVersion =
			`${Number.parseInt(match[2], 10)}.${Number.parseInt(match[3], 10)}.${Number.parseInt(match[4], 10)}`
		;

		if (modVersion === normalizedVersion) {
			continue;
		}

		if (modVersions.get(match[1]) === normalizedVersion) {
			modVersions.set(match[1], modVersion);
		}
	}

	let locale = await exportLocale(server, modVersions, files.zipCache, modOrder, "en");

	let zip = new yazl.ZipFile();
	const addJson = (name: string, value: unknown) => zip.addBuffer(Buffer.from(JSON.stringify(value)), name);
	addJson("export/settings.json", settings satisfies lib.ExportSettings);
	addJson("export/prototypes.json", prototypes satisfies lib.ExportPrototypes);
	addJson("export/locale.json", [...locale.entries()] satisfies lib.ExportLocale);
	addJson("export/defines.json", defines satisfies lib.ExportDefines);

	// Build a single unified spritesheet across all prototypes with icons.
	const state = await createSheetState();
	for (const baseName of prototypesWithIcon) {
		if (!Object.prototype.hasOwnProperty.call(defines.prototypes, baseName)) {
			continue;
		}
		const subTypeNames = new Set(Object.keys(defines.prototypes[baseName]));
		const prototypesWithIcons = filterAndFixPrototypesWithIcons(prototypes, subTypeNames);

		if (prototypesWithIcons.length === 0) {
			continue;
		}
		const count = await exportIcons(server, modVersions, files, prototypesWithIcons, baseName, state);
		server._logger.info(`Exported ${count} ${baseName} icons`);
	}
	checkMissingBuiltinIcons(server, files, state.iconCache);

	// Crop sheet to actual used height and write single spritesheet + metadata.
	const usedRows = Math.ceil(state.pos / (SHEET_WIDTH / SHEET_ICON_SIZE));
	state.sheet.crop({ x: 0, y: 0, w: SHEET_WIDTH, h: Math.max(usedRows * SHEET_ICON_SIZE, 1) });
	zip.addBuffer(await state.sheet.getBuffer("image/png"), "export/spritesheet.png");
	addJson("export/metadata.json", state.metadata satisfies lib.ExportMetadata);
	zip.end();

	server._logger.info(`Export complete: ${state.pos} icons on ${usedRows} row(s)`);

	return zip;
}


// For testing only
export const _createSheetState = createSheetState;
export const _exportIcons = exportIcons;
export const _exportLocale = exportLocale;
export const _loadSimpleIcon = loadSimpleIcon;
export const _loadLayeredIcon = loadLayeredIcon;
export const _checkMissingBuiltinIcons = checkMissingBuiltinIcons;