package grafos;

import com.github.javaparser.ast.Node;
import com.github.javaparser.ast.expr.BooleanLiteralExpr;
import com.github.javaparser.ast.expr.SimpleName;
import com.github.javaparser.ast.stmt.*;

import java.util.*;

/**
 * Represents a control flow graph, offering methods to build it easily with a visitor and an output ready to be
 * converted to a graph by GraphViz ({@link #toStringList(GraphViz)}).
 * <br/>
 * The process should begin with {@link #beginMethod()} and end with {@link #finishMethod()}. Similarly, every opened
 * {@link Block} and section must be closed. The class offers special support for {@link ContinueStmt} and {@link BreakStmt},
 * including both the labeled and unlabeled versions.
 * <br/>
 * The inner class {@link Block} is used to represent a block statement (not only {@link BlockStmt} but also {@link IfStmt}
 * and {@link WhileStmt}, for example). {@link Block} should not be used directly, only through the API provided by CFG.
 */
public class CFG {
	/**
	 * List of nodes that exist in this CFG (can either be {@link com.github.javaparser.ast.expr.Expression} or {@link Statement})
	 */
	private final List<Node> nodes = new LinkedList<>();
	/** List of blocks (statements that contain other statements) in this CFG.
	 * Only includes the "surface" blocks, not the ones inside other blocks.
	 * @see Block */
	private final LinkedList<Block> blocks = new LinkedList<>();
	/** Stack of blocks, representing the position in the method. The current position is inside the topmost block,
	 * which is contained inside the previous one. If empty, the position is outside all blocks, in the main level
	 * of the method. */
	private final Stack<Block> blockStack = new Stack<>();
	/** List of {@link Edge}, or connections between nodes. */
	private final List<Edge> edges = new LinkedList<>();
	/** First and last node of a method, represented by an {@link EmptyStmt} instance. */
	private final Node beginNode = new EmptyStmt(), endNode = new EmptyStmt();
	/** List of nodes that are immediately before the next instruction, and therefore will be linked to it.
	 * May be empty after a break/return/continue, or contain more than one element after a control-flow instruction. */
	private final LinkedList<Node> nodeList = new LinkedList<>();
	/** Map of nodes to the name that needs to be applied to the edge when that node is connected to the next one.
	 * Used to distinguish the {@code true} and {@code false} branches in conditions. */
	private final Map<Node, String> labelMap = new HashMap<>();
	/** Map of loop labels to lists of break statements (all labeled). Each list begins empty and {@link BreakStmt} are
	 * added as found. When the end of the labeled loop is reached, the list is removed and dumped into {@link #nodeList}. */
	private final Map<SimpleName, List<BreakStmt>> breakMap = new HashMap<>();
	/** Stack of lists of break statements (without label). A new list is added when entering any structure that may be broken out of,
	 * and the list is dumped into the {@link #nodeList} when exiting the loop. */
	private final Stack<List<BreakStmt>> breakStack = new Stack<>();
	/** Map of loop labels to the node that the loop "continues" to, i.e., the node to which a continue stmt will be
	 * connected. */
	private final Map<SimpleName, Node> continueMap = new HashMap<>();
	/** Stack of the node that the loop "continues" to.
	 * @see #continueMap */
	private final Stack<Node> continueStack = new Stack<>();

	/** Last labeled visited, if a {@link LabeledStmt} has been visited, but the corresponding loop hasn't been yet.
	 * Otherwise, {@code null}. */
	private SimpleName lastLabel;

	/** Adds a node to the list of nodes known by this tree */
	private void registerNode(Node stmt) {
		nodes.add(stmt);
		if (!blockStack.isEmpty())
			blockStack.peek().add(stmt);
		System.out.println("NODO: " + node2str(stmt));
	}

	/** Clears the list of nodes and adds the given node. */
	public void setNode(Node node) {
		clearNode();
		appendNode(node);
	}

	/** Clears the list of nodes. */
	public void clearNode() {
		nodeList.clear();
	}

	/** Connect each node in the list of nodes to the current node, then set the node as the only element in the list. */
	public void connectTo(Node node) {
		connectTo(node, true);
	}

	/** Connect each node in the list of nodes to the current node, then optionally set the node as the only
	 * element in the list. */
	public void connectTo(Node node, boolean overwrite) {
		if (!nodesContains(node))
			registerNode(node);
		for (Node begin : nodeList)
			connect(begin, node);
		if (overwrite)
			setNode(node);
	}

	/** Checks if an object is in the node registration list. Implemented to use {@code ==} instead of
	 * {@link List#contains(Object)}, which uses {@link Object#equals(Object)}. */
	private boolean nodesContains(Object object) {
		for (Node n : nodes)
			if (n == object)
				return true;
		return false;
	}

