/*
 * Copyright 20089 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 Sequenic.T2ext.Instrumenter;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;

import org.apache.bcel.Repository;
import org.apache.bcel.classfile.Attribute;
import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.ClassGen;

/**
 * Provide functionality to instrument a class C. Currently this means
 * that all the declared methods of C will be instrumented. 
 * 
 * We can try to extend this in the future, e.g. by instrumenting all
 * classes in the same package.
 * 
 * Each method in-scope will be injected with sensors, and a CFG for the
 * method will be constructed.
 *
 * @author Wishnu Prasetya
 * @author Original code by Maaike Gerritsen
 *
 */
public class ClassInstrumenter {

	
	/**
	 * Full path of the target class.
	 */
	String tpath ;
	
	/**
	 * BCEL representation of a target class C; it allows a class file to
	 * be analyzed, but not transformed.
	 */
	JavaClass original ;
	
	/**
	 * BCEL interface to manipulate a class.
	 */
	ClassGen classGen ;
	
	
	private List<MethodInstrumenter> instrumenters = new LinkedList<MethodInstrumenter>() ;
	
	
	/**
	 * Instrument all methods in a JavaClass, including constructors 
	 * and static initializer, and return the CFGs.
	 * 
	 * Does not require class_path to actually be the current classpath.
	 * 
	 * If already instrumented the .cfgs file is loaded and returned.
	 * 
	 * If doInstrument is false the file system is not modified.
	 * 
	 * @param clazz
	 * @return
	 * @throws ClassNotFoundException 
	 * @throws IOException 
	 */
	public static CFGBunchFile allMethods(String file_name, boolean doInstrument) throws IOException
	{	
		File classFile  = new File(file_name);
		if(!classFile.isFile())
			throw new IllegalArgumentException(classFile.getCanonicalPath()+" is not a file");
		//First load the Java Class file
		JavaClass clazz = new ClassParser(file_name).parse();		
		
		//This gets all the declared methods, including <init> (constructors) and <cinit> (static initializers)
		Method[] m = clazz.getMethods();
		File cfgsFile   = new File(stripClass(file_name) + ".cfgs");
		
		if(isAnyInstrumented(m)) 
		{
			System.out.print(classFile.getCanonicalPath()+" is already instrumented,");
			if(cfgsFile.isFile()) {
				System.out.println(" attempting to load " + cfgsFile.getName());
				return CFGBunchFile.loadAndClose(cfgsFile);
			} else {
				System.out.println("but "+cfgsFile.getName()+" is not a file, returning null.");
				return null;
			}
		}
		
		else 
		{			
			CFGBunchFile CFGs = new CFGBunchFile();
			clazz = allMethods(clazz,CFGs.cfgs,false);
			
			if(doInstrument)
			{
				//Rename original
				classFile.renameTo(new File(file_name + ".original"));
				
				//Dump instrumented version
				clazz.dump(classFile);
				
				//Save CFGs
				CFGs.saveAndClose(cfgsFile);
			}
						
			return CFGs;
		}
	}
	
	/** Instrument all methods in a JavaClass, including constructors 
	 * and static initializer.
	 * 
	 * Returns injected version.
	 * 
	 * If cfgs is not null then the CFGs of the method will be added to it.
	 * 
	 * @param original
	 * @param cfgs
	 * @param debug
	 * @return
	 */
	public static JavaClass allMethods(JavaClass original, Collection<CFG> cfgs, boolean debug) 
	{
		String className = original.getClassName();
		ClassGen cg = new ClassGen(original);
		if(cfgs==null)
			cfgs = new ArrayList<CFG>();
		
		//Instrument all and generate CFGs
		//getMethodsL all the declared methods, including <init> (constructors) and <cinit> (static initializers)
		for(Method mezod : original.getMethods())
			cfgs.add(
					new MethodInstrumenter(mezod, className, cg).setInject().doInstrumenting(debug)
			);
		
		return cg.getJavaClass(); 
	}

