/**
 * 简化 switch 控制流平坦化
 * 将 while(true) { switch(arr[i++]) { case '0': ... } } 转换为顺序执行
 */
const fs = require('fs');
const babel = require('@babel/core');
const traverse = require('@babel/traverse').default;
const generate = require('@babel/generator').default;
const t = require('@babel/types');

const inputPath = 'D:/temp/破解/cursorpro-0.4.5/deobfuscated_full/extension/out/webview/provider_simplified.js';
const outputPath = 'D:/temp/破解/cursorpro-0.4.5/deobfuscated_full/extension/out/webview/provider_final.js';

console.log('╔════════════════════════════════════════════════════╗');
console.log('║       Switch 控制流还原工具                        ║');
console.log('╚════════════════════════════════════════════════════╝');

const code = fs.readFileSync(inputPath, 'utf8');
console.log(`读取文件: ${(code.length / 1024).toFixed(2)} KB`);

console.log('\n[1] 解析 AST...');
let ast;
try {
    ast = babel.parseSync(code, {
        sourceType: 'script',
        plugins: []
    });
    console.log('AST 解析成功');
} catch (e) {
    console.error('AST 解析失败:', e.message);
    process.exit(1);
}

console.log('\n[2] 查找并还原 switch 控制流...');
let restoredCount = 0;

// 查找并收集变量定义
function findVariableValue(path, varName) {
    let value = null;

    // 在当前作用域向上查找
    const binding = path.scope.getBinding(varName);
    if (binding && binding.path.isVariableDeclarator()) {
        const init = binding.path.node.init;
        if (t.isCallExpression(init) &&
            t.isMemberExpression(init.callee) &&
            t.isStringLiteral(init.callee.object) &&
            t.isIdentifier(init.callee.property, { name: 'split' })) {
            // "1|4|0|3|2".split('|')
            value = init.callee.object.value.split('|');
        }
    }

    return value;
}

traverse(ast, {
    WhileStatement(path) {
        // 检查是否是 while(true) 或 while(1)
        const test = path.node.test;
        if (!t.isBooleanLiteral(test, { value: true }) &&
            !t.isNumericLiteral(test, { value: 1 })) {
            return;
        }

        const body = path.node.body;
        if (!t.isBlockStatement(body)) return;

        // 查找 switch 语句
        const switchStmt = body.body.find(stmt => t.isSwitchStatement(stmt));
        if (!switchStmt) return;

        // 检查 discriminant 是否是 arr[i++] 形式
        const disc = switchStmt.discriminant;
        if (!t.isMemberExpression(disc)) return;
        if (!t.isUpdateExpression(disc.property) || disc.property.operator !== '++') return;

        const arrName = disc.object.name;
        const indexName = disc.property.argument.name;

        // 获取顺序数组
        const order = findVariableValue(path, arrName);
        if (!order) return;

        // 收集 case 语句
        const cases = new Map();
        for (const caseStmt of switchStmt.cases) {
            if (!caseStmt.test) continue; // default case

            let caseValue;
            if (t.isStringLiteral(caseStmt.test)) {
                caseValue = caseStmt.test.value;
            } else if (t.isNumericLiteral(caseStmt.test)) {
                caseValue = String(caseStmt.test.value);
            } else {
                continue;
            }

            // 过滤掉 continue 和 break
            const consequent = caseStmt.consequent.filter(stmt =>
                !t.isContinueStatement(stmt) && !t.isBreakStatement(stmt)
            );

            cases.set(caseValue, consequent);
        }

        // 按顺序组装语句
        const newStatements = [];
        for (const key of order) {
            const stmts = cases.get(key);
            if (stmts) {
                newStatements.push(...stmts);
            }
        }

        if (newStatements.length > 0) {
            // 查找并移除相关变量声明
            const parentBody = path.parentPath;
            if (t.isBlockStatement(parentBody.node) || t.isProgram(parentBody.node)) {
                // 替换 while 语句
                path.replaceWithMultiple(newStatements);
                restoredCount++;
            }
        }
    }
});

console.log(`还原了 ${restoredCount} 处 switch 控制流`);

console.log('\n[3] 清理冗余变量...');

// 清理只用于 switch 控制流的变量
let cleanedVars = 0;
traverse(ast, {
    VariableDeclarator(path) {
        if (!t.isIdentifier(path.node.id)) return;
        const name = path.node.id.name;
        const init = path.node.init;

        // 检查是否是 "x|x|x".split('|') 形式
        if (t.isCallExpression(init) &&
            t.isMemberExpression(init.callee) &&
            t.isStringLiteral(init.callee.object) &&
            /^\d+(\|\d+)+$/.test(init.callee.object.value)) {

            // 检查是否还在使用
            const binding = path.scope.getBinding(name);
            if (binding && binding.references === 0) {
                path.remove();
                cleanedVars++;
            }
        }

        // 检查是否是用于索引的变量 (初始化为 0)
        if (t.isNumericLiteral(init, { value: 0 })) {
            const binding = path.scope.getBinding(name);
            if (binding && binding.references === 0) {
                path.remove();
                cleanedVars++;
            }
        }
    }
});

console.log(`清理了 ${cleanedVars} 个冗余变量`);

console.log('\n[4] 生成代码...');
const output = generate(ast, {
    comments: false,
    compact: false,
    concise: false
}, code);

console.log('\n[5] 保存文件...');
fs.writeFileSync(outputPath, output.code);
console.log(`保存到: ${outputPath}`);
console.log(`新文件大小: ${(output.code.length / 1024).toFixed(2)} KB`);

console.log('\n✅ 完成！');