/*
 * Copyright 2008 Wishnu Prasetya.
 *
 * This file is part of T2.
 * T2 is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License (GPL) as published by the
 * Free Software Foundation; either version 3 of the License, or any
 * later version.
 * 
 * T2 is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * for more details.
 *
 * A copy of the GNU General Public License can be found in T2 distribution.
 * If it is missing, see http://www.gnu.org/licenses.
 */
package t2ext.bcel;

import java.io.File;
import java.io.FileOutputStream;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;

import t2ext.bcel.utils.MethodCallInsertion;
import t2ext.bcel.utils.Node;
import t2ext.util.InfoFile;
import t2ext.util.Util;
import t2ext.xml.XMLBuilder;

import com.sun.org.apache.bcel.internal.Repository;
import com.sun.org.apache.bcel.internal.classfile.JavaClass;
import com.sun.org.apache.bcel.internal.classfile.LineNumberTable;
import com.sun.org.apache.bcel.internal.classfile.Method;
import com.sun.org.apache.bcel.internal.generic.ATHROW;
import com.sun.org.apache.bcel.internal.generic.BranchHandle;
import com.sun.org.apache.bcel.internal.generic.BranchInstruction;
import com.sun.org.apache.bcel.internal.generic.ClassGen;
import com.sun.org.apache.bcel.internal.generic.ConstantPoolGen;
import com.sun.org.apache.bcel.internal.generic.GotoInstruction;
import com.sun.org.apache.bcel.internal.generic.IfInstruction;
import com.sun.org.apache.bcel.internal.generic.Instruction;
import com.sun.org.apache.bcel.internal.generic.InstructionHandle;
import com.sun.org.apache.bcel.internal.generic.InstructionList;
import com.sun.org.apache.bcel.internal.generic.MethodGen;
import com.sun.org.apache.bcel.internal.generic.ReturnInstruction;

/**
 * 
 * 
 * @author Maaike Gerritsen
 */
public class ControlFlow {
	private Hashtable<InstructionHandle, Integer> _instructionHandleToId;
	private Hashtable<Integer, InstructionHandle> _idToInstructionHandle;
	private Hashtable<Integer, Node> _idToNode;
	private ClassGen _classGen;
	private MethodCallInsertion _callMethods;
	private final String Extension = "_id";
	private Hashtable<String, List<List<Integer>>> _paths;
	private String _byteClass;
	private XMLBuilder _builder;
	private InfoFile _infoFile;
	private static final String Aux = "aux_";

	@SuppressWarnings("unchecked")
	public ControlFlow(String byteClass, InfoFile infoFile) {
		_byteClass = byteClass;
		_infoFile = infoFile;
		try {
			Class myClass = Class.forName(byteClass);
			JavaClass clazz = Repository.lookupClass(myClass);
			_classGen = new ClassGen(clazz);
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		}
		_instructionHandleToId = new Hashtable<InstructionHandle, Integer>();
		_idToInstructionHandle = new Hashtable<Integer, InstructionHandle>();
		_idToNode = new Hashtable<Integer, Node>();
		_callMethods = new MethodCallInsertion(_classGen);
		_paths = new Hashtable<String, List<List<Integer>>>();
		_builder = new XMLBuilder();
		_builder.setClassInfoName(_byteClass);
	}