	/**
	 * This is the standard constructor of this class. It will set the entire
	 * set of declared methods of the target class as targets to be
	 * instrumented.
	 * 
	 * @param tclass The target class to instrument.
	 * @param path   The full path to the target class.
	 */
	public ClassInstrumenter(Class tclass, String path){
		
		tpath = path ;
		java.lang.reflect.Method[] declaredMethods = tclass.getDeclaredMethods() ;
		List<java.lang.reflect.Method> tmethods = new LinkedList<java.lang.reflect.Method>() ;
		for (int i=0; i<declaredMethods.length; i++) {
			java.lang.reflect.Method M = declaredMethods[i] ;
			//if (! M.getName().equals("main")) 
			tmethods.add(M) ;
		}
		
		// Have BCEL to read the target class and method:
		try {  original = Repository.lookupClass(tclass); }
		catch (Exception e) {
			throw new Error(e) ;
		}			
		// Set-up the class-gen to allow manipulation on the class:
		classGen = new ClassGen(original) ;
		
		for (java.lang.reflect.Method m : tmethods) {
			instrumenters.add(new MethodInstrumenter(original.getMethod(m),tclass.getName(),classGen)) ;
		}
	}
	
	/**
	 * This constructor is just used for testing.
	 */
	ClassInstrumenter(final Class tclass, List<java.lang.reflect.Method> tmethods){
		
		// Have BCEL to read the target class and method:
		try {  original = Repository.lookupClass(tclass); }
		catch (Exception e) {
			throw new Error(e) ;
		}			
		// Set-up the class-gen to allow manipulation on the class:
		classGen = new ClassGen(original) ;
		
		for (java.lang.reflect.Method m : tmethods) {
			instrumenters.add(new MethodInstrumenter(original.getMethod(m),tclass.getName(),classGen)) ;
		}
	}
	
	//Never Actually used...
//	public CFGBunchFile onlyGenerateCFG() {
//		return doInstrumenting(false, false);
//	}
	
	
	public CFGBunchFile doInstrumentingInject() {
		return doInstrumenting(true, false);
	}
	
	public CFGBunchFile doInstrumentingInjectDebug() {
		return doInstrumenting(true, true);
	}
	

	private CFGBunchFile doInstrumenting(boolean doInject, boolean debug)
	{
						
		CFGBunchFile CFGs = new CFGBunchFile() ;
		// Now instrument:
		for (MethodInstrumenter X : instrumenters) 
		{
			if(doInject) 
				X.setInject();
			CFGs.cfgs.add(X.doInstrumenting(debug)) ;
		}
		//Finalize the transformation, and save it in a new class file.
		if (doInject) 
		{
    		try {
    		    classGen.getJavaClass().dump(tpath);
    		  } 
    		catch(Exception e) { 
    			throw new Error(e) ;
    		}		
		}
		String savefile = stripClass(tpath) + ".cfgs";
		CFGs.save(savefile) ;
		System.out.println("\nSaving CFGs in " + savefile + "\n") ;
		
		System.out.println("Instrumenting these methods:") ;
		for (MethodInstrumenter X : instrumenters) 
			System.out.println("  " + X.methodOriginalName) ;
		
		//System.out.println("The CFGs:") ;
		for (CFG G : CFGs.cfgs) {
			G.simplePrint() ;
		}
		
		return CFGs ;
	}
	
	private static String stripClass(String classFile) {
		final String dotClass = ".class";
		if(classFile.endsWith(dotClass))
			return classFile.substring(0, classFile.length() - dotClass.length());
		throw new Error("Class filename does not end with .class: "+classFile);
	}
	
	/**
	 * Check if a Class is already instrumented. Currently only checks for 
	 * occurrences of Sequenic.T2ext.Instrumenter.Sensor.tick in the class file.
	 */
	public static boolean isInstrumented(Class C) {
		JavaClass clazz = null;
		try {
			// TODO check classpath stuff?
			clazz = Repository.lookupClass(C);
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
			System.exit(1);
		}
		return isAnyInstrumented(clazz.getMethods());
	}
	
	/**
	 * Check if a Class is already instrumented. Currently only checks for 
	 * occurrences of Sequenic.T2ext.Instrumenter.Sensor.tick in the class file.
	 */
	private static boolean isAnyInstrumented(Method[] methods)
	{
		for(Method m : methods) {
			Attribute[] attributes = m.getAttributes();
			//System.out.println(m);
			for(Attribute a : attributes) {
				//System.out.println(a);
				// FIXME: maybe a better way to determine if it's already instrumented?
				if(a.toString().indexOf("Sequenic.T2ext.Instrumenter.Sensor.tick") != -1)
					return true;
			}
		}
		return false;
	}
}