package accrue.analysis.callgraph;

import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.Set;

import polyglot.ast.Block;
import polyglot.ast.Call;
import polyglot.ast.Catch;
import polyglot.ast.CodeDecl;
import polyglot.ast.CodeNode;
import polyglot.ast.ConstructorCall;
import polyglot.ast.Initializer;
import polyglot.ast.JLang;
import polyglot.ast.New;
import polyglot.ast.Node;
import polyglot.ast.Return;
import polyglot.ast.Term;
import polyglot.ast.Try;
import polyglot.types.ClassType;
import polyglot.types.MethodInstance;
import polyglot.types.SemanticException;
import polyglot.types.Type;
import polyglot.types.TypeSystem;
import polyglot.util.InternalCompilerError;
import polyglot.visit.NodeVisitor;
import accrue.analysis.ast.Suppress;
import accrue.analysis.ext.AccrueExt;
import accrue.analysis.ext.AccrueExt_c;
import accrue.analysis.ext.ExtCall;
import accrue.analysis.ext.ExtTerm;
import accrue.analysis.interprocanalysis.EdgeIdentifier;
import accrue.analysis.interprocanalysis.EdgeIdentifier.EdgeType;
import accrue.analysis.interprocanalysis.EdgeIdentifier.ExceptionEdge;
import accrue.analysis.interprocanalysis.EdgeIdentifier_c;
import accrue.analysis.interprocanalysis.NodeIdentifier;
import accrue.analysis.interprocanalysis.NodeIdentifier_c;
import accrue.analysis.preciseExceptions.ThrowableType;

public class PreciseCallGraphVisitor extends NodeVisitor {
    protected final PreciseCallGraphAnalysisUtil autil;
    protected final TypeSystem ts;

    protected final Set<ThrowableType> throwTypes =
            new HashSet<ThrowableType>();
    protected final Set<EdgeIdentifier> throwSites =
            new HashSet<EdgeIdentifier>();
    protected final Set<EdgeIdentifier> returnSites =
            new HashSet<EdgeIdentifier>();

    public PreciseCallGraphVisitor(PreciseCallGraphAnalysisUtil autil) {
        this.autil = autil;
        if (this.autil == null) throw new InternalCompilerError("null autil");
        this.ts = autil.typeSystem();
        if (this.ts == null) throw new InternalCompilerError("null ts");
    }

    public PreciseCallGraphVisitor(TypeSystem ts) {
        if (ts == null) throw new InternalCompilerError("null ts");
        this.autil = null;
        this.ts = ts;
    }

    private PreciseCallGraphVisitor(PreciseCallGraphVisitor v) {
        this.autil = v.autil;
        this.ts = v.ts;
    }

    public PCGAbsVal performVisit(CodeNode n) {
        n.visit(this);
        boolean canTerminateNormally =
                autil.canTerminateNormally(n) || canReturnNormally(n);
        if (n instanceof CodeDecl && ((CodeDecl) n).body().reachable()) {
            returnSites.add(new EdgeIdentifier_c(EdgeIdentifier.EDGE_RETURN,
                                                 new NodeIdentifier_c(autil.currentContext(),
                                                                      n)));
        }
        return new PCGAbsVal(throwTypes,
                             throwSites,
                             returnSites,
                             canTerminateNormally);
    }

    private boolean canReturnNormally(CodeNode n) {
        final Set<Node> result = new HashSet<Node>();
        n.visit(new NodeVisitor() {
            @Override
            public Node leave(Node old, Node n, NodeVisitor v) {
                if (n instanceof Return && autil.canTerminateNormally(n)) {
                    result.add(n);
                }
                return n;
            }
        });

        return !result.isEmpty();

    }

    @Override
    public Node override(Node n) {
        if (n instanceof Try) {
            process((Try) n);
            return n;
        }
        if (n instanceof Suppress) {
            process((Suppress) n);
            return n;
        }
        return super.override(n);
    }

