package re3lib;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;

import ghidra.app.decompiler.DecompileResults;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.pcode.HighFunction;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.PcodeOpAST;
import ghidra.program.model.pcode.Varnode;

public class PCallTracer {
	public class QueueItem {
		Function function;
		int depth;

		QueueItem(Function function, int depth) {
			this.function = function;
			this.depth = depth;
		}
	}

	public List<Function> out = new ArrayList<>();
	public boolean trace = false;
	List<QueueItem> queue = new ArrayList<>();
	HashSet<Address> visited = new HashSet<>();

	DecompileCache decomp;
	GhidraScript script;

	public PCallTracer() {
		this.script = RecompileConfig.INSTANCE.script;
		this.decomp = RecompileConfig.INSTANCE.decompCache;
	}

	public void setBlacklist(HashSet<Address> blacklist) {
		this.visited = new HashSet<>(blacklist);
	}

	void visit(HighFunction highFunction, int depth) {
		Iterator<PcodeOpAST> opIter = highFunction.getPcodeOps();
		while (opIter.hasNext()) {
			PcodeOpAST op = opIter.next();
			if (op.getOpcode() == PcodeOp.CALL) {
				Varnode target = op.getInput(0);
				if (target.isAddress()) {
					Address callAddr = target.getAddress();
					Function calledFunction = script.getFunctionAt(callAddr);
					if (calledFunction == null) {

						script.println("PCallTracer, called function not found: " + op.toString() + " - "
								+ highFunction.getFunction().getName());
						continue;
					}
					if (!visited.contains(calledFunction.getEntryPoint())) {
						queue.add(new QueueItem(calledFunction, depth + 1));
					}
				}
			}
		}
	}

	void visit(Function function, int depth) {
		if (!visited.contains(function.getEntryPoint())) {
			visited.add(function.getEntryPoint());
			if (trace) {
				// script.println("PCallTracer, visiting " + function.getName() + " (depth:" + depth + ")");
			}
			DecompileResults decompRes = decomp.getOrInsert(function);
			visit(decompRes.getHighFunction(), depth);
			out.add(function);
		}
	}

	public void traceCalls(Function inFunction) {
		queue.add(new QueueItem(inFunction, 0));
		while (queue.size() > 0) {
			QueueItem item = queue.remove(0);
			visit(item.function, item.depth);
		}
	}
}