// Library for patching Factorio saves with scenario code.

import fs from "node:fs/promises";
import path from "path";
import semver from "semver";
import { pipeline } from "node:stream/promises";
import { Type, Static } from "@sinclair/typebox";

import * as lib from "@clusterio/lib";

/**
 * Describes a module that can be patched into a save
 */
export class SaveModule {
	constructor(
		/** Module */
		public info: lib.ModuleInfo,
		/** Files and their content that will be patched into the save */
		public files = new Map<string, Buffer>(),
	) { }

	static moduleFilePath(filePath: string, moduleName: string) {
		// Map locale files to the save's locale folder
		if (filePath.startsWith("locale/")) {
			const slashes = filePath.match(/\//g)!.length;
			if (slashes > 1) {
				const secondSlash = filePath.indexOf("/", "locale/".length);
				return `${filePath.slice(0, secondSlash)}/${moduleName}-${filePath.slice(secondSlash + 1)}`;
			}
			const period = `${filePath}.`.indexOf(".");
			return `${filePath.slice(0, period)}/${moduleName}${filePath.slice(period)}`;
		}
		// Map all other files into modules/name in the save.
		return path.posix.join("modules", moduleName, filePath);
	}

	async loadFiles(moduleDirectory: string) {
		let dirs: [string, string][] = [[moduleDirectory, ""]];
		while (dirs.length) {
			let [dir, relativeDir] = dirs.pop()!;
			for (let entry of await fs.readdir(dir, { withFileTypes: true })) {
				let fsPath = path.join(dir, entry.name);
				let relativePath = path.posix.join(relativeDir, entry.name);

				if (entry.isFile()) {
					let savePath = SaveModule.moduleFilePath(relativePath, this.info.name);
					this.files.set(savePath, await fs.readFile(fsPath));
					if (relativePath === "module_exports.lua") {
						this.files.set(
							`modules/${this.info.name}.lua`,
							Buffer.from(`return require("modules/${this.info.name}/module_exports")`, "utf-8")
						);
					}

				} else if (entry.isDirectory()) {
					dirs.push([fsPath, relativePath]);
				}
			}
		}
	}

	static jsonSchema = Type.Object({
		...lib.ModuleInfo.jsonSchema.properties,
		"files": Type.Array(Type.String()),
	});

	toJSON() {
		return {
			...this.info.toJSON(),
			files: [...this.files.keys()],
		};
	}

	static async fromSave(json: Static<typeof SaveModule.jsonSchema>, zip: lib.ZipArchive, root: string) {
		const module = new this(lib.ModuleInfo.fromJSON(json));
		for (const filename of json.files) {
			const content = await zip.readFile(path.posix.join(root, filename));
			if (content === null) {
				lib.logger.warn(`Missing file ${filename} in save`);
				continue;
			}
			module.files.set(filename, content);
		}
		return module;
	}

	static async fromPluginInfo(pluginInfo: lib.PluginNodeEnvInfo) {
		let moduleDirectory = path.join(path.dirname(pluginInfo.packagePath), "module");
		try {
			await fs.access(moduleDirectory);
		} catch (err: any) {
			return null;
		}

		let moduleJsonPath = path.join(moduleDirectory, "module.json");
		let moduleJson;
		try {
			moduleJson = {
				name: pluginInfo.name,
				version: pluginInfo.version,
				dependencies: { "clusterio": "*" },
				...JSON.parse(await fs.readFile(moduleJsonPath, "utf8")),
			};
		} catch (err: any) {
			throw new Error(`Loading module/module.json in plugin ${pluginInfo.name} failed: ${err.message}`);
		}
		if (!lib.ModuleInfo.validate(moduleJson)) {
			throw new Error(
				`module/module.json in plugin ${pluginInfo.name} failed validation:\n` +
				`${JSON.stringify(lib.ModuleInfo.validate.errors, null, "\t")}`
			);
		}

		let module = new SaveModule(lib.ModuleInfo.fromJSON(moduleJson));
		await module.loadFiles(moduleDirectory);
		return module;
	}

	static async fromDirectory(moduleDirectory: string) {
		let name = path.basename(moduleDirectory);
		let moduleJsonPath = path.join(moduleDirectory, "module.json");
		let moduleJson;
		try {
			moduleJson = JSON.parse(await fs.readFile(moduleJsonPath, "utf8"));
		} catch (err: any) {
			throw new Error(`Loading ${name}/module.json failed: ${err.message}`);
		}
		if (!lib.ModuleInfo.validate(moduleJson)) {
			throw new Error(
				`${name}/module.json failed validation:\n` +
				`${JSON.stringify(lib.ModuleInfo.validate.errors, null, "\t")}`
			);
		}

		if (moduleJson.name !== name) {
			throw new Error(`Expected name of module ${moduleJson.name} to match the directory name ${name}`);
		}

		let module = new SaveModule(lib.ModuleInfo.fromJSON(moduleJson));
		await module.loadFiles(moduleDirectory);
		return module;
	}
}

/**
 * Load the modules to patch into a save
 *
 * Collects the modules provided by the given plugins together with the
 * modules bundled with the host package.
 *
 * @param pluginInfos - Plugins to load modules from.
 * @returns Map of module name to module.
 */
export async function loadModules(pluginInfos: Iterable<lib.PluginNodeEnvInfo>) {
	let modules: Map<string, SaveModule> = new Map();
	for (let pluginInfo of pluginInfos) {
		let module = await SaveModule.fromPluginInfo(pluginInfo);
		if (!module) {
			continue;
		}
		modules.set(module.info.name, module);
	}

	// Find stand alone modules to load
	// XXX for now only the included clusterio module is loaded
	let modulesDirectory = path.join(import.meta.dirname, "..", "..", "..", "modules");
	for (let entry of await fs.readdir(modulesDirectory, { withFileTypes: true })) {
		if (entry.isDirectory()) {
			if (modules.has(entry.name)) {
				throw new Error(`Module with name ${entry.name} already exists in a plugin`);
			}
			let module = await SaveModule.fromDirectory(path.join(modulesDirectory, entry.name));
			modules.set(module.info.name, module);
		}
	}
	return modules;
}

export class PatchInfo {
	static currentVersion = 1;