	public void createControlFlow() {
		Method[] methods = _classGen.getJavaClass().getMethods();
		Hashtable<String, String> officialMethodNames = new Hashtable<String, String>();
		for (int i = 0; i < methods.length; i++) {
			_idToNode = new Hashtable<Integer, Node>();
			Method method = methods[i];
			LineNumberTable lines = method.getLineNumberTable();
			String methodName = method.getName();
			
			boolean allowed = checkIfMethodIstoBeUsed(method);
			if (allowed) {
				// System.out.println("----------------");
				//System.out.println("Processing "+methodName);
				
				// Create a new method based on the existing one
				ConstantPoolGen pgen = _classGen.getConstantPool();
				String cname = _classGen.getClassName();
				MethodGen methodGen = new MethodGen(method, cname, pgen);
				
				InstructionList list = methodGen.getInstructionList();
				if (list != null) {
					// Make nodes for the instructions in this method
					setIdToHandle(list);
					
					String methodAndParam = Util.completeMethodName(methodName,
							methodGen.getArgumentTypes());
					String officialName = Util.officialMethodName(_byteClass,
							methodName, methodGen.getArgumentTypes());
					officialMethodNames.put(methodAndParam, officialName);
					
					_builder.registerMethodsForClassInfo(methodName, methodGen
							.getArgumentTypes());
					
					sortInstructionHandles(list, methodAndParam, lines);
					
					// Replace the original method with the newly created one
					methodGen.setInstructionList(createNewList(list));
					_classGen.removeMethod(method);
					_classGen.addMethod(methodGen.getMethod());
					list.dispose();
				}
			}
		}
		// Save the new class... 
		writeToFile();
		// and classInfo.xml
		_builder.createClassInfoFile();
		// and paths.xml
		_builder.createPathFile(getPaths(), officialMethodNames);
	}

	private boolean checkIfMethodIstoBeUsed(Method method) {
		String name = method.getName();
		
		// Why these < and add checks?
		if(name.contains("<") || name.contains("add")) {
			return false;
		}
		if (_infoFile.noArgumentsPresent) {
			if (method.isPublic())
				return true;
			return false;
		}
		if (!_infoFile.methodsToBeChecked.isEmpty()) {
			if (_infoFile.methodsToBeChecked.contains(name))
				return true;
			return false;
		} 
		if (!_infoFile.methodsToBeExcluded.isEmpty()) {
			if (_infoFile.methodsToBeExcluded.contains(name))
				return false;
			// FIXME this seems wrong: if there's an exclude list it doesn't 
			// mean that everything that's not on it automatically should be
			// checked, does it?
			return true;
		}
		if (method.isPublic())
			return true;
		if (method.isStatic() && _infoFile.exclStatic)
			return false;
		if (method.isPrivate() && !_infoFile.inclPrivate)
			return false;
		if (method.isProtected() && _infoFile.exclProtected)
			return false;
		if (!method.isPublic() && _infoFile.publicOnly)
			return false;
		return true;
	}

	private void sortInstructionHandles(InstructionList list,
			String methodName, LineNumberTable lines) {
		createControlFlowGraph(list, methodName, lines);
		rearrangeIfListsToInsert(lines, list);
		PrimePathCalculator pmc = new PrimePathCalculator(_idToNode);
		_paths.put(methodName, pmc.getPrimePaths());
	}

