package summary;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import driver.PTAOptions;
import driver.SootUtils;
import pta.PTA;
import pta.SubPTAInterface;
import soot.Body;
import soot.DoubleType;
import soot.FloatType;
import soot.IntegerType;
import soot.Local;
import soot.LongType;
import soot.MethodSource;
import soot.PatchingChain;
import soot.RefType;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Trap;
import soot.Type;
import soot.Unit;
import soot.Value;
import soot.ValueBox;
import soot.VoidType;
import soot.jimple.AssignStmt;
import soot.jimple.CaughtExceptionRef;
import soot.jimple.DoubleConstant;
import soot.jimple.FloatConstant;
import soot.jimple.IdentityStmt;
import soot.jimple.IntConstant;
import soot.jimple.Jimple;
import soot.jimple.JimpleBody;
import soot.jimple.LongConstant;
import soot.jimple.NullConstant;
import soot.jimple.ParameterRef;
import soot.jimple.ReturnStmt;
import soot.jimple.ReturnVoidStmt;
import soot.jimple.Stmt;
import soot.jimple.ThisRef;
import util.Debug;
import util.PTATimer;

public class Summarizer {

	public static void run() {
		System.out.println("now loading methods...");
		PTATimer.start("loading methods");
		Collection<SootMethod> methods = null;
		if(PTAOptions.summethod!=null) {
			SootMethod method = Scene.v().getMethod(PTAOptions.summethod);
			PTATimer.stop("loading methods");
			PTATimer.start("summary");
			summarize(method);
		}else {
			new PTA().purerun(); // prepare methodstmts
			PTATimer.stop("loading methods");
			// load to-summarize methods
			methods = new AppCallLog(PTAOptions.SUMMARY_LOG).getMethods();
			System.out.println("#method to summarize:"+methods.size());
			// for each to-summarize method: summarize
			PTATimer.start("summary");
			switch (PTAOptions.cores) {
			case PTAOptions.FOREACH:
				summarize(methods);
				break;
			case PTAOptions.STREAM:
				streamSummarize(methods);
				break;
			case PTAOptions.PARALLELSTREAM:
				parallelSummarize(methods);
				break;
			default:
				manullyParallelSummarize(methods,PTAOptions.cores);
			}
		}		
		
		PTATimer.stop("summary");
		// write bytecode
		System.out.println("now writing classes...");
		PTATimer.start("writing classes");
		String analysis = PTAOptions.subptaPattern.toString();
		String dir = PTAOptions.APP_PATH.replace(".jar", "")+"-deps_"+analysis;
		for (SootClass clz : Scene.v().getLibraryClasses()) {
			try {
				SootUtils.writeByteCode(dir, clz);
			}catch (Throwable ve) {
				Debug.println(ve);
			}
		}
		PTATimer.stop("writing classes");
		// prepare jars
		System.out.println("now packaging jars...");
		PTATimer.start("packaging jars");
		JarFile librJarFile = new JarFile(PTAOptions.APP_PATH.replace(".jar", "")+"-deps_"+analysis+".jar");
        try {
			librJarFile.addGeneratedLibraryClassFiles();
		} catch (IOException e) {
			e.printStackTrace();
		}
        PTATimer.stop("packaging jars");
        deleteFolder(new File(dir));
        System.out.println("done.");
	}
	
	public static void deleteFolder(File folder) {
        File[] files = folder.listFiles();
        if (files != null) {
            for (File f : files) {
                if (f.isDirectory()) {
                    deleteFolder(f);
                } else {
                    f.delete();
                }
            }
        }
        folder.delete();
    }
	
	private static void summarize(Collection<SootMethod> methods){
		System.out.println("iterator.");
		methods.forEach(Summarizer::summarize);
	}
	private static void streamSummarize(Collection<SootMethod> methods){
		System.out.println("stream.");
		methods.stream().forEach(Summarizer::summarize);
	}
	private static void parallelSummarize(Collection<SootMethod> methods) {
		System.out.println("parallel stream.");
		methods.parallelStream().forEach(Summarizer::summarize);
	}
	private static void manullyParallelSummarize(Collection<SootMethod> methods, int numThreads) {
		if(numThreads == 0)
			numThreads = Runtime.getRuntime().availableProcessors();
		System.out.println("manul parallel (numThreads: "+numThreads+").");
		ExecutorService executorService = Executors.newFixedThreadPool(numThreads);
		methods.stream().forEach(method ->
			executorService.execute(() -> {
		        summarize(method);
		    }));
		executorService.shutdown();
		try {
		    executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
		} catch (InterruptedException e) {
		    throw new RuntimeException(e);
		}
	}

