package tracer;

import java.io.File;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URL;
import java.net.URLDecoder;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CodeConverter;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.CtConstructor;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.bytecode.BadBytecode;
import javassist.bytecode.analysis.ControlFlow;
import javassist.bytecode.analysis.ControlFlow.Block;
import javassist.expr.ExprEditor;
import javassist.expr.FieldAccess;
import javassist.expr.MethodCall;
import javassist.expr.NewArray;
import javassist.expr.NewExpr;

/**
 * トレース出力文を対象パッケージ内に織り込むクラス
 * 
 * @author Nitta
 *
 */
public class Tracer {
	public static int lineNo = 1;
	public static final String TRACER = "tracer.";
	private static final String STANDARD_CLASSES = "java.util.ListIterator|java.util.Iterator|java.util.List|java.util.Vector|java.util.ArrayList|java.util.Stack|java.util.Map|java.util.HashMap|java.util.Set|java.util.HashSet|java.util.Hashtable|java.util.LinkedList|java.lang.Thread";
	private static final String CONCRETE_STANDARD_CLASSES = "java.util.Vector|java.util.ArrayList|java.util.Stack |java.util.HashMap|java.util.HashSet|java.util.Hashtable|java.util.LinkedList|java.lang.Thread";
	private static final String EXCEPT_FOR_METHODS = "java.lang.Thread.currentThread..|java.lang.Thread.getId..";
	private static final String STANDARD_LIB = "java.";
	private static OutputStatementsGenerator outputStatementsGenerator = null;
	private static CodeConverter conv = new CodeConverter();

	public static void main(String[] args) {
		initialize(new OutputStatementsGenerator(new JSONTraceGenerator()));		// 引数で出力フォーマットを指定する		
		packageInstrumentation("worstCase");										// 指定したパッケージ直下の全クラスにインストゥルメンテーションを行う
	}

	/**
	 * 出力文生成器を指定してインストゥルメンテーションの初期化を行う
	 * 
	 * @param outputStatementsGenerator 出力文生成器(出力フォーマットを指定できる)
	 */
	public static void initialize(OutputStatementsGenerator outputStatementsGenerator) {
		// 配列へのアクセスの検出
		Tracer.outputStatementsGenerator = outputStatementsGenerator;	// 引数で出力フォーマットを指定する
		if (!(Tracer.outputStatementsGenerator.getGenerator() instanceof PlainTextTraceGenerator)) {
			// 
			ClassPool cp = ClassPool.getDefault();
			CtClass cc;
			try {
				cc = cp.get(TRACER + "JSONArrayAdvisor");			// JSONの場合のみ配列アクセスを出力する
				conv.replaceArrayAccess(cc, new CodeConverter.DefaultArrayAccessReplacementMethodNames());
			} catch (NotFoundException e1) {
				e1.printStackTrace();
			}
		}
	}
	
	/**
	 * 初期化済みか?
	 * @return 初期化済み: true, それ以外: false
	 */
	public static boolean isInitialized() {
		return (outputStatementsGenerator != null);
	}

