const parser = require('@babel/parser');
const traverse = require('@babel/traverse').default;
const fs = require('fs');

const code = `
let state = 0, x = 5;
while (1) {
    switch (state) {
        case 0:
            console.log('执行块 0 操作')
            state = 1;
            break;
        case 1:
            state = (x > 0) ? 2 : 3;
            console.log('执行块 1 操作')
            break;
        case 2:
            console.log('执行块 2 操作')
            state = 4;
            break;
        case 3:
            console.log('执行块 3 操作')
            state = 4;
            break;
        case 4:
            console.log('程序执行完毕');
    }
}
`;

const ast = parser.parse(code, { sourceType: 'module' });

let initialState = null;
let switchCases = {};

traverse(ast, {
  VariableDeclarator(path) {
    if (path.node.id.name === 'state') {
      initialState = path.node.init.value;
    }
  },
  SwitchCase(path) {
    const test = path.node.test.value;
    const body = path.node.consequent;
    let logs = [];
    let nextState = null;

    body.forEach(node => {
      if (
        node.type === 'ExpressionStatement' &&
        node.expression.type === 'CallExpression' &&
        node.expression.callee.object.name === 'console'
      ) {
        logs.push(node.expression.arguments[0].value);
      }

      if (
        node.type === 'ExpressionStatement' &&
        node.expression.type === 'AssignmentExpression' &&
        node.expression.left.name === 'state'
      ) {
        nextState = node.expression.right;
      }
    });

    switchCases[test] = { logs, nextState };
  }
});

function formatNextState(nextState) {
  if (nextState.type === 'ConditionalExpression') {
    return `if (${generateCode(nextState.test)}) {
  ${switchCases[nextState.consequent.value].logs.map(log => `console.log('${log}')`).join('\n  ')}
} else {
  ${switchCases[nextState.alternate.value].logs.map(log => `console.log('${log}')`).join('\n  ')}
}`;
  } else if (nextState.type === 'NumericLiteral') {
    const logs = switchCases[nextState.value]?.logs || [];
    return logs.map(log => `console.log('${log}')`).join('\n');
  } else {
    return '// 未识别的跳转';
  }
}

function generateCode(node) {
  if (node.type === 'BinaryExpression') {
    return `${node.left.name} ${node.operator} ${node.right.value}`;
  }
  return '// 未识别表达式';
}

// 🌳 生成结构化代码
let outputCode = '';
outputCode += switchCases[0].logs.map(log => `console.log('${log}')`).join('\n') + '\n';

outputCode += formatNextState(switchCases[1].nextState) + '\n';
outputCode += switchCases[4].logs.map(log => `console.log('${log}')`).join('\n');

console.log('\n✅ 结构化还原后的代码：\n');
console.log(outputCode);