	static void summarize(SootMethod method){
		SubPTAInterface subpta = SubPTAInterface.createSubPTA(method);
		Debug.println(subpta);
		
		((PTA) subpta).purerun(); // run lib-based PTA for the method
		
		Collection<Stmt> stmts = subpta.makeSummary(); // make summary for the method
		
		if(!method.isConcrete()) return;
		
		if (method.toString().contains("<java.lang.System: void <clinit>()>")) {
			System.out.println(stmts);
		}
		
		List<Stmt> stmtList = new ArrayList<>(stmts);
		Collections.sort(stmtList, new Comparator<Stmt>() {

			@Override
			public int compare(Stmt o1, Stmt o2) {
				return lvl(o1)-lvl(o2);
			}

			private int lvl(Stmt o1) {
				if(o1 instanceof IdentityStmt)
					return 0;
				else if (o1 instanceof AssignStmt) 
					return 1;
				return 0;
			}
			
		});
		
		JimpleBody body = replaceBody(method); // replace jimple body
		
		validateAddAll(body,stmts); // add stmts and ensure validity
	}
	
	private static JimpleBody replaceBody(SootMethod method) {
		JimpleBody body = Jimple.v().newBody(method);
		method.setActiveBody(body);
		method.setSource(new MethodSource() {
			@Override
			public Body getBody(SootMethod m, String phaseName) {
				return body;
			}
		});
		return body;
	}


	private static void validateAddAll(JimpleBody body, Collection<Stmt> stmts) {
		PatchingChain<Unit> units = body.getUnits();
		
    	Map<Integer, Unit> parameterRefs = new HashMap<>();
    	Set<IdentityStmt> catches = new HashSet<>();
    	boolean hasReturn = false;
		for(Stmt stmt:stmts) {
			if(stmt instanceof IdentityStmt istmt) {
				Value v = istmt.getRightOp();
	    		if (v instanceof ParameterRef ref)
	    			parameterRefs.put(ref.getIndex(), stmt);
	    		else if(v instanceof ThisRef ref)// thisRef
	    			parameterRefs.put(-1, stmt);
	    		else if(v instanceof CaughtExceptionRef ref) {// catch
	    			catches.add(istmt);
	    			units.add(stmt);
	    		}else
	    			throw new RuntimeException("unhanlded identities:"+v.getClass());		
			}else {
				units.add(stmt);
				if(stmt instanceof ReturnStmt || stmt instanceof ReturnVoidStmt)
					hasReturn = true;
			}
		}
    	
		
		insertIdentities(body,parameterRefs);
		
		createDummyTraps(body,catches);
		
		if(!hasReturn)
			addReturnStmt(body); // add return stmt
		
		registerLocals(body); // register used locals for declaration
	}

	private static void insertIdentities(JimpleBody body, Map<Integer, Unit> parameterRefs) {
		PatchingChain<Unit> units = body.getUnits();
		SootMethod method = body.getMethod();
		
		Unit lastUnit = null;
    	
        //add this-ref before everything else
        if (!method.isStatic()){
        	Unit s = parameterRefs.get(-1);
        	if(s==null) {
        		Local l = Jimple.v().newLocal("this", RefType.v(method.getDeclaringClass()));
        		s = Jimple.v().newIdentityStmt(l, Jimple.v().newThisRef((RefType)l.getType()));
        	}
        	units.addFirst(s);
        	lastUnit = s;
        }
        
        int i = 0;
        for (Type t : method.getParameterTypes()) {
        	Unit s = parameterRefs.get(i);
        	if(s==null) {
        		Local l = Jimple.v().newLocal("parameter"+i, t);
                s = Jimple.v().newIdentityStmt(l, Jimple.v().newParameterRef(l.getType(), i));
        	}
            
            if (lastUnit == null)
            	units.addFirst(s);
            else
            	units.insertAfter(s, lastUnit);
            
            lastUnit = s;
            i++;
        }
	}

	private static void createDummyTraps(JimpleBody body, Set<IdentityStmt> catches) {
		Collection<Trap> traps = body.getTraps();
		for(IdentityStmt catchUnit:catches) {
			RefType refType = (RefType)catchUnit.getLeftOp().getType();
			traps.add(Jimple.v().newTrap(refType.getSootClass(), catchUnit, catchUnit, catchUnit));
		}
	}

	static void registerLocals(Body body) {
		Collection<Local> locals = body.getLocals();
		body.getUseAndDefBoxes().stream().map(ValueBox::getValue).filter(v->v instanceof Local).map(v->(Local)v).collect(Collectors.toSet()).forEach(locals::add);
	}
	
	static void addReturnStmt(JimpleBody body) {
		Type returnType = body.getMethod().getReturnType(); 
		if(returnType==VoidType.v())
			body.getUnits().add(Jimple.v().newReturnVoidStmt());
		else {
			Value retValue;
			if (returnType instanceof IntegerType) {
				retValue = IntConstant.v(1);
			} else if (returnType instanceof LongType) {
				retValue = LongConstant.v(1);
			} else if (returnType instanceof FloatType) {
				retValue = FloatConstant.v(1);
			} else if (returnType instanceof DoubleType) {
			    retValue = DoubleConstant.v(1);
			} else
				retValue = NullConstant.v();
			body.getUnits().add(Jimple.v().newReturnStmt(retValue));
		}
	}

}