	constructor(
		public patchNumber: number,
		public scenario: lib.ModuleInfo,
		public modules: SaveModule[],
		public version = PatchInfo.currentVersion,
	) { }

	static jsonSchema = Type.Object({
		"version": Type.Number(),
		"patch_number": Type.Number(),
		"scenario": lib.ModuleInfo.jsonSchema,
		"modules": Type.Array(SaveModule.jsonSchema),
	});

	toJSON() {
		return {
			version: PatchInfo.currentVersion,
			patch_number: this.patchNumber,
			scenario: this.scenario.toJSON(),
			modules: this.modules.map(m => m.toJSON()),
		};
	}

	static async fromSave(json: Static<typeof PatchInfo.jsonSchema>, zip: lib.ZipArchive, root: string) {
		if (json.version === undefined) {
			interface ScenarioInfoV0 {
				name: string;
				modules: string[];
			}
			interface PatchInfoV0 {
				patch_number: number;
				scenario: ScenarioInfoV0;
				modules: { name: string, files: { path: string, load: boolean, require: boolean }[], }[];
			}

			const info = json as unknown as PatchInfoV0;
			return new this(
				info.patch_number,
				new lib.ModuleInfo(info.scenario.name, "0.0.0", info.scenario.modules),
				[],
				0,
			);
		}

		return new this(
			json.patch_number,
			lib.ModuleInfo.fromJSON(json.scenario),
			await Promise.all(json.modules.map(m => SaveModule.fromSave(m, zip, root))),
			json.version,
		);
	}
}

/**
 * Generates control.lua code for loading the Clusterio modules
 *
 * @param patchInfo - The patch info files's json content
 * @returns Generated control.lua code.
 * @internal
 */
function generateLoader(patchInfo: PatchInfo) {
	let lines = [
		"-- Auto generated scenario module loader created by Clusterio",
		"-- Modifications to this file will be lost when loaded in Clusterio",
		`clusterio_patch_number = ${patchInfo.patchNumber}`,
		"",
		'local event_handler = require("event_handler")',
		"",
		"-- Scenario modules",
	];

	for (let requirePath of patchInfo.scenario.load) {
		lines.push(`event_handler.add_lib(require("${requirePath}"))`);
	}
	for (let requirePath of patchInfo.scenario.require) {
		lines.push(`require("${requirePath}")`);
	}

	lines.push(...[
		"",
		"-- Clusterio modules",
	]);

	for (let module of patchInfo.modules) {
		for (let file of module.info.load) {
			let requirePath = `modules/${module.info.name}/${file.replace(/\.lua$/i, "")}`;
			lines.push(`event_handler.add_lib(require("${requirePath}"))`);
		}
		for (let file of module.info.require) {
			let requirePath = `modules/${module.info.name}/${file.replace(/\.lua$/i, "")}`;
			lines.push(`require("${requirePath}")`);
		}
	}

	// End last line with a newline
	lines.push("");

	return lines.join("\n");
}

/**
 * Reorders modules to satisfy their dependencies
 *
 * Looks through and reorders the array of module definitions in order to
 * satisfy the property that dependencies are earlier in the array than
 * their dependents.  Throws an error if this is not possible.
 *
 * @param modules - Array of modules to reorder
 * @internal
 */
function reorderDependencies(modules: SaveModule[]) {
	let index = 0;
	let present = new Map<string, string>();
	let hold = new Map<string, [SaveModule]>();
	reorder: while (index < modules.length) {
		let module = modules[index];
		if (semver.valid(module.info.version) === null) {
			throw new Error(`Invalid version '${module.info.version}' for module ${module.info.name}`);
		}

		for (let [dependency, requirement] of module.info.dependencies) {
			if (semver.validRange(requirement) === null) {
				throw new Error(
					`Invalid version range '${requirement}' for dependency ${dependency} on module ${module.info.name}`
				);
			}

			if (present.has(dependency)) {
				if (!semver.satisfies(present.get(dependency)!, requirement)) {
					throw new Error(`Module ${module.info.name} requires ${dependency} ${requirement}`);
				}

			// We have an unmet dependency, take it out and continue
			} else {
				if (hold.has(dependency)) {
					hold.get(dependency)!.push(module);
				} else {
					hold.set(dependency, [module]);
				}
				modules.splice(index, 1);
				continue reorder;
			}
		}

		// No unmet dependencies, record and continue
		present.set(module.info.name, module.info.version);
		index += 1;

		if (hold.has(module.info.name)) {
			modules.splice(index, 0, ...hold.get(module.info.name)!);
			hold.delete(module.info.name);
		}
	}

	if (!hold.size) {
		return;
	}

	// There are three reasons for a module to end up being held: The module depends
	// on a module that is missing, the module is part of a dependency loop, or the
	// the module depends on a module that satisfy any of these conditions.

	let remaining = new Map<string, SaveModule>();
	for (let heldModules of hold.values()) {
		for (let module of heldModules) {
			remaining.set(module.info.name, module);
		}
	}

	// Start with a random module from the remaining modules
	for (let module of remaining.values()) {
		let cycle: string[] = [];
		while (true) {
			// Find an unmet dependency
			let dependency = [...module.info.dependencies.keys()].find(name => !present.has(name))!;

			if (!remaining.has(dependency)) {
				// There's no module being held up by this dependency, the
				// dependency is missing.
				throw new Error(`Missing dependency ${dependency} for module ${module.info.name}`);
			}

			if (cycle.includes(module.info.name)) {
				cycle.push(module.info.name);
				cycle.splice(0, cycle.indexOf(module.info.name));
				throw new Error(`Module dependency loop detected: ${cycle.join(" -> ")}`);
			}

			cycle.push(module.info.name);
			module = remaining.get(dependency)!;
		}
	}
}

const knownScenarios: Record<string, lib.ModuleInfo> = {
	// First seen in 0.17.63
	"4e866186ebe297f1038fd325b09df1a1f5e2fdd1": new lib.ModuleInfo("freeplay", "0.17.63", [], ["scenario"]),
	// First seen in 2.0
	"bcbdde18ce4ec16ebfd93bd694cd9b12ef969c9a": new lib.ModuleInfo("freeplay", "2.0.0", [], ["scenario"]),
	// First seen in 2.0.29
	"3af547d2f4db3728b75ac5d1b3a4f47830dc48e7": new lib.ModuleInfo("freeplay", "2.0.29", [], ["scenario"]),
};

/**
 * Names of the zip entries added by patching for the modules in the patch info
 */
function moduleEntries(zip: lib.ZipArchive, patchInfo: PatchInfo, root: string) {
	const entries = new Set<string>();
	if (patchInfo.version === 0) {
		const modulesPath = path.posix.join(root, "modules/");
		for (const name of zip.entries.keys()) {
			if (name.startsWith(modulesPath)) {
				entries.add(name);
			}
		}
	} else {
		for (const module of patchInfo.modules) {
			for (const filepath of module.files.keys()) {
				entries.add(path.posix.join(root, filepath));
			}
		}
	}
	return entries;
}

/**
 * Write a copy of the save next to it with entries removed and added
 *
 * Keeps the order of the entries in the original.
 *
 * @returns Path to the temporary file written.
 */
async function rewriteSave(
	zip: lib.ZipArchive, savePath: string, removed: Set<string>, added: Map<string, Buffer>
) {
	added = new Map(added);
	const output = new lib.ZipWriter();
	const [tempSavePath, writeStream] = await lib.createTempWriteStream(savePath);
	const writing = pipeline(output.outputStream, writeStream);
	// Awaited below, but a write error can reject it before then.
	writing.catch(() => {});
	try {
		for (const entry of zip.entries.values()) {
			const name = entry.fileName;
			if (removed.has(name)) {
				continue;
			}
			const contents = added.get(name);
			if (contents !== undefined) {
				await output.addBuffer(contents, name);
				added.delete(name);
			} else {
				await output.addEntry(zip, entry);
			}
		}
		for (const [name, contents] of added) {
			await output.addBuffer(contents, name);
		}
		await output.end();
		await writing;
	} catch (err) {
		output.outputStream.destroy();
		await writing.catch(() => {});
		await fs.rm(tempSavePath, { force: true });
		throw err;
	}
	return tempSavePath;
}

/**
 * Patch a save with the given modules
 *
 * Adds the modules given by the modules parameter to the save located
 * at savePath and rewrites the control.lua in the save to load the
 * modules that were added.  Will also remove any previous module
 * located in the save.
 *
 * @param savePath - Path to the Factorio save to patch.
 * @param modules - Description of the modules to patch.
 */
export async function patch(savePath: string, modules: SaveModule[]) {
	const zip = await lib.ZipArchive.fromFile(savePath);
	let tempSavePath: string;
	try {
		const root = lib.findRoot(zip);
		const rootPath = (relativePath: string) => path.posix.join(root, relativePath);

		// Files to add or replace in the save, relative to the root
		const files = new Map<string, Buffer>();
		let patchInfo: PatchInfo;
		const patchInfoContent = await zip.readFile(rootPath("clusterio.json"));
		if (patchInfoContent !== null) {
			patchInfo = await PatchInfo.fromSave(JSON.parse(patchInfoContent.toString("utf8")), zip, root);
			if (patchInfo.version > PatchInfo.currentVersion) {
				throw new Error(
					`Save patch version ${patchInfo.version} is newer than the patch version this ` +
					`version of Clusterio can load (${PatchInfo.currentVersion})`
				);
			}

			// Unpatching put the original control.lua back in place of scenario.lua
			const scenarioRequired = [...patchInfo.scenario.load, ...patchInfo.scenario.require].includes("scenario");
			if (scenarioRequired && !zip.file(rootPath("scenario.lua"))) {
				const control = await zip.readFile(rootPath("control.lua"));
				if (control === null) {
					throw new Error("Unable to patch save, missing control.lua file.");
				}
				files.set("scenario.lua", control);
			}

		// No info file present, try to detect if it's a known compatible scenario.
		} else {
			const controlFile = zip.file(rootPath("control.lua"));
			if (!controlFile) {
				throw new Error("Unable to patch save, missing control.lua file.");
			}
			const controlHash = await lib.hashStream(await zip.openStream(controlFile));

			if (controlHash in knownScenarios) {
				patchInfo = new PatchInfo(0, knownScenarios[controlHash], []);
				files.set("scenario.lua", await zip.readEntry(controlFile));
			} else {
				throw new Error(`Unable to patch save, unknown scenario (${controlHash})`);
			}
		}

		// Increment patch number
		patchInfo.patchNumber = patchInfo.patchNumber + 1;

		// Remove any existing modules from the save
		const removed = moduleEntries(zip, patchInfo, root);
		patchInfo.modules = [];

		reorderDependencies(modules);

		// Add the modules to the save.
		for (const module of modules) {
			for (const [relativePath, contents] of module.files) {
				files.set(relativePath, contents);
			}
			patchInfo.modules.push(module);
		}

		// Add loading code and patch info
		files.set("control.lua", Buffer.from(generateLoader(patchInfo), "utf8"));
		files.set("clusterio.json", Buffer.from(JSON.stringify(patchInfo, null, "\t"), "utf8"));

		const added = new Map([...files].map(([relativePath, contents]) => [rootPath(relativePath), contents]));
		tempSavePath = await rewriteSave(zip, savePath, removed, added);
	} finally {
		zip.close();
	}
	await fs.rename(tempSavePath, savePath);
}

/**
 * Remove the modules patched into a save
 *
 * Moves the original control.lua that patching kept as scenario.lua back
 * and removes the module files.  The patch info is kept with no modules so that
 * the patch number keeps counting if the save is patched again.
 *
 * @param savePath - Path to the Factorio save to unpatch.
 * @returns true if the save was changed.
 */
export async function unpatch(savePath: string) {
	const zip = await lib.ZipArchive.fromFile(savePath);
	let tempSavePath: string;
	try {
		const root = lib.findRoot(zip);
		const rootPath = (relativePath: string) => path.posix.join(root, relativePath);
		const patchInfoContent = await zip.readFile(rootPath("clusterio.json"));
		if (patchInfoContent === null) {
			return false;
		}
		const patchInfo = await PatchInfo.fromSave(JSON.parse(patchInfoContent.toString("utf8")), zip, root);
		if (patchInfo.version > PatchInfo.currentVersion) {
			throw new Error(
				`Save patch version ${patchInfo.version} is newer than the patch version this ` +
				`version of Clusterio can unpatch (${PatchInfo.currentVersion})`
			);
		}
		if (patchInfo.version !== 0 && !patchInfo.modules.length) {
			return false;
		}

		const removed = moduleEntries(zip, patchInfo, root);
		patchInfo.modules = [];
		const scenario = await zip.readFile(rootPath("scenario.lua"));
		if (scenario !== null) {
			removed.add(rootPath("scenario.lua"));
		}
		// Scenarios made for Clusterio have no scenario.lua, keep loading them without modules
		const control = scenario ?? Buffer.from(generateLoader(patchInfo), "utf8");
		tempSavePath = await rewriteSave(zip, savePath, removed, new Map([
			[rootPath("control.lua"), control],
			[rootPath("clusterio.json"), Buffer.from(JSON.stringify(patchInfo, null, "\t"), "utf8")],
		]));
	} finally {
		zip.close();
	}
	await fs.rename(tempSavePath, savePath);
	return true;
}

// For testing only
export const _generateLoader = generateLoader;
export const _reorderDependencies = reorderDependencies;