	/**
	 * 指定したパッケージ内のクラスにインストゥルメンテーションを行う
	 * 
	 * @param packageName パッケージ名
	 */
	public static void packageInstrumentation(String packageName) {
		ClassLoader loader = Thread.currentThread().getContextClassLoader();
		URL resource = loader.getResource(packageName);
		File dir;
		try {
			dir = new File(URLDecoder.decode(resource.getPath(), "UTF-8"));
			for (String file : dir.list()) {
				if (file.endsWith(".class")) {
					String className = packageName + "." + file.replace(".class", "");
					classInstrumentation(className);
				}
			}
		} catch (UnsupportedEncodingException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 指定したクラスにインストゥルメンテーションを行う
	 * 
	 * @param className クラス名
	 */
	public static void classInstrumentation(String className) {
		ClassPool cp = ClassPool.getDefault();
		CtClass cc;
		try {
			cc = cp.get(className);
			classInstrumentation(cc);
		} catch (NotFoundException | BadBytecode | CannotCompileException | IOException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 指定したクラスにインストゥルメンテーションを行う
	 * 
	 * @param cc Javassistのクラスオブジェクト
	 */
	public static void classInstrumentation(CtClass cc) throws BadBytecode, NotFoundException, CannotCompileException, IOException {
		for (final CtConstructor c : cc.getDeclaredConstructors()) {
			methodInstrumentation(cc, c);				
		}
		for (final CtMethod m : cc.getDeclaredMethods()) {
			methodInstrumentation(cc, m);
		}
		cc.instrument(conv);
		cc.writeFile("bin");
	}

	private static void methodInstrumentation(final CtClass cc, final CtBehavior m) throws BadBytecode, NotFoundException, CannotCompileException {
		// メソッド本体内の各ブロックの最初に出力文を挿入する(出力文の挿入でブロックが増えてしまうので、先に挿入しておく)
		Block[] blocks = null;
		if (m instanceof CtMethod && !m.isEmpty()) {
			ControlFlow cf = new ControlFlow((CtMethod)m);
			blocks = cf.basicBlocks();
			for (int i = blocks.length - 1; i >= 0; i--) {
				int blockPos = m.getMethodInfo().getLineNumber(blocks[i].position());
				m.insertAt(blockPos, outputStatementsGenerator.generateInsertStatementsForBlockEntry((CtMethod)m, i, blocks[i], blockPos));
			}
		}
		
		// メソッド本体内のフィールドアクセスとメソッド呼び出しを置き換える
		m.instrument(new ExprEditor() {
			public void edit(FieldAccess f) throws CannotCompileException {
				if (f.isReader()) {
					if (!f.getFieldName().contains("$")) {		// AspectJでは final local 変数からのゲットは無視されるので、それに合わせて除外する
						f.replace(outputStatementsGenerator.generateReplaceStatementsForFieldGet(cc, m, f, f.getLineNumber()));
					}
				} else {
					if (!f.getFieldName().contains("$")) {		// この条件がないとなぜか落ちる（無名フィールド?へのセットがあって、それを拾って落ちてる?）
						f.replace(outputStatementsGenerator.generateReplaceStatementsForFieldSet(cc, f, f.getLineNumber()));
					}
				}
			}
			public void edit(MethodCall c) throws CannotCompileException {
				try {
					CtMethod m = c.getMethod();
					String className = m.getDeclaringClass().getName();
					if (!className.startsWith(STANDARD_LIB) && !className.startsWith(TRACER)) {
						// 通常のメソッドの呼び出し
						c.replace(outputStatementsGenerator.generateReplaceStatementsForCall(m.getDeclaringClass(), m, c.getLineNumber(), true));
					} else if (className.matches(STANDARD_CLASSES) && !m.getLongName().matches(EXCEPT_FOR_METHODS)) {
						// 標準クラスのメソッド呼び出し（呼び出し先にバイトコードを埋め込めない）							
						c.replace(outputStatementsGenerator.generateReplaceStatementsForCall(m.getDeclaringClass(), m, c.getLineNumber(), false));
					}
				} catch (NotFoundException e) {
					e.printStackTrace();
				}
			}
			public void edit(NewExpr n) throws CannotCompileException {
				try {
					CtConstructor m = n.getConstructor();
					String className = m.getDeclaringClass().getName();
					if (!className.startsWith(STANDARD_LIB) && !className.startsWith(TRACER)) {
						// 通常のコンストラクタの呼び出し
						n.replace(outputStatementsGenerator.generateReplaceStatementsForCall(m.getDeclaringClass(), m, n.getLineNumber(), true));
					} else if (m.getDeclaringClass().getName().matches(CONCRETE_STANDARD_CLASSES)) {
						// 標準クラスのコンストラクタ呼び出し（呼び出し先にバイトコードを埋め込めない）							
						n.replace(outputStatementsGenerator.generateReplaceStatementsForCall(m.getDeclaringClass(), m, n.getLineNumber(), false));
					}
				} catch (NotFoundException e) {
					e.printStackTrace();
				}
			}
			public void edit(NewArray a) throws CannotCompileException {
				a.replace(outputStatementsGenerator.generateReplaceStatementsForNewArray(a, a.getLineNumber()));
			}
//			public void edit(ConstructorCall c) throws CannotCompileException {
//				try {
//					CtConstructor m = c.getConstructor();
//					if (m.getDeclaringClass().getName().matches(CONCRETE_STANDARD_CLASSES)) {
//						c.replace(generateReplaceStatementsForCall(m.getDeclaringClass(), m));
//					}
//				} catch (NotFoundException e) {
//					e.printStackTrace();
//				}
//			}
		});
		
		// メソッド用の出力文を生成する
		if (!m.isEmpty()) {
			// メソッドの実行前後に出力文を挿入する
			m.insertBefore(outputStatementsGenerator.generateInsertBeforeStatementsForMethodBody(cc, m));
			m.insertAfter(outputStatementsGenerator.generateInsertAfterStatementsForMethodBody(cc, m));
		} else {
			// メソッド本体が空のときはコンストラクタの場合のみ(=デフォルトコンストラクタ)本体に出力文を設定する
			if (m instanceof CtConstructor) {
				m.setBody(outputStatementsGenerator.generateInsertAfterStatementsForMethodBody(cc, m));
				m.insertBefore(outputStatementsGenerator.generateInsertBeforeStatementsForMethodBody(cc, m));
			}
		}
	}
}