package alterrs.jbct.deob;

import static alterrs.asm.Opcodes.ATHROW;
import static alterrs.asm.Opcodes.NEW;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;

import alterrs.asm.tree.AbstractInsnNode;
import alterrs.asm.tree.ClassNode;
import alterrs.asm.tree.LabelNode;
import alterrs.asm.tree.MethodNode;
import alterrs.asm.tree.TryCatchBlockNode;
import alterrs.asm.tree.TypeInsnNode;
import alterrs.commons.LoggingUtil;
import alterrs.jbct.JBCT;
import alterrs.jbct.utility.AbstractTransformer;
import alterrs.jbct.utility.ID;

/**
 * @author Lazaro Brito
 */
public class TryCatchDeobfuscator extends AbstractTransformer {
	private static final Logger logger = LoggingUtil.log();

	private int count = 0, count2 = 0;

	public TryCatchDeobfuscator() {
		super(true);
	}

	@Override
	public void visitMethod(ClassNode c, MethodNode m) {
		List<TryCatchBlockNode> removeList = new LinkedList<TryCatchBlockNode>();

		tryCatchLoop: for (Object tcbo : m.tryCatchBlocks) {
			TryCatchBlockNode tcb = (TryCatchBlockNode) tcbo;
			if(removeList.contains(tcb)) {
				continue;
			}
			
			if (tcb.type != null
					&& tcb.type.equals("java/lang/RuntimeException")) {
				boolean stringBuilder = false;

				AbstractInsnNode insn = tcb.handler;
				while (true) {
					if (insn == null) {
						continue tryCatchLoop;
					}

					if (insn.getOpcode() == NEW) {
						if (((TypeInsnNode) insn).desc
								.equals("java/lang/StringBuilder")) {
							stringBuilder = true;
							break;
						}
					}

					if (insn.getOpcode() == ATHROW || m.instructions.indexOf(insn) > m.instructions.indexOf(tcb.end)) {
						break;
					} else {
						insn = insn.getNext();
					}
				}

				if (!stringBuilder) {
					continue tryCatchLoop;
				}

				removeList.add(tcb);
				for (Object tcbo2 : m.tryCatchBlocks) {
					TryCatchBlockNode tcb2 = (TryCatchBlockNode) tcbo2;
					if(!removeList.contains(tcb2) && tcb2.handler == tcb.handler) {
						removeList.add(tcb);
					}
				}

				insn = tcb.handler;
				while (insn != null) {
					AbstractInsnNode nextInsn = insn.getNext();
					m.instructions.remove(insn);

					if (insn.getOpcode() == ATHROW) {
						break;
					} else {
						insn = nextInsn;
					}
				}

				count++;
			}
		}
		for (TryCatchBlockNode tcb : removeList) {
			m.tryCatchBlocks.remove(tcb);
		}

		Map<Integer, List<TryCatchBlockNode>> tryCatchHandlerPositions = new HashMap<Integer, List<TryCatchBlockNode>>();
		for (Object tcbo : m.tryCatchBlocks) {
			TryCatchBlockNode tcb = (TryCatchBlockNode) tcbo;

			int handlerPos = m.instructions.indexOf(tcb.handler);
			if (!tryCatchHandlerPositions.containsKey(handlerPos)) {
				List<TryCatchBlockNode> handlers = new LinkedList<TryCatchBlockNode>();
				handlers.add(tcb);

				tryCatchHandlerPositions.put(handlerPos, handlers);
			} else {
				tryCatchHandlerPositions.get(handlerPos).add(tcb);
			}
		}

		for (Map.Entry<Integer, List<TryCatchBlockNode>> entry : tryCatchHandlerPositions
				.entrySet()) {
			List<TryCatchBlockNode> tryCatchHandlers = entry.getValue();
			if (tryCatchHandlers.size() > 1) {
				int startPos = Integer.MAX_VALUE;
				int endPos = 0;

				LabelNode handler = tryCatchHandlers.get(0).handler;
				String type = tryCatchHandlers.get(0).type;

				for (TryCatchBlockNode tcb : tryCatchHandlers) {
					int startPos_ = m.instructions.indexOf(tcb.start);
					if (startPos_ < startPos) {
						startPos = startPos_;
					}

					int endPos_ = m.instructions.indexOf(tcb.end);
					if (endPos_ > endPos) {
						endPos = endPos_;
					}

					m.tryCatchBlocks.remove(tcb);
					
				}

				m.tryCatchBlocks.add(new TryCatchBlockNode(
						(LabelNode) m.instructions.get(startPos),
						(LabelNode) m.instructions.get(endPos), handler, type));
				count2++;
			}
		}
	}

	@Override
	public void onFinish() {
		logger.info("Removed " + count + " faulty try-catches and fixed "
				+ count2 + " try-catch tables!");
	}
}