	/** Removes a node from the list of nodes */
	public void removeNode(Node node) {
		nodeList.remove(node);
	}

	/** Adds a node to the list of nodes */
	public void appendNode(Node node) {
		nodeList.add(node);
	}

	/** Returns an unmodifiable copy of the list of nodes */
	public List<Node> getNodeList() {
		return Collections.unmodifiableList(nodeList);
	}

	/** Replace the list of nodes with this list. The argument may be reused without affecting the list of nodes. */
	public void setNodeList(List<Node> nodes) {
		nodeList.clear();
		nodeList.addAll(nodes);
	}

	// Method operations

	/** Begin a method: clear the list of nodes and add the "Start" node */
	public void beginMethod() {
		nodeList.clear();
		nodeList.add(beginNode);
	}

	/** Finish a method: connect each element of the list of nodes to the end (unless they are the end already) */
	public void finishMethod() {
		for (Node begin : nodeList)
			if (begin != endNode)
				connect(begin, endNode);
	}

	// Block operations

	/** Start a new block, with a node that it represents and its condition. */
	public void beginBlock(Node node, Node condition) {
		Block b = new Block(node, condition);
		if (condition != null)
			registerNode(condition);
		if (blockStack.isEmpty())
			blocks.addFirst(b);
		else
			blockStack.peek().addSubBlock(b);
		blockStack.push(b);
	}

	/** End a block */
	public void endBlock() {
		blockStack.pop();
	}

	/** Set the label to be used when the last element of the list of nodes is connected to the next node. Run this
	 * AFTER adding the node with {@link #appendNode(Node)}. */
	public void setNextLabel(boolean b) {
		labelMap.put(nodeList.getLast(), String.valueOf(b));
	}

	/** Connects one node with another, checking the map of labels and clearing the appropriate entry */
	private void connect(Node begin, Node end) {
		edges.add(new Edge(begin, end, labelMap.get(begin)));
		labelMap.remove(begin);
	}

	public Node getEndNode() {
		return endNode;
	}

	// Break operations

	/** Begin a new list of break statements associated with the given label */
	public void beginBreakSection(SimpleName name) {
		breakMap.put(name, new LinkedList<>());
	}

	/** Begin a new list of unlabeled break statements */
	public void beginBreakSection() {
		breakStack.push(new LinkedList<>());
	}

	/** Connect a break statement to the chain, then add it to the appropriate list, so that when the switch/loop ends
	 * it can be connected to the proper node, which is currently unknown. */
	public void connectBreak(BreakStmt stmt) {
		connectTo(stmt);
		if (stmt.getLabel().isPresent()) {
			breakMap.get(stmt.getLabel().get()).add(stmt);
		} else {
			breakStack.peek().add(stmt);
		}
		clearNode();
	}

	/** Ends a list of labeled break statements, adding all the break statements collected to the list of nodes. */
	public void endBreakSection(SimpleName name) {
		nodeList.addAll(breakMap.remove(name));
	}

	/** Ends a list of unlabeled break statements, adding all the break statements collected to the list of nodes.
	 * Restores the previous list (if it exists), so that nested loops/switch statements work properly */
	public void endBreakSection() {
		nodeList.addAll(breakStack.pop());
	}

	// Continue operations

	/** Begin a labeled continue section (loop). Inside the loop, {@link #beginContinueSection(Node)} must be called for
	 * this method to work properly. */
	public void beginContinueSection(SimpleName name) {
		assert lastLabel == null;
		lastLabel = name;
	}

	/** Begin a continue section (loop). Will register the given node as the to-go node when a continue statement is
	 * found, and optionally register it with the proper label if {@link #beginContinueSection(SimpleName)} has been
	 * called immediately before. */
	public void beginContinueSection(Node node) {
		continueStack.push(node);
		if (lastLabel != null) {
			assert !continueMap.containsKey(lastLabel);
			continueMap.put(lastLabel, node);
			lastLabel = null;
		}
	}

	/** Connect each statement in the list of nodes to the continue statement, then connect the continue to the proper
	 * statement, depending on whether it has a label or not. */
	public void connectContinue(ContinueStmt stmt) {
		connectTo(stmt);
		if (stmt.getLabel().isPresent()) {
			connectTo(continueMap.get(stmt.getLabel().get()));
		} else {
			connectTo(continueStack.peek());
		}
		clearNode();
	}

	/** Ends a labeled continue section */
	public void endContinueSection(SimpleName name) {
		continueMap.remove(name);
	}

	/** Ends an unlabeled continue section (loop). Works with nested loops */
	public void endContinueSection() {
		continueStack.pop();
	}

