"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.transformModule = transformModule;
const error_1 = require("./utils/error");
const make_node_1 = require("./utils/make-node");
const random_identifier_1 = require("./utils/random-identifier");
const resolve_import_1 = require("./utils/resolve-import");
const resolve_pattern_1 = require("./utils/resolve-pattern");
function transformByType(node, type, filter) {
    for (const key of Array.isArray(node) ? node.keys() : Object.keys(node)) {
        if (["span", "type"].includes(key))
            continue;
        if (!node[key])
            continue;
        if (typeof node[key] === "object")
            node[key] = transformByType(node[key], type, filter);
    }
    if (node["type"] === type)
        return filter(node);
    return node;
}
function declarationToExpression(decl) {
    if (decl.type === "FunctionDeclaration") {
        return {
            ...decl,
            identifier: null,
            type: "FunctionExpression"
        };
    }
    else if (decl.type === "ClassDeclaration") {
        return {
            ...decl,
            identifier: null,
            type: "ClassExpression"
        };
    }
    else {
        /* istanbul ignore next */
        (0, error_1.raiseUnexpectedNode)("declaration", decl.type);
    }
}
function expressionToDeclaration(expr) {
    if (expr.type === "FunctionExpression") {
        return {
            ...expr,
            type: "FunctionDeclaration"
        };
    }
    else if (expr.type === "ClassExpression") {
        return {
            ...expr,
            type: "ClassDeclaration"
        };
    }
    else {
        /* istanbul ignore next */
        (0, error_1.raiseUnexpectedNode)("expression", expr.type);
    }
}
function transformModule(code, ast, moduleName, bundleInfo, options) {
    var _a, _b;
    const randomIdentifier = new random_identifier_1.RandomIdentifierGenerator(code);
    // Extract import declarations
    const imports = ast.body.filter((item) => item.type === "ImportDeclaration");
    // Handle `export { named as renamed } from "module";`
    const exportFroms = ast.body.filter((item) => item.type === "ExportNamedDeclaration" && !!item.source);
    const newImportsByExportFroms = [];
    const exportMap = {};
    // Extract export declarations
    // In Rollup's output, there should be only one, and as the last top-level statement
    // But some plugins (e.g. @vitejs/plugin-legacy) may inject others like "export function"
    const namedExports = ast.body.filter((item, i) => {
        switch (item.type) {
            /* istanbul ignore next */
            case "ExportAllDeclaration":
                (0, error_1.raiseUnexpectedNode)("top-level statement", item.type);
            case "ExportDefaultExpression":
                // Convert to a variable
                const identifier = randomIdentifier.generate();
                ast.body[i] = (0, make_node_1.makeVariableInitDeclaration)(identifier, item.expression);
                exportMap["default"] = identifier;
                return false;
            case "ExportDefaultDeclaration":
                if (item.decl.type === "FunctionExpression" || item.decl.type === "ClassExpression") {
                    // Convert to a declaration or variable
                    if (item.decl.identifier) {
                        ast.body[i] = expressionToDeclaration(item.decl);
                        exportMap["default"] = item.decl.identifier.value;
                    }
                    else {
                        const identifier = randomIdentifier.generate();
                        ast.body[i] = (0, make_node_1.makeVariableInitDeclaration)(identifier, item.decl);
                        exportMap["default"] = identifier;
                    }
                }
                else {
                    /* istanbul ignore next */
                    (0, error_1.raiseUnexpectedNode)("top-level export declaration", item.decl.type);
                }
                return false;
            case "ExportDeclaration":
                if (item.declaration.type === "FunctionDeclaration" || item.declaration.type === "ClassDeclaration") {
                    // Remove the "export" keyword from this statement
                    ast.body[i] = item.declaration;
                    exportMap[item.declaration.identifier.value] = item.declaration.identifier.value;
                }
                else {
                    /* istanbul ignore next */
                    (0, error_1.raiseUnexpectedNode)("top-level export declaration", item.declaration.type);
                }
                return false;
            // Handle `export { named as renamed };` without "from"
            case "ExportNamedDeclaration":
                if (!item.source) {
                    item.specifiers.forEach(specifier => {
                        /* istanbul ignore if */
                        if (specifier.type !== "ExportSpecifier") {
                            (0, error_1.raiseUnexpectedNode)("export specifier", specifier.type);
                        }
                        exportMap[(specifier.exported || specifier.orig).value] = specifier.orig.value;
                    });
                }
                return true;
        }
        return false;
    });
    const exportedNameSet = new Set(Object.values(exportMap));
    const exportedNames = Array.from(exportedNameSet);
    /*
     * Move ALL top-level statements to an async IIFE:
     *
     * ```js
     * export let __tla = Promise.all([
     *   // imported TLA promises
     * ]).then(async () => {
     *   // original top-level statements here
     * });
     * ```
     *
     * And add variable declarations for exported names to new top-level, before the IIFE.
     *
     * ```js
     * let x;
     * export let __tla = Promise.all([
     *   // imported TLA promises
     * ]).then(async () => {
     *   // const x = 1;
     *   x = 1;
     * });
     * export { x as someExport };
     * ```
     */
    const topLevelStatements = ast.body.filter((item) => !imports.includes(item) && !namedExports.includes(item));
    const importedNames = new Set(imports.flatMap(importStmt => importStmt.specifiers.map(specifier => specifier.local.value)));
    const exportFromedNames = new Set(exportFroms.flatMap(exportStmt => exportStmt.specifiers.map(specifier => {
        if (specifier.type === "ExportNamespaceSpecifier") {
            return specifier.name.value;
        }
        else if (specifier.type === "ExportDefaultSpecifier") {
            // When will this happen?
            return specifier.exported.value;
        }
        else {
            return (specifier.exported || specifier.orig).value;
        }
    })));
    const exportedNamesDeclaration = (0, make_node_1.makeVariablesDeclaration)(exportedNames.filter(name => !importedNames.has(name) && !exportFromedNames.has(name)));
    const warppedStatements = topLevelStatements.flatMap(stmt => {
        if (stmt.type === "VariableDeclaration") {
            const declaredNames = stmt.declarations.flatMap(decl => (0, resolve_pattern_1.resolvePattern)(decl.id));
            const exportedDeclaredNames = declaredNames.filter(name => exportedNameSet.has(name));
            const unexportedDeclaredNames = declaredNames.filter(name => !exportedNameSet.has(name));
            // None is exported in the declared names, no need to transform
            if (exportedDeclaredNames.length === 0)
                return stmt;
            // Generate assignment statements for init-ed declarators
            const assignmentStatements = stmt.declarations
                .filter(decl => decl.init)
                .map(decl => (0, make_node_1.makeAssignmentStatement)(decl.id, decl.init));
            // Generate variable declarations for unexported variables
            const unexportedDeclarations = (0, make_node_1.makeVariablesDeclaration)(unexportedDeclaredNames);
            return unexportedDeclarations ? [unexportedDeclarations, ...assignmentStatements] : assignmentStatements;
        }
        else if (stmt.type === "FunctionDeclaration" || stmt.type === "ClassDeclaration") {
            const name = stmt.identifier.value;
            if (!exportedNameSet.has(name))
                return stmt;
            return (0, make_node_1.makeAssignmentStatement)((0, make_node_1.makeIdentifier)(name), declarationToExpression(stmt));
        }
        else {
            return stmt;
        }
    });
    /*
     * Process dynamic imports.
     *
     * ```js
     * [
     *   import("some-module-with-tla"),
     *   import("some-module-without-tla"),
     *   import(dynamicModuleName)
     * ]
     * ```
     *
     * The expression evaluates to a promise, which will resolve after module loaded, but not after
     * out `__tla` promise resolved.
     *
     * We can check the target module. If the argument is string literial and the target module has NO
     * top-level await, we won't need to transform it.
     *
     * ```js
     * [
     *   import("some-module-with-tla").then(async m => { await m.__tla; return m; }),
     *   import("some-module-without-tla"),
     *   import(dynamicModuleName).then(async m => { await m.__tla; return m; })
     * ]
     * ```
     */
    transformByType(warppedStatements, "CallExpression", (call) => {
        var _a;
        if (call.callee.type === "Import") {
            const argument = call.arguments[0].expression;
            if (argument.type === "StringLiteral") {
                const importedModuleName = (0, resolve_import_1.resolveImport)(moduleName, argument.value);
                // Skip transform
                if (importedModuleName && !((_a = bundleInfo[importedModuleName]) === null || _a === void 0 ? void 0 : _a.transformNeeded))
                    return call;
            }
            return (0, make_node_1.makeCallExpression)((0, make_node_1.makeMemberExpression)(call, "then"), [
                (0, make_node_1.makeArrowFunction)(["m"], [
                    (0, make_node_1.makeStatement)((0, make_node_1.makeAwaitExpression)((0, make_node_1.makeMemberExpression)("m", options.promiseExportName))),
                    (0, make_node_1.makeReturnStatement)((0, make_node_1.makeIdentifier)("m"))
                ], true)
            ]);
        }
        return call;
    });
    /*
     * Import and await the promise "__tla" from each imported module with TLA transform enabled.
     *
     * ```js
     * import { ..., __tla as __tla_0 } from "...";
     * import { ..., __tla as __tla_1 } from "...";
     * ```
     *
     * To work with circular dependency, wrap each imported promise with try-catch.
     * Promises from circular dependencies will not be imported and awaited.
     *
     * ```js
     * export let __tla = Promise.all([
     *   (() => { try { return __tla_0; } catch {} })(),
     *   (() => { try { return __tla_1; } catch {} })()
     * ]).then(async () => {
     *   // original top-level statements here
     * });
     * ```
     */
    // Add import of TLA promises from imported modules
    let importedPromiseCount = 0;
    for (const declaration of [...imports, ...exportFroms]) {
        const importedModuleName = (0, resolve_import_1.resolveImport)(moduleName, declaration.source.value);
        if (!importedModuleName || !bundleInfo[importedModuleName])
            continue;
        if (bundleInfo[importedModuleName].transformNeeded) {
            let targetImportDeclaration;
            if (declaration.type === "ImportDeclaration") {
                targetImportDeclaration = declaration;
            }
            else {
                targetImportDeclaration = (0, make_node_1.makeImportDeclaration)(declaration.source);
                newImportsByExportFroms.push(targetImportDeclaration);
            }
            targetImportDeclaration.specifiers.push((0, make_node_1.makeImportSpecifier)(options.promiseExportName, options.promiseImportName(importedPromiseCount)));
            importedPromiseCount++;
        }
    }
    const importedPromiseArray = importedPromiseCount === 0
        ? null
        : (0, make_node_1.makeArrayExpression)([...Array(importedPromiseCount).keys()].map(i => (0, make_node_1.makeCallExpression)((0, make_node_1.makeArrowFunction)([], [(0, make_node_1.makeTryCatchStatement)([(0, make_node_1.makeReturnStatement)((0, make_node_1.makeIdentifier)(options.promiseImportName(i)))], [])]))));
    // The `async () => { /* original top-level statements */ }` function
    const wrappedTopLevelFunction = (0, make_node_1.makeArrowFunction)([], warppedStatements, true);
    // `Promise.all([ /* ... */]).then(async () => { /* ... */ })` or `(async () => {})()`
    const promiseExpression = importedPromiseArray
        ? (0, make_node_1.makeCallExpression)((0, make_node_1.makeMemberExpression)((0, make_node_1.makeCallExpression)((0, make_node_1.makeMemberExpression)("Promise", "all"), [importedPromiseArray]), "then"), [wrappedTopLevelFunction])
        : (0, make_node_1.makeCallExpression)(wrappedTopLevelFunction);
    /*
     * New top-level after transformation:
     *
     * import { ..., __tla as __tla_0 } from "some-module-with-TLA";
     * import { ... } from "some-module-without-TLA";
     *
     * let some, variables, exported, from, original, top, level;
     *
     * let __tla = Promise.all([ ... ]).then(async () => {
     *   ...
     * });
     *
     * export { ..., __tla };
     */
    const newTopLevel = [
        ...imports,
        ...newImportsByExportFroms,
        ...exportFroms,
        exportedNamesDeclaration
    ];
    if (exportedNames.length > 0 || ((_b = (_a = bundleInfo[moduleName]) === null || _a === void 0 ? void 0 : _a.importedBy) === null || _b === void 0 ? void 0 : _b.length) > 0) {
        // If the chunk is being imported, append export of the TLA promise to export list
        const promiseDeclaration = (0, make_node_1.makeVariableInitDeclaration)(options.promiseExportName, promiseExpression);
        exportMap[options.promiseExportName] = options.promiseExportName;
        newTopLevel.push(promiseDeclaration, (0, make_node_1.makeExportListDeclaration)(Object.entries(exportMap)));
    }
    else {
        // If the chunk is an entry, just execute the promise expression
        newTopLevel.push((0, make_node_1.makeStatement)(promiseExpression));
    }
    ast.body = newTopLevel.filter(x => x);
    return ast;
}
