utils.js

/**
 * ⚠️ PRIVACY RESEARCH USE ONLY
 * For authorized testing environments that comply with all applicable laws.
 * See: https://github.com/botswin/FunCaptcha-Privacy-Research/blob/main/DISCLAIMER.md
 */

// @ts-check

import { traverse, types } from '@babel/core';
import { generate } from '@babel/generator';
import { parse } from '@babel/parser';

/**
 * Inline immutable constants and single‑assignment variables in the AST.
 * For example:
 *   var a = 1;
 *   b[a] = 2;
 * becomes:
 *   b[1] = 2;
 *
 * @param {object} ast - The Babel AST root node.
 */
export function inlineImmutableBindings(ast) {
    //
    // Iterate through the AST, dealing with invariant constants and single-assignment variables
    // For example:
    // var a = 1;
    // b[a] = 2;
    // Replace with:
    // b[1] = 2;
    traverse(ast, {
        // Handling constants initialized at declaration
        VariableDeclarator(path) {
            const { id, init } = path.node;
            if (!types.isIdentifier(id)) return;
            if (!(types.isStringLiteral(init) || types.isNumericLiteral(init))) return;

            const binding = path.scope.getBinding(id.name);
            if (!binding) return;
            // Ensure that variables are not reassigned
            if (binding.constant && binding.constantViolations.length === 0) {
                // Replace all references with literals
                binding.referencePaths.forEach((refPath) => {
                    refPath.replaceWith(types.cloneDeep(init));
                });
                // Delete statement
                path.remove();
            }
        },

        // Handling constants that are declared and then assigned
        AssignmentExpression(path) {
            const { node } = path;
            if (node.operator !== '=') return;
            const left = node.left;
            const right = node.right;
            if (!types.isIdentifier(left)) return;
            if (!(types.isStringLiteral(right) || types.isNumericLiteral(right))) return;

            const binding = path.scope.getBinding(left.name);
            if (!binding) return;
            // This is a one-time assignment only.
            if (binding.constantViolations.length === 1 && binding.constantViolations[0].node === node) {
                // Inline all references
                binding.referencePaths.forEach((refPath) => {
                    refPath.replaceWith(types.cloneDeep(right));
                });
                // Delete assignment statement: if the parent node is an ExpressionStatement, then delete the statement
                const parentStmt = path.parentPath;
                if (parentStmt && parentStmt.isExpressionStatement()) {
                    parentStmt.remove();
                } else {
                    path.remove();
                }

                // Delete the original statement
                const declPath = binding.path;
                if (declPath && declPath.remove) {
                    declPath.remove();
                }
            }
        },
    });
}

