package accrue.analysis.interprocanalysis;

import java.io.PrintStream;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

import polyglot.main.Report;
import polyglot.types.CodeInstance;
import accrue.analysis.AnalysisTopics;

/**
 * Representation of the interprocedural call graph.
 *
 */
public class CallGraph {
    private final Map<AnalysisUnit, Set<AnalysisUnit>> callerToCallee;
    private final Map<AnalysisUnit, Set<AnalysisUnit>> calleeToCaller;
    private final Set<AnalysisUnit> startProcs;
    private final Set<AnalysisUnit> allProcs;

    private final Map<AnalysisUnit, Set<AnalysisUnit>> recursiveCalls;

    private final Map<CodeInstance, Set<AnalysisUnit>> byCodeInstance;

    public CallGraph(Set<AnalysisUnit> startProcs) {
        this.callerToCallee =
                new LinkedHashMap<AnalysisUnit, Set<AnalysisUnit>>();
        this.calleeToCaller =
                new LinkedHashMap<AnalysisUnit, Set<AnalysisUnit>>();
        this.byCodeInstance =
                new LinkedHashMap<CodeInstance, Set<AnalysisUnit>>();
        this.recursiveCalls = new HashMap<AnalysisUnit, Set<AnalysisUnit>>();
        this.startProcs = Collections.unmodifiableSet(startProcs);

        // Initialize allProcs, and add all the startProcs to it.
        this.allProcs = new LinkedHashSet<AnalysisUnit>(startProcs);

    }

    public Set<AnalysisUnit> callees() {
        return calleeToCaller.keySet();
    }

    public Set<AnalysisUnit> calledBy(AnalysisUnit a) {
        Set<AnalysisUnit> dependees = calleeToCaller.get(a);
        if (dependees == null) {
            return Collections.emptySet();
        }
        return dependees;

    }

    public Set<AnalysisUnit> allProcs() {
        return allProcs;
    }

    public Set<AnalysisUnit> startProcs() {
        return startProcs;
    }

    private void recordByCodeInstance(AnalysisUnit au) {
        CodeInstance ci = au.cn.codeInstance();
        Set<AnalysisUnit> s = byCodeInstance.get(ci);
        if (s == null) {
            s = new LinkedHashSet<AnalysisUnit>();
            byCodeInstance.put(ci, s);
        }
        s.add(au);
    }

    public boolean addEdge(AnalysisUnit caller, AnalysisUnit callee) {
        recordByCodeInstance(caller);
        recordByCodeInstance(callee);

        Set<AnalysisUnit> s = callerToCallee.get(caller);

        if (s == null) {
            s = new LinkedHashSet<AnalysisUnit>();
            callerToCallee.put(caller, s);
        }

        boolean changed = s.add(callee);

        allProcs.add(callee);
        allProcs.add(caller); // defensively add the caller, in case it wasn't
        // in there already

        Set<AnalysisUnit> t = calleeToCaller.get(callee);

        if (t == null) {
            t = new LinkedHashSet<AnalysisUnit>();
            calleeToCaller.put(callee, t);
        }

        changed = t.add(caller) || changed;

        // Detect backedges by depth-first search from callee to caller
        if (existsPathBetween(callee, caller)) {
            // this call is a back edge
            addRecursiveCall(caller, callee);
        }

        if (changed && Report.should_report(AnalysisTopics.POINTSTO, 4))
            Report.report(4, "Adding call graph edge " + caller + " -> "
                    + callee);

        return changed;
    }

    /**
     * Depth first search of the CallGraph. Returns true if a path exists from source to destination.
     * @param from Source node in the CallGraph
     * @param to Destination node in the CallGraph
     * @return
     */
    private boolean existsPathBetween(AnalysisUnit source, AnalysisUnit dest) {
        return existsPathBetween(source,
                                 dest,
                                 new LinkedHashSet<AnalysisUnit>());
    }

    /**
     * Helper method for existsPathBetween, tracks set of previously visited nodes.
     * @param from
     * @param to
     * @param visited
     * @return
     */
    private boolean existsPathBetween(AnalysisUnit from, AnalysisUnit to,
            Set<AnalysisUnit> visited) {
        if (from.equals(to)) {
            return true;
        }
        else {
            visited.add(from);
            Set<AnalysisUnit> succ = callees(from);
            for (AnalysisUnit next : succ) {
                if (visited.contains(next)) {
                    continue;
                }
                if (existsPathBetween(next, to, visited)) {
                    return true;
                }
            }
            return false;
        }
    }

    private void addRecursiveCall(AnalysisUnit caller, AnalysisUnit callee) {
        Set<AnalysisUnit> s = recursiveCalls.get(caller);
        if (s == null) {
            s = new LinkedHashSet<AnalysisUnit>();
            recursiveCalls.put(caller, s);
        }
        s.add(callee);
    }

    public boolean existsPathOfLength(CodeInstance ci, int length) {
        Set<AnalysisUnit> s = byCodeInstance.get(ci);
        if (s == null) return false;
        for (AnalysisUnit au : s) {
            if (existsPathOfLength(au, length)) {
                return true;
            }
        }
        return false;
    }

    public boolean existsPathOfLength(AnalysisUnit au, int length) {
        return existsPathOfLength(au, new LinkedHashSet<AnalysisUnit>(), length);
    }

    private boolean existsPathOfLength(AnalysisUnit a,
            Set<AnalysisUnit> onPath, int length) {
        if (length == 0) return true;
        onPath.add(a);
        for (AnalysisUnit b : callees(a)) {
            // a calls b
            if (onPath.contains(b)) {
                // a -> b is a back-edge.
                // there is a cycle!
                return true;
            }
            else {
                if (existsPathOfLength(a, onPath, length - 1)) {
                    return true;
                }
            }
        }
        onPath.remove(a);
        return false;
    }

    public Set<AnalysisUnit> callees(AnalysisUnit r) {
        Set<AnalysisUnit> s = callerToCallee.get(r);
        if (s == null) {
            return Collections.emptySet();
        }
        return s;
    }

    public void dump(PrintStream out) {
        for (AnalysisUnit n : callerToCallee.keySet()) {
            out.println(n);
            for (AnalysisUnit s : callerToCallee.get(n)) {
                out.println("   " + s);
            }
        }
    }

    public int nodeSize() {
        return callerToCallee.size();
    }

    public long totalSize() {
        long sum = 0;
        for (Set<?> s : this.callerToCallee.values()) {
            sum += s.size();
        }
        return sum;
    }

    public boolean maybeRecursiveCall(AnalysisUnit caller, AnalysisUnit callee) {
        Set<AnalysisUnit> s = recursiveCalls.get(caller);
        if (s == null) return false;
        return s.contains(callee);
    }
}
