import * as path from "path";

import { SourceNode } from "source-map";

import * as ts from "typescript";

import { CompilerOptions, LuaTarget } from "../CompilerOptions";

import { escapeString, tstlHeader } from "../LuaPrinter";

import { cast, formatPathToLuaPath, isNonNull, trimExtension } from "../utils";

import { couldNotFindBundleEntryPoint } from "./diagnostics";

import { getEmitOutDir, getEmitPathRelativeToOutDir, getProjectRoot } from "./transpiler";

import { EmitFile, ProcessedFile } from "./utils";

const createModulePath = (pathToResolve: string, program: ts.Program) =>

escapeString(formatPathToLuaPath(trimExtension(getEmitPathRelativeToOutDir(pathToResolve, program))));

// Override `require` to read from ____modules table.

function requireOverride(options: CompilerOptions) {

const runModule =

options.luaTarget === LuaTarget.Lua50

? "if (table.getn(arg) > 0) then value = module(unpack(arg)) else value = module(file) end"

: 'if (select("#", ...) > 0) then value = module(...) else value = module(file) end';

return `

local ____modules = {}

local ____moduleCache = {}

local ____originalRequire = require

local function require(file, ...)

if ____moduleCache[file] then

return ____moduleCache[file].value

end

if ____modules[file] then

local module = ____modules[file]

local value = nil

${runModule}

____moduleCache[file] = { value = value }

return value

else

if ____originalRequire then

return ____originalRequire(file)

else

error("module '" .. file .. "' not found")

end

end

end

`;

}

export const sourceMapTracebackBundlePlaceholder = "{#SourceMapTracebackBundle}";

type SourceMapLineData = number | { line: number; file: string };

export function printStackTraceBundleOverride(rootNode: SourceNode): string {

const map: Record<number, SourceMapLineData> = {};

const getLineNumber = (line: number, fallback: number) => {

const data: SourceMapLineData | undefined = map[line];

if (data === undefined) {

return fallback;

}

if (typeof data === "number") {

return data;

}

return data.line;

};

const transformLineData = (data: SourceMapLineData) => {

if (typeof data === "number") {

return data;

}

return `{line = ${data.line}, file = "${data.file}"}`;

};

let currentLine = 1;

rootNode.walk((chunk, mappedPosition) => {

if (mappedPosition.line !== undefined && mappedPosition.line > 0) {

const line = getLineNumber(currentLine, mappedPosition.line);

map[currentLine] = {

line,

file: path.basename(mappedPosition.source),

};

}

currentLine += chunk.split("\n").length - 1;

});

const mapItems = Object.entries(map).map(([line, original]) => `["${line}"] = ${transformLineData(original)}`);

const mapString = "{" + mapItems.join(",") + "}";

return `__TS__SourceMapTraceBack(debug.getinfo(1).short_src, ${mapString});`;

}

export function getBundleResult(program: ts.Program, files: ProcessedFile[]): [ts.Diagnostic[], EmitFile] {

const diagnostics: ts.Diagnostic[] = [];

const options = program.getCompilerOptions() as CompilerOptions;

const bundleFile = cast(options.luaBundle, isNonNull);

const entryModule = cast(options.luaBundleEntry, isNonNull);

// Resolve project settings relative to project file.

const resolvedEntryModule = path.resolve(getProjectRoot(program), entryModule);

const outputPath = path.resolve(getEmitOutDir(program), bundleFile);

const entryModuleFilePath =

program.getSourceFile(entryModule)?.fileName ?? program.getSourceFile(resolvedEntryModule)?.fileName;

if (entryModuleFilePath === undefined) {

diagnostics.push(couldNotFindBundleEntryPoint(entryModule));

}

// For each file: ["<module path>"] = function() <lua content> end,

const moduleTableEntries = files.map(f => moduleSourceNode(f, createModulePath(f.fileName, program)));

// Create ____modules table containing all entries from moduleTableEntries

const moduleTable = createModuleTableNode(moduleTableEntries);

// return require("<entry module path>")

const args = options.luaTarget === LuaTarget.Lua50 ? "unpack(arg == nil and {} or arg)" : "...";

// Avoid producing a tail-call (which removes the bundle's stack frame) by assigning the `require` result to a local and returning it.

const entryPath = createModulePath(entryModuleFilePath ?? entryModule, program);

const entryPoint = `local ____entry = require(${entryPath}, ${args})\nreturn ____entry\n`;

const footers: string[] = [];

if (options.sourceMapTraceback) {

// Generates SourceMapTraceback for the entire file

footers.push('local __TS__SourceMapTraceBack = require("lualib_bundle").__TS__SourceMapTraceBack\n');

footers.push(`${sourceMapTracebackBundlePlaceholder}\n`);

}

const sourceChunks = [requireOverride(options), moduleTable, ...footers, entryPoint];

if (!options.noHeader) {

sourceChunks.unshift(tstlHeader);

}

const bundleNode = joinSourceChunks(sourceChunks);

let { code, map } = bundleNode.toStringWithSourceMap();

code = code.replace(sourceMapTracebackBundlePlaceholder, printStackTraceBundleOverride(bundleNode));

return [

diagnostics,

{

outputPath,

code,

sourceMap: map.toString(),

sourceFiles: files.flatMap(x => x.sourceFiles ?? []),

},

];

}

function moduleSourceNode({ code, sourceMapNode }: ProcessedFile, modulePath: string): SourceNode {

const tableEntryHead = `[${modulePath}] = function(...) \n`;

const tableEntryTail = " end,\n";

return joinSourceChunks([tableEntryHead, sourceMapNode ?? code, tableEntryTail]);

}

function createModuleTableNode(fileChunks: SourceChunk[]): SourceNode {

const tableHead = "____modules = {\n";

const tableEnd = "}\n";

return joinSourceChunks([tableHead, ...fileChunks, tableEnd]);

}

type SourceChunk = string | SourceNode;

function joinSourceChunks(chunks: SourceChunk[]): SourceNode {

return new SourceNode(null, null, null, chunks);

}