// Enhanced helper: follow alias chain via bindings until no further alias
// This version handles more complex alias patterns
function resolveRootBinding(binding) {
    const seen = new Set();
    let cur = binding;

    while (cur && !seen.has(cur)) {
        seen.add(cur);
        const path = cur.path;

        // case A: var X = Y;
        if (path.isVariableDeclarator()) {
            const init = path.node.init;
            if (types.isIdentifier(init)) {
                const next = path.scope.getBinding(init.name);
                if (next) {
                    cur = next;
                    continue;
                }
            }
        }

        // case B: X = Y;
        let jumped = false;
        for (const viol of cur.constantViolations) {
            // B.1: Direct assignment expressions
            if (
                viol.isAssignmentExpression() &&
                types.isIdentifier(viol.node.left, { name: cur.identifier.name }) &&
                types.isIdentifier(viol.node.right)
            ) {
                const next = viol.scope.getBinding(viol.node.right.name);
                if (next) {
                    cur = next;
                    jumped = true;
                    break;
                }
            }

            // B.2: Assignments in sequence expressions (comma operators)
            if (viol.isExpressionStatement() && viol.get('expression').isSequenceExpression()) {
                const expressions = viol.get('expression').node.expressions;
                for (const expr of expressions) {
                    if (
                        types.isAssignmentExpression(expr) &&
                        types.isIdentifier(expr.left, { name: cur.identifier.name }) &&
                        types.isIdentifier(expr.right)
                    ) {
                        const next = viol.scope.getBinding(expr.right.name);
                        if (next) {
                            cur = next;
                            jumped = true;
                            break;
                        }
                    }
                }
                if (jumped) break;
            }

            // B.3: Assignments in conditions
            if (viol.isIfStatement() || viol.isWhileStatement() || viol.isDoWhileStatement()) {
                jumped = tryFindAssignmentInCondition(viol.get('test'), cur, (next) => {
                    cur = next;
                    jumped = true;
                });
                if (jumped) break;
            }

            // B.4: For statements
            if (viol.isForStatement()) {
                // Check initialization
                if (viol.node.init && types.isAssignmentExpression(viol.node.init)) {
                    jumped = checkAndHandleAssignment(viol.node.init, viol.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                    if (jumped) break;
                }

                // Check test condition
                if (viol.node.test) {
                    jumped = tryFindAssignmentInCondition(viol.get('test'), cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                    if (jumped) break;
                }

                // Check update expression
                if (viol.node.update && types.isAssignmentExpression(viol.node.update)) {
                    jumped = checkAndHandleAssignment(viol.node.update, viol.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                    if (jumped) break;
                }
            }

            // B.5: Destructuring assignments
            if (viol.isAssignmentExpression() && types.isObjectPattern(viol.node.left)) {
                for (const prop of viol.node.left.properties) {
                    if (types.isIdentifier(prop.key, { name: cur.identifier.name }) && types.isIdentifier(prop.value)) {
                        const next = viol.scope.getBinding(prop.value.name);
                        if (next) {
                            cur = next;
                            jumped = true;
                            break;
                        }
                    }
                }
                if (jumped) break;
            }

            // B.6: Complex property access patterns - e.g., (o = (e = t)[(n = I)(1014) + "ension"])
            if (viol.isAssignmentExpression() && types.isMemberExpression(viol.node.right)) {
                // Check for assignments in the object of the member expression
                const objectPath = viol.get('right').get('object');
                if (objectPath.isAssignmentExpression()) {
                    jumped = checkAndHandleAssignment(objectPath.node, objectPath.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                    if (jumped) break;
                }

                // Check for assignments in the property of the member expression
                const propertyPath = viol.get('right').get('property');
                if (propertyPath.isAssignmentExpression()) {
                    jumped = checkAndHandleAssignment(propertyPath.node, propertyPath.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                    if (jumped) break;
                }

                // Check for complex cases like (n = I)(1014) + "ension" in property
                if (propertyPath.isBinaryExpression() && propertyPath.node.operator === '+') {
                    const leftProperty = propertyPath.get('left');
                    const rightProperty = propertyPath.get('right');

                    // Check left side of the binary expression
                    if (leftProperty.isCallExpression()) {
                        const callee = leftProperty.get('callee');
                        if (callee.isParenthesizedExpression()) {
                            const expr = callee.get('expression');
                            if (expr.isAssignmentExpression()) {
                                jumped = checkAndHandleAssignment(expr.node, expr.scope, cur, (next) => {
                                    cur = next;
                                    jumped = true;
                                });
                                if (jumped) break;
                            }
                        }
                    }

                    // Check right side of the binary expression
                    if (!jumped && rightProperty.isAssignmentExpression()) {
                        jumped = checkAndHandleAssignment(rightProperty.node, rightProperty.scope, cur, (next) => {
                            cur = next;
                            jumped = true;
                        });
                        if (jumped) break;
                    }
                }
            }
        }

        if (jumped) continue;

        // Case C: Check parent expressions for assignments in complex contexts
        let parent = path.parentPath;
        while (parent && !jumped) {
            // C.1: Handle assignments in sequence expressions (comma operators)
            if (parent.isSequenceExpression()) {
                const expressions = parent.node.expressions;
                for (let i = 0; i < expressions.length; i++) {
                    const expr = expressions[i];
                    if (
                        types.isAssignmentExpression(expr) &&
                        types.isIdentifier(expr.left, { name: cur.identifier.name }) &&
                        types.isIdentifier(expr.right)
                    ) {
                        const next = parent.scope.getBinding(expr.right.name);
                        if (next) {
                            cur = next;
                            jumped = true;
                            break;
                        }
                    }
                }
            }

            // C.2: Handle assignments in logical expressions (&&, ||, ??)
            if (parent.isLogicalExpression()) {
                // Check left side
                jumped = tryFindAssignmentInExpression(parent.get('left'), cur, (next) => {
                    cur = next;
                    jumped = true;
                });

                // Check right side if not jumped
                if (!jumped) {
                    jumped = tryFindAssignmentInExpression(parent.get('right'), cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }
            }

            // C.3: Handle assignments in conditional expressions (ternary)
            if (parent.isConditionalExpression()) {
                // Check test
                jumped = tryFindAssignmentInExpression(parent.get('test'), cur, (next) => {
                    cur = next;
                    jumped = true;
                });

                // Check consequent if not jumped
                if (!jumped) {
                    jumped = tryFindAssignmentInExpression(parent.get('consequent'), cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }

                // Check alternate if not jumped
                if (!jumped) {
                    jumped = tryFindAssignmentInExpression(parent.get('alternate'), cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }
            }

            // C.4: Handle assignment and immediate call pattern: (n = I)(1014)
            if (parent.isCallExpression()) {
                const callee = parent.node.callee;
                // Check if callee is a parenthesized expression
                if (types.isParenthesizedExpression(callee) && types.isAssignmentExpression(callee.expression)) {
                    jumped = checkAndHandleAssignment(callee.expression, parent.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }
                // Also check for assignments in a sequence expression used as the callee
                else if (types.isSequenceExpression(callee)) {
                    for (const expr of callee.expressions) {
                        if (types.isAssignmentExpression(expr)) {
                            jumped = checkAndHandleAssignment(expr, parent.scope, cur, (next) => {
                                cur = next;
                                jumped = true;
                            });
                            if (jumped) break;
                        }
                    }
                }
            }

            // C.5: Handle assignments in member expressions - (o = (e = t)[(n = I)(1014) + "ension"])
            if (parent.isMemberExpression()) {
                // Check for assignments in the object part
                const objectPath = parent.get('object');
                if (objectPath.isAssignmentExpression()) {
                    jumped = checkAndHandleAssignment(objectPath.node, objectPath.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }

                // Check for assignments in the property part
                if (!jumped) {
                    const propertyPath = parent.get('property');
                    if (propertyPath.isAssignmentExpression()) {
                        jumped = checkAndHandleAssignment(propertyPath.node, propertyPath.scope, cur, (next) => {
                            cur = next;
                            jumped = true;
                        });
                    }

                    // Check for complex patterns in property access - (n = I)(1014) + "ension"
                    if (!jumped && propertyPath.isBinaryExpression() && propertyPath.node.operator === '+') {
                        const leftProperty = propertyPath.get('left');

                        // Check left side of the binary expression for function calls with assignment
                        if (leftProperty.isCallExpression()) {
                            const callee = leftProperty.get('callee');
                            if (callee.isParenthesizedExpression()) {
                                const expr = callee.get('expression');
                                if (expr.isAssignmentExpression()) {
                                    jumped = checkAndHandleAssignment(expr.node, expr.scope, cur, (next) => {
                                        cur = next;
                                        jumped = true;
                                    });
                                }
                            }
                        }
                    }
                }
            }

            // C.6: Handle assignments within return statements
            if (parent.isReturnStatement() && parent.node.argument) {
                const argPath = parent.get('argument');

                // Check for direct assignment in return
                if (argPath.isAssignmentExpression()) {
                    jumped = checkAndHandleAssignment(argPath.node, argPath.scope, cur, (next) => {
                        cur = next;
                        jumped = true;
                    });
                }

                // Check for property access with assignments
                if (!jumped && argPath.isMemberExpression()) {
                    const objectPath = argPath.get('object');
                    if (objectPath.isAssignmentExpression()) {
                        jumped = checkAndHandleAssignment(objectPath.node, objectPath.scope, cur, (next) => {
                            cur = next;
                            jumped = true;
                        });
                    }

                    if (!jumped) {
                        const propertyPath = argPath.get('property');
                        if (
                            propertyPath.isAssignmentExpression() ||
                            propertyPath.isCallExpression() ||
                            propertyPath.isBinaryExpression()
                        ) {
                            jumped = tryFindAssignmentInExpression(propertyPath, cur, (next) => {
                                cur = next;
                                jumped = true;
                            });
                        }
                    }
                }

                // Check for sequence expressions in return
                if (!jumped && argPath.isSequenceExpression()) {
                    for (const expr of argPath.node.expressions) {
                        if (types.isAssignmentExpression(expr)) {
                            jumped = checkAndHandleAssignment(expr, argPath.scope, cur, (next) => {
                                cur = next;
                                jumped = true;
                            });
                            if (jumped) break;
                        }
                    }
                }
            }

            // Move up to next parent
            parent = parent.parentPath;
        }

        if (jumped) continue;
        break;
    }

    return cur;
}

// Helper function to check for assignments in conditions
function tryFindAssignmentInCondition(path, curBinding, callback) {
    if (!path) return false;

    // Check for direct assignment in condition
    if (path.isAssignmentExpression()) {
        return checkAndHandleAssignment(path.node, path.scope, curBinding, callback);
    }

    // Check for assignments in sequence expressions (comma operators)
    if (path.isSequenceExpression()) {
        const expressions = path.node.expressions;
        for (const expr of expressions) {
            if (types.isAssignmentExpression(expr)) {
                const result = checkAndHandleAssignment(expr, path.scope, curBinding, callback);
                if (result) return true;
            }
        }
    }

    // Check for assignments in parenthesized expressions
    if (path.isParenthesizedExpression()) {
        return tryFindAssignmentInCondition(path.get('expression'), curBinding, callback);
    }

    // Check for assignments in logical expressions (&&, ||, ??)
    if (path.isLogicalExpression()) {
        // Check left side
        let result = tryFindAssignmentInCondition(path.get('left'), curBinding, callback);
        if (result) return true;

        // Check right side
        return tryFindAssignmentInCondition(path.get('right'), curBinding, callback);
    }

    // Check for assignments in call expressions - handles (n = I)(1014) pattern
    if (path.isCallExpression()) {
        const callee = path.get('callee');

        if (callee.isParenthesizedExpression()) {
            const expr = callee.get('expression');
            if (expr.isAssignmentExpression()) {
                return checkAndHandleAssignment(expr.node, expr.scope, curBinding, callback);
            }
        }
    }

    return false;
}

// Helper function to check for assignments in expressions
function tryFindAssignmentInExpression(path, curBinding, callback) {
    if (!path) return false;

    // Check for direct assignment
    if (path.isAssignmentExpression()) {
        return checkAndHandleAssignment(path.node, path.scope, curBinding, callback);
    }

    // Check for assignments in sequence expressions (comma operators)
    if (path.isSequenceExpression()) {
        const expressions = path.node.expressions;
        for (const expr of expressions) {
            if (types.isAssignmentExpression(expr)) {
                const result = checkAndHandleAssignment(expr, path.scope, curBinding, callback);
                if (result) return true;
            }
        }
    }

    // Check for assignments in parenthesized expressions
    if (path.isParenthesizedExpression()) {
        return tryFindAssignmentInExpression(path.get('expression'), curBinding, callback);
    }

    // Check for assignments in call expressions where callee is an assignment
    if (path.isCallExpression()) {
        const callee = path.get('callee');

        // Handle (n = I)(1014) pattern
        if (callee.isParenthesizedExpression()) {
            const expr = callee.get('expression');
            if (expr.isAssignmentExpression()) {
                return checkAndHandleAssignment(expr.node, expr.scope, curBinding, callback);
            }
        }

        // Handle other complex call expressions that might contain assignments
        if (callee.isSequenceExpression()) {
            const expressions = callee.node.expressions;
            for (const expr of expressions) {
                if (types.isAssignmentExpression(expr)) {
                    const result = checkAndHandleAssignment(expr, path.scope, curBinding, callback);
                    if (result) return true;
                }
            }
        }
    }

    return false;
}

// Helper function to check an assignment and handle it if it's relevant
function checkAndHandleAssignment(assignmentNode, scope, curBinding, callback) {
    if (
        types.isIdentifier(assignmentNode.left, { name: curBinding.identifier.name }) &&
        types.isIdentifier(assignmentNode.right)
    ) {
        const next = scope.getBinding(assignmentNode.right.name);
        if (next) {
            callback(next);
            return true;
        }
    }
    return false;
}

/**
 * Bind functions called with a single numeric argument to window.
 * Supports alias chains like a → b → c → Ke.
 *
 * @param {object} ast - The Babel AST root node.
 * @returns {object} A new AST with window bindings injected.
 */
export function bindSingleArgCalls(ast) {
    // 1. collect root bindings behind calls like fn(123)
    const rootBindings = new Set();
    traverse(ast, {
        CallExpression(path) {
            const callee = path.node.callee;
            const args = path.node.arguments;
            if (
                !(types.isIdentifier(callee) || types.isAssignmentExpression(callee)) ||
                args.length !== 1 ||
                !types.isNumericLiteral(args[0])
            ) {
                return;
            }

            let binding;
            if (callee.left?.name) {
                binding = path.scope.getBinding(callee.left.name);
            } else {
                binding = path.scope.getBinding(callee.name);
            }

            if (!binding) return;
            const root = resolveRootBinding(binding);
            rootBindings.add(root);
        },
    });

    // 2. inject window bindings via string replace (unchanged)
    let codeStr = generate(ast, { retainLines: true }).code;
    rootBindings.forEach((root) => {
        const fnName = root.identifier.name;
        const re = new RegExp(`return\\s+\\(?${fnName}\\s*=\\s*function\\s*\\([^)]*\\)\\s*\\{`, 'g');
        codeStr = codeStr.replace(re, (match) => `${match}window.__bind_${fnName} = ${fnName};`);
    });

    // 3. parse back into AST and return
    const ast2 = parse(codeStr, { sourceType: 'module' });
    return ast2;
}

/**
 * 1. Retrieve the alias -> realName map for this AST.
 * 2. Find all calls aliasName(<number>), look up realName,
 *    invoke window.__bind_realName(<number>) in page,
 *    and replace the call node with the returned literal.
 *
 * @param {import('puppeteer').Page} page
 * @param {object} ast - the same AST passed to bindSingleArgCalls
 */
export async function evaluateAndInlineCalls(page, ast) {
    const callPaths = [];
    traverse(ast, {
        CallExpression(path) {
            const callee = path.node.callee;
            const args = path.node.arguments;
            if (
                !(types.isIdentifier(callee) || types.isAssignmentExpression(callee)) ||
                args.length !== 1 ||
                !types.isNumericLiteral(args[0])
            ) {
                return;
            }

            let binding;
            if (callee.left?.name) {
                binding = path.scope.getBinding(callee.left.name);
            } else {
                binding = path.scope.getBinding(callee.name);
            }

            if (!binding) return;
            const root = resolveRootBinding(binding);

            // call window.__bind_<root>(<number>)
            callPaths.push({
                path,
                root,
                arg: args[0],
            });
        },
    });

    // Invoke the function in the page context
    for (const { path, root, arg } of callPaths) {
        try {
            const result = await page.evaluate(
                (name, num) => {
                    return window[`__bind_${name}`](num);
                },
                root.identifier.name,
                arg.value
            );

            // Replace the call node with the returned literal
            if (result) {
                path.replaceWith(types.stringLiteral(result));
            }
        } catch (err) {
            // Ignore
        }
    }
}

function evaluateStringConcat(expr) {
    if (types.isStringLiteral(expr)) {
        return expr.value;
    }
    if (types.isBinaryExpression(expr, { operator: '+' })) {
        const left = evaluateStringConcat(expr.left);
        const right = evaluateStringConcat(expr.right);
        if (typeof left === 'string' && typeof right === 'string') {
            return left + right;
        }
    }
    return null;
}

export function mergeStringConcat(ast) {
    // Merge string concatenations into literals
    traverse(ast, {
        BinaryExpression: {
            exit(path) {
                if (path.node.operator === '+') {
                    const evaluated = evaluateStringConcat(path.node);
                    if (typeof evaluated === 'string') {
                        path.replaceWith(types.stringLiteral(evaluated));
                    }
                }
            },
        },
    });
}

export function convertComputedToProperty(ast) {
    // Convert computed properties to regular properties
    // For example: p['abc'] = 123 becomes p.abc = 123
    traverse(ast, {
        MemberExpression(path) {
            if (path.node.computed && types.isStringLiteral(path.node.property)) {
                const propName = path.node.property.value;
                if (/^[a-zA-Z_$][a-zA-Z0-9_$]*$/.test(propName)) {
                    path.node.computed = false;
                    path.node.property = types.identifier(propName);
                }
            }
        },
    });
}

/**
 * Removes all functions from AST that match a specific pattern:
 * Functions with a large string array and a self-redefining return pattern
 *
 * @param {Object} ast - The Babel AST to process
 */
export function removeArrayFunctions(ast) {
    const visitor = {
        // Handle function declarations
        FunctionDeclaration(path) {
            if (isTargetArrayFunction(path, types)) {
                path.remove();
            }
        },

        // Handle function expressions
        FunctionExpression(path) {
            // Function expression in variable declaration: const x = function() {...}
            if (path.parent && path.parent.type === 'VariableDeclarator' && isTargetArrayFunction(path, types)) {
                path.parentPath.parentPath.remove();
                return;
            }

            // Function expression in assignment: x = function() {...}
            if (
                path.parent &&
                path.parent.type === 'AssignmentExpression' &&
                path.parent.right === path.node &&
                isTargetArrayFunction(path, types)
            ) {
                path.parentPath.remove();
                return;
            }

            // Function expression in object property: { x: function() {...} }
            if (path.parent && path.parent.type === 'ObjectProperty' && isTargetArrayFunction(path, types)) {
                path.parent.remove();
            }
        },
    };

    traverse(ast, visitor);
    return ast;
}

/**
 * Determines if a function matches the target pattern
 *
 * @param {NodePath} path - Babel function path object
 * @param {Object} types - Babel types object
 * @returns {boolean} - Whether the function matches the target pattern
 */
function isTargetArrayFunction(path, types) {
    const { node } = path;

    // Ensure we have a function with a body
    if (!node.body || !node.body.body) {
        return false;
    }

    const body = node.body.body;

    // Function body needs at least two statements (array declaration and return statement)
    if (body.length < 2) {
        return false;
    }

    // 1. Look for a large string array variable declaration
    let hasLargeStringArray = false;
    let arrayVarName = null;

    for (const stmt of body) {
        // Look for variable declarations
        if (types.isVariableDeclaration(stmt)) {
            for (const declarator of stmt.declarations) {
                // Check if it's an array expression
                if (types.isArrayExpression(declarator.init)) {
                    const elements = declarator.init.elements;

                    // Check if it has at least 10 elements, and all are string literals
                    if (
                        elements &&
                        elements.length >= 10 &&
                        elements.every((el) => el && (types.isStringLiteral(el) || types.isTemplateLiteral(el)))
                    ) {
                        hasLargeStringArray = true;
                        arrayVarName = declarator.id.name;
                        break;
                    }
                }
            }
        }

        if (hasLargeStringArray) break;
    }

    // If no matching array found, return false
    if (!hasLargeStringArray || !arrayVarName) {
        return false;
    }

    // 2. Check if the last statement is a return statement
    const lastStmt = body[body.length - 1];
    if (!types.isReturnStatement(lastStmt)) {
        return false;
    }

    // 3. Check if the return statement matches the pattern: return (Func = function() { return array; })();
    const returnArg = lastStmt.argument;

    // Handle standard pattern: return (Func = function() { return array; })();
    if (types.isCallExpression(returnArg)) {
        const callee = returnArg.callee;

        // Check if it's a parenthesized expression
        if (types.isParenthesizedExpression(callee)) {
            const inner = callee.expression;

            // Check if the inner expression is an assignment
            if (types.isAssignmentExpression(inner)) {
                const right = inner.right;

                // Check if the right side of the assignment is a function expression
                if (types.isFunctionExpression(right)) {
                    const funcBody = right.body.body;

                    // Function body has only one return statement
                    if (funcBody.length === 1 && types.isReturnStatement(funcBody[0])) {
                        const innerReturnArg = funcBody[0].argument;

                        // The return is the previously defined array variable
                        if (types.isIdentifier(innerReturnArg) && innerReturnArg.name === arrayVarName) {
                            return true;
                        }
                    }
                }
            }
        }

        // Handle assignment expression without parentheses: return (Func = function(){...})()
        if (types.isAssignmentExpression(callee)) {
            const right = callee.right;

            if (types.isFunctionExpression(right)) {
                const funcBody = right.body.body;

                if (funcBody.length === 1 && types.isReturnStatement(funcBody[0])) {
                    const innerReturnArg = funcBody[0].argument;

                    if (types.isIdentifier(innerReturnArg) && innerReturnArg.name === arrayVarName) {
                        return true;
                    }
                }
            }
        }
    }

    return false;
}

/**
 * Removes all shuffle IIFE functions from the AST using string pattern matching
 *
 * @param {Object} ast - The Babel AST to process
 */
export function removeShuffleIIFE(ast) {
    const visitor = {
        // Match both !(function(){})() and (function())()
        UnaryExpression(path) {
            if (
                path.node.operator === '!' &&
                types.isCallExpression(path.node.argument) &&
                types.isFunctionExpression(path.node.argument.callee)
            ) {
                const code = generate(path.node).code;

                if (isShuffleIIFEPattern(code)) {
                    path.remove();
                }
            }
        },

        CallExpression(path) {
            if (
                (types.isFunctionExpression(path.node.callee) ||
                    (types.isParenthesizedExpression(path.node.callee) &&
                        types.isFunctionExpression(path.node.callee.expression))) &&
                !types.isUnaryExpression(path.parent)
            ) {
                const code = generate(path.node).code;

                if (isShuffleIIFEPattern(code)) {
                    path.remove();
                }
            }
        },
    };

    traverse(ast, visitor);

    return ast;
}

/**
 * Checks if the code string matches the shuffle IIFE pattern
 *
 * @param {string} code - Generated code string
 * @returns {boolean} - Whether it matches the pattern
 */
function isShuffleIIFEPattern(code) {
    // Check for required patterns in the code string
    const hasForLoop = code.includes('for');
    const hasParseInt = (code.match(/parseInt/g) || []).length >= 2;
    const hasPush = (code.match(/\.push/g) || []).length >= 2;
    const hasShift = (code.match(/\.shift/g) || []).length >= 2;
    const hasFunction = (code.match(/function/g) || []).length == 1;

    return hasForLoop && hasParseInt && hasPush && hasShift && hasFunction;
}