    @Override
    public Node leave(Node parent, Node old, Node n, NodeVisitor v) {
        if (n instanceof Call) {
            process((Call) n);
        }
        else if (n instanceof New) {
            process((New) n);
        }
        else if (n instanceof ConstructorCall) {
            process((ConstructorCall) n);
        }
        else if (n instanceof Try) {
            throw new InternalCompilerError("Should be impossible");
        }
        else if (n instanceof Term) {
            if (n instanceof Return) {
                returnSites.add(new EdgeIdentifier_c(EdgeIdentifier.EDGE_RETURN,
                                                     new NodeIdentifier_c(autil.currentContext(),
                                                                          (Return) n)));
            }
            else if (n instanceof Block && parent instanceof CodeDecl) {
                Block b = (Block) n;
                if (b.reachable() && autil.canTerminateNormally(b)) {
                    // code body exits
                    EdgeType et =
                            (parent instanceof Initializer) ? EdgeIdentifier.EDGE_INIT
                                    : EdgeIdentifier.EDGE_RETURN;
                    returnSites.add(new EdgeIdentifier_c(et,
                                                         new NodeIdentifier_c(autil.currentContext(),
                                                                              b)));
                }
            }
            for (ThrowableType tt : ((ExtTerm) AccrueExt_c.ext(n)).throwTypes(ts,
                                                                              autil.currentContext())) {
                throwTypes.add(tt);
                addThrowSite(tt, n);
            }
        }
        else {
            for (Type t : ((JLang) lang()).throwTypes(n, ts)) {
                this.throwTypes.add(new ThrowableType(t, false));
            }
        }
        return n;
    }

    private void addThrowSite(ThrowableType tt, Node n) {
        EdgeType excEdge = new ExceptionEdge((ClassType) tt.type());
        throwSites.add(new EdgeIdentifier_c(excEdge,
                                            new NodeIdentifier_c(autil.currentContext(),
                                                                 (Term) n)));
    }

    private void process(Try n) {
        PreciseCallGraphVisitor v = new PreciseCallGraphVisitor(this);

        // visit the try block with v
        n.tryBlock().visit(v);

        // remove exceptions based on the catch blocks
        for (polyglot.ast.Catch c : n.catchBlocks()) {
            Set<EdgeIdentifier> catchSites = new HashSet<EdgeIdentifier>();
            Type catchType = c.catchType();
            for (Iterator<ThrowableType> iter = v.throwTypes.iterator(); iter.hasNext();) {
                ThrowableType k = iter.next();
                if (k.type().isSubtype(catchType)) {
                    // the exception will always be caught.
                    iter.remove();
                }
            }
            for (Iterator<EdgeIdentifier> iter = v.throwSites.iterator(); iter.hasNext();) {
                EdgeIdentifier edge = iter.next();
                ExceptionEdge ex = (ExceptionEdge) edge.getKey();
                if (ex.type().isSubtype(catchType)) {
                    catchSites.add(edge);
                    // the exception will always be caught.
                    iter.remove();
                }
            }
            recordCatchSites(n, catchSites);
        }

        // visit the catch blocks with v
        for (Catch c : n.catchBlocks()) {
            c.body().visit(v);
        }

        if (n.finallyBlock() != null) {
            // now visit the finally block with w.
            PreciseCallGraphVisitor w = new PreciseCallGraphVisitor(this.autil);
            n.finallyBlock().visit(w);

            // if the finally block can terminate normally, then add all of v to this
            if (autil.canTerminateNormally(n.finallyBlock())) {
                this.throwTypes.addAll(v.throwTypes);
            }

            // add w to this
            this.throwTypes.addAll(w.throwTypes);
        }
        else {
            // no finally block
            this.throwTypes.addAll(v.throwTypes);
        }
    }

    private void process(Call n) {
        ExtCall ext = (ExtCall) AccrueExt_c.ext(n);
        if (autil == null) {
            if (ext.targetMayBeNull(null)) {
                this.throwTypes.add(new ThrowableType(ts.NullPointerException(),
                                                      true));
            }
            for (Type t : ((JLang) lang()).throwTypes(n, ts)) {
                this.throwTypes.add(new ThrowableType(t, false));
            }
        }
        else {
            MethodInstance mi = n.methodInstance();
            PCGAbsVal dfi = autil.call(mi, n);
            this.throwTypes.addAll(dfi.throwTypes());
            if (ext.targetMayBeNull(autil.currentContext())) {
                ThrowableType npe =
                        new ThrowableType(ts.NullPointerException(), true);
                this.throwTypes.add(npe);
                addThrowSite(npe, n);
            }

            recordReturnSources(n, dfi.returnSites());
        }
    }