	/** Outputs the contents of this object as a list of strings, ready to be parsed by dot and turned into a graph. */
	public List<String> toStringList(GraphViz gv) {
		List<String> res = new LinkedList<>();
		res.add(gv.start_graph());
		for (Block b : blocks)
			res.addAll(b.toStringList(gv));
		for (Edge e : edges) {
			res.add(e.toString());
		}
		res.add(gv.end_graph());
		return res;
	}

	/** Converts a node to string, prepending an index and line number, and escaping any quotations. */
	private String node2str(Node n) {
		if (n == beginNode) return "Start";
		if (n == endNode) return "Stop";
		return String.format("\"(%d, %d) %s\"", indexOfNode(n) + 1, n.getRange().get().begin.line, escape(n.toString()));
	}

	/** Finds the index of a node in the node registration list. Throws an exception if unsuccessful.
	 * This implementation is necessary to use {@code ==} instead of {@link Object#equals(Object)}. */
	private int indexOfNode(Node n) {
		for (int i = 0; i < nodes.size(); i++)
			if (nodes.get(i) == n)
				return i;
		throw new RuntimeException("Internal error, can't find node in list");
	}

	/** Escapes a string. Currently only escapes quotation marks. */
	private String escape(String s) {
		return s.replace("\"", "\\\"");
	}

	/** Id for each subgraph/block (incremented with each new block) */
	private static int clusterId = 0;

	/** Representation of a statement containing {@link Node}, such as an if, loop or switch statement */
	class Block extends LinkedList<Node> {
		/**
		 * The node that is being represented by this block
		 */
		private final Node container;
		/**
		 * The condition to enter the current block, or the condition of the if
		 */
		private final Node condition;
		/** The sub-blocks present in this Block */
		private final List<Block> inners = new LinkedList<>();

		Block(Node container, Node condition) {
			this.container = container;
			this.condition = condition;
		}

		/** Add a sub-block to this block */
		void addSubBlock(Block block) {
			inners.add(block);
		}

		/** Get a representation ready to be used in dot (of the subgraph) */
		List<String> toStringList(GraphViz gv) {
			List<String> res = new LinkedList<>();
			res.add(gv.start_subgraph(clusterId++)); // start
			res.add("label = \"" + getTitle() + "\""); // title of the sub-graph
			// Add the condition node with a special shape, then reset the shape back to the default
			if (condition != null) {
				res.add("node [ shape = rectangle ]; " + node2str(condition));
				res.add("node [ shape = ellipse ]");
			}
			// Add all inner blocks
			for (Block b : inners) {
				res.addAll(b.toStringList(gv));
			}
			// Add all nodes as part of the sub-graph
			for (Node n : this) {
				res.add(node2str(n));
			}
			res.add(gv.end_subgraph()); // end
			return res;
		}

		/** Generate a title for the sub-graph based on the block instruction it represents. */
		String getTitle() {
			// Variables to be filled for String#format(String, Object[])
			String template;
			Object[] filler;
			if (container instanceof IfStmt) {
				template = "if (%s)";
				filler = new Object[]{((IfStmt) container).getCondition()};
			} else if (container instanceof WhileStmt) {
				template = "while (%s)";
				filler = new Object[]{((WhileStmt) container).getCondition()};
			} else if (container instanceof DoStmt) {
				template = "do-while (%s)";
				filler = new Object[]{((DoStmt) container).getCondition()};
			} else if (container instanceof ForStmt) {
				template = "for (%s; %s; %s)";
				filler = new Object[]{
						((ForStmt) container).getInitialization(),
						((ForStmt) container).getCompare().orElse(new BooleanLiteralExpr(true)),
						((ForStmt) container).getUpdate()};
			} else if (container instanceof ForeachStmt) {
				template = "for (%s : %s)";
				filler = new Object[]{
						((ForeachStmt) container).getVariable(),
						((ForeachStmt) container).getIterable()};
			} else if (container instanceof BlockStmt) {
				template = "";
				filler = new Object[0];
			} else if (container instanceof SwitchStmt) {
				template = "switch (%s)";
				filler = new Object[]{((SwitchStmt) container).getSelector()};
			} else {
				throw new RuntimeException("Unimplemented!!!");
			}
			return String.format(template, filler);
		}
	}

	/** Representation of a connection between two nodes, with an optional label */
	private class Edge {
		private final Node origin, destination;
		private final String label;

		private Edge(Node origin, Node destination, String label) {
			this.origin = origin;
			this.destination = destination;
			this.label = label;
		}

		@Override
		public String toString() {
			String s = node2str(origin) + " -> " + node2str(destination);
			if (label != null) {
				s += " [ label = \"" + label + "\" ]";
			}
			return s;
		}
	}
}