	@SuppressWarnings("unchecked")
	private void rearrangeIfListsToInsert(LineNumberTable lines,
			InstructionList list) {
		for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
				.hasNext();) {
			InstructionHandle handle = iterator.next();
			int id = _instructionHandleToId.get(handle);
			Node node = _idToNode.get(id);
			if(node.isBranchNode())
				setListToInsertForIfNode(handle, lines, node.getListToInsert());
		}
	}

	@SuppressWarnings("unchecked")
	private void createControlFlowGraph(InstructionList list,
			String methodName, LineNumberTable lines) {
		for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
				.hasNext();) {
			InstructionHandle handle = iterator.next();
			int id = _instructionHandleToId.get(handle);
			Node node = _idToNode.get(id);
			int lineNr = lines
					.getSourceLine(node.getNodeHandle().getPosition());
			node.setLineNr(lineNr);
			Instruction instr = handle.getInstruction();
			if (id == list.size() - 1 || instr instanceof ReturnInstruction
					|| instr instanceof ATHROW) // found last instruction
				node.setIsLast(true);
			else if (instr instanceof GotoInstruction) {
				node = createTargetChild(handle, node, true);
			} else if (instr instanceof IfInstruction) {
				Node root = _idToNode.get(0);
				if (root.getListToInsert() == null) {
					String nameRoot = Aux + methodName + Extension + 0
							+ Extension;
					root.setListToInsert(_callMethods.callMethod(nameRoot));
				}
				node.setIsBranchNode(true);
				node = createTargetChild(handle, node, false);
				node = createNextChild(id, node);
				int leftId = node.getChildLeft().getId();
				int rightId = node.getChildRight().getId();
				String nameLeft = Aux + methodName + Extension + leftId
						+ Extension;
				String nameRight = Aux + methodName + Extension + rightId
						+ Extension;
				String nodeName = Aux + methodName + Extension + id + Extension;
				node.setListToInsert(_callMethods.callMethod(nodeName));
				int lineNrParent = lines.getSourceLine(node.getNodeHandle()
						.getPosition());
				int lineNrChildLeft = lines.getSourceLine(node.getChildLeft()
						.getNodeHandle().getPosition());
				int lineNrChildRight = lines.getSourceLine(node.getChildRight()
						.getNodeHandle().getPosition());
				if (lineNrParent != lineNrChildLeft
						|| lineNrChildLeft == lineNrChildRight)
					node.getChildLeft().setListToInsert(
							_callMethods.callMethod(nameLeft));
				if (lineNrChildLeft == lineNrChildRight) {
					node.getChildLeft().addToCFG();
					node.getChildRight().addToCFG();
				} else if (node.getChildLeft().getNodeHandle().getInstruction() instanceof GotoInstruction)
					node.getChildLeft().setListToInsert(
							_callMethods.callMethod(nameLeft));
				node.getChildRight().setListToInsert(
						_callMethods.callMethod(nameRight));
			} else {
				node = createNextChild(id, node);
			}
			_idToNode.put(id, node);
		}
	}

	private void setListToInsertForIfNode(InstructionHandle handle,
			LineNumberTable lines, InstructionList callMethod) {
		int lineNr = lines.getSourceLine(handle.getPosition());
		InstructionHandle previousHandle = handle.getPrev();
		int previousLineNr = lines.getSourceLine(previousHandle.getPosition());
		while (previousHandle != null && !isBranchHandle(previousHandle)
				&& previousLineNr == lineNr) {
			handle = previousHandle;
			int prevId = _instructionHandleToId.get(handle);
			Node prevNode = _idToNode.get(prevId);
			if(prevNode.isTargetOfGoto())
				break;
			previousHandle = previousHandle.getPrev();
			if (previousHandle != null)
				previousLineNr = lines.getSourceLine(previousHandle
						.getPosition());
		}
		int id = _instructionHandleToId.get(handle);
		Node node = _idToNode.get(id);
		if (node.getListToInsert() == null)
			node.setListToInsert(callMethod);
		else {
			InstructionList list = node.getListToInsert();
			list.append(callMethod);
			node.setListToInsert(list);
		}
		_idToNode.put(id, node);
	}

	private boolean isBranchHandle(InstructionHandle previousHandle) {
		if (previousHandle.getInstruction() instanceof BranchInstruction)
			return true;
		return false;
	}

	private Node createNextChild(int id, Node node) {
		Node child = _idToNode.get(new Integer(id + 1));
		node.setChildLeft(child);
		child.addToParents(node);
		_idToNode.put(id + 1, child);
		return node;
	}

	private Node createTargetChild(InstructionHandle handle, Node node,
			boolean leftChild) {
		InstructionHandle target = ((BranchHandle) handle).getTarget();
		int childId = _instructionHandleToId.get(target);
		Node child = _idToNode.get(childId);
		child.addToParents(node);
		if (childId < node.getId()) {
			child.setCycleBegin(true);
			node.setCycleEnd(true);
		}
		if (leftChild)
			node.setChildLeft(child);
		else {
			node.setChildRight(child);
			child.setIsRightChild(true);
		}
		if (handle.getInstruction() instanceof GotoInstruction)
			child.targetOfGoto();
		_idToNode.put(childId, child);
		return node;
	}

	@SuppressWarnings("unchecked")
	private InstructionList createNewList(InstructionList list) {
		for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
				.hasNext();) {
			InstructionHandle handle = iterator.next();
			int id = _instructionHandleToId.get(handle);
			Node node = _idToNode.get(id);
			if (node != null && node.getListToInsert() != null) {
				InstructionList newList = node.getListToInsert();
				InstructionHandle newHandle = newList.getStart();
				list.insert(handle, newList);
				// if (node.isRightChild())
				// if(node.hasIfParent())
				list.redirectBranches(handle, newHandle);
			}
		}
		// System.out.println("LIST -------------");
		// for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
		// .hasNext();) {
		// Instruction instr = iterator.next().getInstruction();
		// System.out.println(instr.toString());
		// }
		// System.out.println("LIST -------------");
		list.setPositions();
		return list;
	}

	public Node getNode() {
		return _idToNode.get(0);
	}

	public Hashtable<String, List<List<Integer>>> getPaths() {
		return _paths;
	}

	private void writeToFile() {
		File file = Util.getClassLocation(_byteClass);
		System.out.println("Overwriting class "+file);
		try {
			FileOutputStream fos = new FileOutputStream(file);
			_classGen.getJavaClass().dump(fos);
			fos.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * Makes Nodes for each of the instructions in the list.
	 * @param list The list with instructions
	 */
	@SuppressWarnings("unchecked")
	private void setIdToHandle(InstructionList list) {
		int counter = 0;
		_instructionHandleToId = new Hashtable<InstructionHandle, Integer>();
		_idToInstructionHandle = new Hashtable<Integer, InstructionHandle>();
		Iterator<InstructionHandle> iterator = list.iterator();
		while(iterator.hasNext()) {
			InstructionHandle handle = iterator.next();
			_instructionHandleToId.put(handle, counter);
			_idToInstructionHandle.put(counter, handle);
			Node node = new Node();
			node.setId(counter);
			node.setNodeHandle(handle);
			_idToNode.put(counter, node);
			counter++;
		}
	}

	@SuppressWarnings("unchecked")
	public void printInstructions() {
		Method[] methods = _classGen.getJavaClass().getMethods();
		for (int i = 0; i < methods.length; i++) {
			Method method = methods[i];
			LineNumberTable lines = method.getLineNumberTable();
			String name = method.getName();
			System.out.println("METHOD NAME: " + name);
			// if ("maaike".equals(name.trim())) {
			ConstantPoolGen pgen = _classGen.getConstantPool();
			String cname = _classGen.getClassName();
			MethodGen methodGen = new MethodGen(method, cname, pgen);
			InstructionList list = methodGen.getInstructionList();
			System.err.println("++++ Instructions ++++");
			int counter = 0;
			for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
					.hasNext();) {
				InstructionHandle handle = iterator.next();
				Instruction instr = handle.getInstruction();
				// System.out.println(counter + ": " + instr.toString());
				int bla = -1;
				if (lines != null)
					try{
						bla = lines.getSourceLine(handle.getPosition());
					}catch(Exception e){
						System.out.println(e.getMessage());
					}
				System.out.println(counter + ": " + instr.toString() + "\t "
						+ bla);
				counter++;
				// System.out.println("at line: "
				// + lines.getSourceLine(handle.getPosition()));
			}
			// }
		}
	}

	@SuppressWarnings("unchecked")
	public void printInstructions2() {
		Method[] methods = _classGen.getJavaClass().getMethods();
		for (int i = 0; i < methods.length; i++) {
			Method method = methods[i];
			@SuppressWarnings("unused")
			LineNumberTable lines = method.getLineNumberTable();
			String name = method.getName();
			System.out.println("METHOD NAME: " + name);
			// if ("maaike".equals(name.trim())) {
			ConstantPoolGen pgen = _classGen.getConstantPool();
			String cname = _classGen.getClassName();
			MethodGen methodGen = new MethodGen(method, cname, pgen);
			InstructionList list = methodGen.getInstructionList();
			System.err.println("++++ Instructions ++++");
			int counter = 0;
			for (Iterator<InstructionHandle> iterator = list.iterator(); iterator
					.hasNext();) {
				InstructionHandle handle = iterator.next();
				Instruction instr = handle.getInstruction();
				System.out.println(counter + ": " + instr.toString());
				// int bla = -1;
				// if (lines != null)
				// bla = lines.getSourceLine(handle.getPosition());
				// System.out.println(counter + ": " + instr.toString() + "\t "
				// + bla);
				counter++;
				// System.out.println("at line: "
				// + lines.getSourceLine(handle.getPosition()));
			}
			// }
		}
	}
}