package ch.usi.inf.sp.callgraph;

import org.objectweb.asm.Opcodes;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;


/**
 * Dump out information about the given ClassHierarchy.
 * 
 * @author maggicl@usi.ch
 * @author Matthias.Hauswirth@usi.ch
 */
public final class CallGraphRenderer {

	private final Map<String, Integer> identifiers = new HashMap<>();
	private final Map<Integer, Boolean> nodeBuilt = new HashMap<>();

	private String identifier(String what, boolean fromNode) {
		return "node" + identifiers.computeIfAbsent(what, (k) -> {
			int key = identifiers.size() + 1;
			nodeBuilt.compute(key, (ke, v) -> v == null ? fromNode : (v || fromNode));
			return key;
		});
	}

	private String classNodeName(ClassType type) {
		return identifier(type.getInternalName(), true);
	}

	private static String methodDescription(Method e) {
		return (e.getVerboseModifiers() + " " +
				e.getName().replaceAll("<", "&lt;").replaceAll(">", "&gt;") +
				e.getDescriptor()).trim();
	}

	private static String classNodeDescriptor(ClassType type) {
		final String className = type.getInternalName();

		if (!type.isResolved()) {
			return "[shape=ellipse,style=dashed,label=\"" + className + "\"]";
		}

		final String methods = type.getMethods().stream()
				.map(CallGraphRenderer::methodDescription)
				.collect(Collectors.joining("\\n"));

		final String style = type.isAbstract() ? "dashed" : "solid";

		return "[shape=record,style=" + style + ",label=\"{" + className + "|" + methods + "}\"]";
	}

	private String methodNodeName(ClassHierarchy hierarchy, String className, String name, String descriptor, boolean fromNode) {
		try {
			if (!hierarchy.getOrCreateClass(className).isResolved()) {
				return identifier(className, true);
			}
		} catch (TypeInconsistencyException e) {
			throw new RuntimeException(e);
		}

		return identifier(className + " " + name + " " + descriptor, fromNode);
	}

	private static String methodNodeDescriptor(Method method) {
		return "[shape=rectangle,style=filled,fillcolor=lightgreen,label=\"" + method.getDeclaringClassName() +
				"\\n" + methodDescription(method) + "\"]";
	}


	private static String getEdge(String from, String to, int opCode, String color) {
		String style;
		switch (opCode) {
			case Opcodes.INVOKEVIRTUAL: style = "dashed"; break;
			case Opcodes.INVOKEINTERFACE: style = "dotted"; break;
			default: style = "solid";
		}

		return from + " -> " + to + " [style=" + style + ",color=" + color + "]";
	}

	public void dumpDot(final ClassHierarchy hierarchy, final String fileName) throws IOException {
		final PrintWriter pw = new PrintWriter(new FileWriter(fileName));
		pw.println("digraph CallGraph {");
		pw.println("  rankdir=\"BT\"");

		final ArrayList<String> edges = new ArrayList<>();

		for (final Type type : hierarchy.getTypes()) {
			if (type instanceof ClassType) {
				final ClassType classType = (ClassType) type;
				final ClassType superClassType = classType.getSuperClass();

				pw.println(classNodeName(classType) + " " + classNodeDescriptor(classType) + ";");

				if (superClassType != null) {
					edges.add(getEdge(classNodeName(classType), classNodeName(superClassType), 0, "black"));
				}

				for (final ClassType c : classType.getInterfaces()) {
					edges.add(getEdge(classNodeName(classType), classNodeName(c), Opcodes.INVOKEVIRTUAL, "black"));
				}

				for (final Method m : classType.getMethods()) {
					final String mName = methodNodeName(hierarchy, m.getDeclaringClassName(), m.getName(), m.getDescriptor(), true);

					pw.println(mName + " " + methodNodeDescriptor(m) + ";");
					edges.add(getEdge(classNodeName(classType), mName, 0, "lightgreen"));

					for (final CallSite c : m.getCallSites()) {
						final String cName = methodNodeName(hierarchy,
								c.getDeclaredTargetClassName(),
								c.getTargetMethodName(),
								c.getTargetMethodDescriptor(), false);

						edges.add(getEdge(mName, cName, c.getOpcode(), "blue"));

						for (final ClassType p : c.getPossibleTargetClasses()) {
							final String pName = methodNodeName(hierarchy,
									p.getInternalName(),
									c.getTargetMethodName(),
									c.getTargetMethodDescriptor(), false);

							edges.add(getEdge(mName, pName, c.getOpcode(), "red"));
						}
					}
				}
			}
		}

		// Add missing node declarations (happens with methods in classes not scanned)
		final Map<Integer, String> reverseIdentifiers = identifiers.entrySet().stream()
				.collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
		for (Map.Entry<Integer, Boolean> e : nodeBuilt.entrySet()) {
			if (!e.getValue()) {
				pw.println("node" + e.getKey() + " [label=\"" + reverseIdentifiers.get(e.getKey()) + "\"]");
			}
		}

		pw.append('\n');
		for (final String edge : edges) {
			pw.println(edge);
		}
		pw.println("}");
		pw.close();
	}

}