    private void process(ConstructorCall n) {
        PCGAbsVal s = autil.constructorCall(n.constructorInstance(), n);
        this.throwTypes.addAll(s.throwTypes());
        Set<ThrowableType> throwTypes = s.throwTypes();

        if (n.kind().equals(ConstructorCall.SUPER)) {
            PCGAbsVal t = null;
            try {
                // Add in the exceptions from all of the initializers
                for (PCGAbsVal v : autil.analyzeInitializers(n, PCGAbsVal.EMPTY)
                                        .allItems()) {
                    t = (t == null) ? v : t.upperBound(v);
                }
            }
            catch (SemanticException e) {
                throw new InternalCompilerError(e);
            }
            this.throwTypes.addAll(t.throwTypes());

            if (!t.throwTypes().isEmpty()) {
                throwTypes = new LinkedHashSet<ThrowableType>();
                throwTypes.addAll(s.throwTypes());
                throwTypes.addAll(t.throwTypes());
            }
        }
        recordReturnSources(n, s.returnSites());
    }

    private void process(New n) {
        PCGAbsVal s = autil.call(n.constructorInstance(), n);
        this.throwTypes.addAll(s.throwTypes());
        recordReturnSources(n, s.returnSites());
    }

    private void process(Suppress n) {
        PreciseCallGraphVisitor v = new PreciseCallGraphVisitor(this.autil);

        // visit the try block with v
        n.expr().visit(v);

        // remove exceptions based on the suppress list
        for (ClassType ct : n.suppressExceptionTypes()) {
            for (Iterator<ThrowableType> iter = v.throwTypes().iterator(); iter.hasNext();) {
                ThrowableType k = iter.next();
                if (k.type().isSubtype(ct)) {
                    // the exception is suppressed
                    iter.remove();
                }
            }
        }

        this.throwTypes.addAll(v.throwTypes());

    }

    public Set<ThrowableType> throwTypes() {
        return this.throwTypes;
    }

    protected void recordReturnSources(Term n, Set<EdgeIdentifier> returnSites) {
        // call -> return
        AccrueExt_c.ext(n)
                   .recordAnalysisResult(PreciseCallGraphFactory.RETURN_SOURCES,
                                         this.autil.currentContext(),
                                         returnSites);
        // return -> call
        for (EdgeIdentifier e : returnSites) {
            NodeIdentifier enid = e.getTarget();
            if (!enid.isNative()) {
                Node target = enid.getNode();
                AccrueExt returnExt = AccrueExt_c.ext(target);
                Set<EdgeIdentifier> targets =
                        returnExt.getAnalysisResult(PreciseCallGraphFactory.RETURN_TARGETS,
                                                    enid.getContext());
                if (targets == null) {
                    targets = new LinkedHashSet<EdgeIdentifier>();
                    returnExt.recordAnalysisResult(PreciseCallGraphFactory.RETURN_TARGETS,
                                                   enid.getContext(),
                                                   targets);
                }
                NodeIdentifier nid =
                        new NodeIdentifier_c(autil.currentContext(), n);
                targets.add(new EdgeIdentifier_c(EdgeIdentifier.EDGE_RETURN,
                                                 nid));
            }
        }

    }

    protected void recordCatchSites(Term n, Set<EdgeIdentifier> catchSites) {
        // catch -> throw
        AccrueExt_c.ext(n)
                   .recordAnalysisResult(PreciseCallGraphFactory.CATCH_SOURCES,
                                         this.autil.currentContext(),
                                         catchSites);
        // throw -> catch
        for (EdgeIdentifier e : catchSites) {
            NodeIdentifier enid = e.getTarget();
            if (!enid.isNative()) {
                Node target = enid.getNode();
                AccrueExt returnExt = AccrueExt_c.ext(target);
                Set<EdgeIdentifier> targets =
                        returnExt.getAnalysisResult(PreciseCallGraphFactory.CATCH_TARGETS,
                                                    enid.getContext());
                if (targets == null) {
                    targets = new LinkedHashSet<EdgeIdentifier>();
                    returnExt.recordAnalysisResult(PreciseCallGraphFactory.CATCH_TARGETS,
                                                   enid.getContext(),
                                                   targets);
                }
                ExceptionEdge exEdge = (ExceptionEdge) e.getKey();
                NodeIdentifier nid =
                        new NodeIdentifier_c(autil.currentContext(), n);
                targets.add(new EdgeIdentifier_c(exEdge, nid));
            }
        }

    }

}
