package accrue.analysis.preciseExceptions;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import polyglot.ast.ArrayInit;
import polyglot.ast.Assign;
import polyglot.ast.Binary;
import polyglot.ast.Block;
import polyglot.ast.Call;
import polyglot.ast.Case;
import polyglot.ast.Cast;
import polyglot.ast.Catch;
import polyglot.ast.CodeNode;
import polyglot.ast.CompoundStmt;
import polyglot.ast.Conditional;
import polyglot.ast.ConstructorCall;
import polyglot.ast.Eval;
import polyglot.ast.Expr;
import polyglot.ast.Field;
import polyglot.ast.If;
import polyglot.ast.Instanceof;
import polyglot.ast.JLang;
import polyglot.ast.Labeled;
import polyglot.ast.Local;
import polyglot.ast.LocalDecl;
import polyglot.ast.Loop;
import polyglot.ast.New;
import polyglot.ast.NewArray;
import polyglot.ast.Node;
import polyglot.ast.ProcedureCall;
import polyglot.ast.Return;
import polyglot.ast.Stmt;
import polyglot.ast.Synchronized;
import polyglot.ast.Term;
import polyglot.ast.Throw;
import polyglot.ast.Try;
import polyglot.ast.TypeNode;
import polyglot.ast.Unary;
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.DataFlow;
import polyglot.visit.FlowGraph;
import polyglot.visit.FlowGraph.Edge;
import polyglot.visit.FlowGraph.EdgeKey;
import polyglot.visit.FlowGraph.ExceptionEdgeKey;
import polyglot.visit.FlowGraph.Peer;
import polyglot.visit.NodeVisitor;
import accrue.analysis.ast.Suppress;
import accrue.analysis.ext.AccrueExt_c;
import accrue.analysis.ext.ExtCall;
import accrue.analysis.ext.ExtConstructorCall;
import accrue.analysis.ext.ExtNew;
import accrue.analysis.ext.ExtProcedureCall;
import accrue.analysis.ext.ExtTerm;
import accrue.analysis.interprocanalysis.AnalysisUtil;
import accrue.analysis.interprocanalysis.PreciseDataFlow;

public class PreciseExVisitor extends NodeVisitor {
    protected final PreciseExAnalysisUtil autil;
    protected final TypeSystem ts;

    protected boolean trackImplicitErrors = false;
    protected final Set<ThrowableType> throwTypes =
            new HashSet<ThrowableType>();

    public PreciseExVisitor(PreciseExAnalysisUtil autil) {
        super();
        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 PreciseExVisitor(TypeSystem ts) {
        super();
        if (ts == null) throw new InternalCompilerError("null ts");
        this.autil = null;
        this.ts = ts;
    }

    protected PreciseExVisitor(PreciseExVisitor v) {
        super();
        this.autil = v.autil;
        this.ts = v.ts;
    }

    public boolean removeThrowableType(ThrowableType tt) {
        return this.throwTypes.remove(tt);
    }

    public PreciseExAbsVal performVisit(CodeNode n) throws SemanticException {
        n.visit(this);

        // Perform a dataflow on the code decl to figure out if it can terminate or return
        // normally.        
        // this is probably a bit redundant (in that we could change the dataflow to also perform
        // the work that this visitor does), but we will do it separately for
        // legacy reasons.
        PreciseExDataflow df = new PreciseExDataflow(autil);
        boolean canTerminateNormally = df.performDataflow(n);
        return new PreciseExAbsVal(this.throwTypes, canTerminateNormally);
    }

    @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 old, Node n, NodeVisitor v) {

        if (trackImplicitErrors
                && ((n instanceof Stmt && !(n instanceof CompoundStmt)) || (n instanceof Block && ((Block) n).statements()
                                                                                                             .isEmpty()))) {
            // n can throw an error. It can just happen.
            // It sucks, but it's true.
            this.throwTypes.add(new ThrowableType(ts.Error(), false));
        }

        if (process(n)) {
            // let extensions handle their own nodes specially,
            // if they want to.
        }
        else 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 && this.autil != null) {
            this.throwTypes.addAll(((ExtTerm) AccrueExt_c.ext(n)).throwTypes(ts,
                                                                             autil.currentContext()));
        }
        else {
            for (Type t : ((JLang) lang()).throwTypes(n, ts)) {
                this.throwTypes.add(new ThrowableType(t, false));
            }
        }
        return n;
    }

    /**
     * An extension point, so that subclasses can handle any custom nodes
     * easily. 
     * @param n
     * @return true if the node has been processed, false 
     *        if the node has not been processed.
     */
    protected boolean process(Node n) {
        return false;
    }

    protected void process(Try n) {
        PreciseExVisitor v = new PreciseExVisitor(this);

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

        // remove exceptions based on the catch blocks
        for (Catch c : n.catchBlocks()) {
            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();
                }
            }
        }

        // 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.
            PreciseExVisitor w = new PreciseExVisitor(this.autil);
            n.finallyBlock().visit(w);

            // if the finally block can terminate normally, then add all of v to this
            if (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);
        }
    }

    protected boolean canTerminateNormally(Node n) {
        if (n instanceof ProcedureCall) {
            ProcedureCall pc = (ProcedureCall) n;
            if (!canAllTerminateNormally(pc.arguments())) {
                // one or more of the arguments can't terminate normally.
                return false;
            }
            if (pc instanceof Call) {
                if (!canTerminateNormally(((Call) pc).target())) {
                    // the receiver of the call can't terminate normally.
                    return false;
                }
            }
            // handles calls, new, and constructor calls.
            if (this.autil != null) {
                // check whether the call itself can terminate normally.
                ExtProcedureCall ext = (ExtProcedureCall) AccrueExt_c.ext(n);
                return ext.canTerminateNormally(this.autil.currentContext());
            }
            else {
                // conservatively assume that it can
                return true;
            }
        }
        if (n instanceof Assign) {
            Assign a = (Assign) n;
            return canTerminateNormally(a.left())
                    && canTerminateNormally(a.right());
        }
        if (n instanceof ArrayInit) {
            ArrayInit ai = (ArrayInit) n;
            return canAllTerminateNormally(ai.elements());
        }
        if (n instanceof Binary) {
            Binary b = (Binary) n;
            if (!canTerminateNormally(b.left())) {
                return false;
            }
            if (b.operator() == Binary.COND_AND
                    || b.operator() == Binary.COND_OR) {
                // it's a short circuit operator, so if the left can terminate
                // normally, then the whole thing can.
                return true;
            }
            return canTerminateNormally(b.right());
        }
        if (n instanceof Unary) {
            return canTerminateNormally(((Unary) n).expr());
        }
        if (n instanceof Conditional) {
            Conditional c = (Conditional) n;
            if (!canTerminateNormally(c.cond())) {
                return false;
            }
            return canTerminateNormally(c.consequent())
                    || canTerminateNormally(c.alternative());

        }
        if (n instanceof Instanceof) {
            return canTerminateNormally(((Instanceof) n).expr());
        }
        if (n instanceof NewArray) {
            NewArray na = (NewArray) n;
            return canAllTerminateNormally(na.dims())
                    && (na.init() == null || canTerminateNormally(na.init()));
        }
        if (n instanceof Field) {
            return canTerminateNormally(((Field) n).target());
        }
        if (n instanceof Cast) {
            // could be more precise (by checking the results of the pointer analysis), 
            // but let's assume that a cast
            // can terminate normally if the sub expression can.
            return canTerminateNormally(((Cast) n).expr());
        }

        if (n instanceof Local) return true;
        if (n instanceof Expr) return true;

        if (n instanceof Eval) return canTerminateNormally(((Eval) n).expr());
        if (n instanceof Return) {
            Return ret = (Return) n;
            if (ret.expr() == null) return true;
            return canTerminateNormally(ret.expr());
        }

        if (n instanceof Throw) return false;
        if (n instanceof CodeNode)
            return canTerminateNormally(((CodeNode) n).codeBody());

        if (n instanceof TypeNode) return true;

        if (n instanceof LocalDecl) {
            LocalDecl ld = (LocalDecl) n;
            return (ld.init() == null) || canTerminateNormally(ld.init());
        }
        if (n instanceof Case) {
            return canTerminateNormally(((Case) n).expr());
        }

        // should have handled all non-statements by now, and all non-compound statements
        // that may not terminate normally.

        if (!(n instanceof CompoundStmt)) return true;

        if (n instanceof Loop) return true;
        if (n instanceof Block) {
            Block b = (Block) n;
            return canAllTerminateNormally(b.statements());
        }
        if (n instanceof Labeled)
            return canTerminateNormally(((Labeled) n).statement());
        if (n instanceof Synchronized)
            return canTerminateNormally(((Synchronized) n).body());
        if (n instanceof If) {
            If i = (If) n;
            return canTerminateNormally(i.cond())
                    && (canTerminateNormally(i.consequent())
                            || i.alternative() == null || canTerminateNormally(i.alternative()));
        }
        if (n instanceof Try) {
            Try t = (Try) n;
            if (t.finallyBlock() != null
                    && !canTerminateNormally(t.finallyBlock())) {
                // finally block can't terminate normally
                return false;
            }
            if (!canTerminateNormally(t.tryBlock())) {
                for (Catch catchBlock : t.catchBlocks()) {
                    if (canTerminateNormally(catchBlock.body())) {
                        // If any catch block can return normally we'll assume the 
                        // try can return normally
                        return true;
                    }
                }

                return false;
            }

            // let's be imprecise and just say it can
            return true;
        }
        throw new InternalCompilerError("Haven't implemented handling of " + n);
    }

    protected boolean canAllTerminateNormally(List<? extends Node> ns) {
        for (Node n : ns) {
            if (!canTerminateNormally(n)) {
                return false;
            }
        }
        return true;
    }

    protected 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();
            PreciseExAbsVal pexs = autil.call(mi, n);
            this.throwTypes.addAll(pexs.throwTypes());
            if (ext.targetMayBeNull(autil.currentContext())) {
                this.throwTypes.add(new ThrowableType(ts.NullPointerException(),
                                                      true));
            }
            ext.recordPreciseExResults(this.autil.currentContext(),
                                       pexs.throwTypes(),
                                       pexs.canTerminateNormally()
                                               && canTerminateNormally(n.target())
                                               && canAllTerminateNormally(n.arguments()));
        }
    }

    protected void process(ConstructorCall n) {
        if (autil == null) {
            for (Type t : ((JLang) lang()).throwTypes(n, ts)) {
                this.throwTypes.add(new ThrowableType(t, false));
            }
        }
        else {
            PreciseExAbsVal pexs =
                    autil.constructorCall(n.constructorInstance(), n);
            this.throwTypes.addAll(pexs.throwTypes());
            Set<ThrowableType> throwTypes = pexs.throwTypes();

            if (n.kind().equals(ConstructorCall.SUPER)) {
                PreciseExAbsVal t = null;
                try {
                    // Add in the exceptions from all of the initializers
                    for (PreciseExAbsVal v : autil.analyzeInitializers(PreciseExAbsVal.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(pexs.throwTypes());
                    throwTypes.addAll(t.throwTypes());
                }
            }
            ExtConstructorCall ext = (ExtConstructorCall) AccrueExt_c.ext(n);
            ext.recordPreciseExResults(this.autil.currentContext(),
                                       throwTypes,
                                       pexs.canTerminateNormally()
                                               && canAllTerminateNormally(n.arguments()));
        }
    }

    protected void process(New n) {
        if (autil == null) {
            for (Type t : ((JLang) lang()).throwTypes(n, ts)) {
                this.throwTypes.add(new ThrowableType(t, false));
            }
        }
        else {
            PreciseExAbsVal pexs = autil.call(n.constructorInstance(), n);
            this.throwTypes.addAll(pexs.throwTypes());
            ExtNew ext = (ExtNew) AccrueExt_c.ext(n);
            ext.recordPreciseExResults(this.autil.currentContext(),
                                       pexs.throwTypes(),
                                       pexs.canTerminateNormally()
                                               && canAllTerminateNormally(n.arguments()));
        }
    }

    protected void process(Suppress n) {
        PreciseExVisitor v = new PreciseExVisitor(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;
    }

    static class Reachable extends DataFlow.Item {
        public static final Reachable REACHABLE = new Reachable(true);
        public static final Reachable NOT_REACHABLE = new Reachable(false);
        public final boolean isReachable;

        private Reachable(boolean isReachable) {
            this.isReachable = isReachable;
        }

        @Override
        public int hashCode() {
            final int prime = 31;
            int result = 1;
            result = prime * result + (isReachable ? 1231 : 1237);
            return result;
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) return true;
            if (obj == null) return false;
            if (getClass() != obj.getClass()) return false;
            Reachable other = (Reachable) obj;
            if (isReachable != other.isReachable) return false;
            return true;
        }

        public static Reachable reachable(boolean isReachable) {
            if (isReachable) {
                return Reachable.REACHABLE;
            }
            return Reachable.NOT_REACHABLE;
        }

        @Override
        public String toString() {
            if (isReachable) {
                return "REACHABLE";
            }
            return "NOT_REACHABLE";
        }
    }

    class PreciseExDataflow extends PreciseDataFlow<Reachable> {

        private boolean isExitReachable = false;

        public PreciseExDataflow(AnalysisUtil<?> autil) {
            super(autil,
                  autil.extensionInfo().scheduler().currentJob(),
                  autil.extensionInfo().typeSystem(),
                  autil.extensionInfo().nodeFactory(),
                  true);
        }

        public boolean performDataflow(CodeNode n) throws SemanticException {
            this.isExitReachable = false;
            dataflow(n);
            return this.isExitReachable;
        }

        @Override
        protected Reachable createInitialItem(FlowGraph<Reachable> graph,
                Term node, boolean entry) {
            return Reachable.REACHABLE;
        }

        @Override
        protected Reachable confluence(List<Reachable> items,
                Peer<Reachable> peer, FlowGraph<Reachable> graph) {
            // if any of the incoming edges are reachable, then the confluence should be
            for (Reachable r : items) {
                if (r.isReachable) {
                    return Reachable.REACHABLE;
                }
            }
            return Reachable.NOT_REACHABLE;
        }

        @Override
        protected Map<EdgeKey, Reachable> flow(Reachable in,
                FlowGraph<Reachable> graph, Peer<Reachable> peer) {
            if (!in.isReachable) {
                // we can't reach this node.
                return reachableMap(peer.succEdgeKeys(),
                                    Reachable.NOT_REACHABLE,
                                    Reachable.NOT_REACHABLE);
            }
            // this node is reachable. See whether it can terminate (or return) normally
            Map<EdgeKey, Reachable> rm =
                    reachableMap(peer.succEdgeKeys(),
                                 Reachable.reachable(peer.isEntry()
                                         || canTerminateNormally(peer.node())),
                                 Reachable.NOT_REACHABLE);

            return rm;
        }

        private Map<EdgeKey, Reachable> reachableMap(Set<EdgeKey> edgeKeys,
                Reachable normalEdges, Reachable exceptionEdges) {
            Map<EdgeKey, Reachable> m = new HashMap<EdgeKey, Reachable>();
            for (EdgeKey ek : edgeKeys) {
                if (ek instanceof ExceptionEdgeKey) {
                    m.put(ek, exceptionEdges);
                }
                else {
                    m.put(ek, normalEdges);
                }
            }
            return m;
        }

        @Override
        protected void check(FlowGraph<Reachable> graph, Peer<Reachable> p)
                throws SemanticException {
            if (graph.exitPeers().contains(p)) {
                // it's an exit peer! see if there is an incoming non-exception edge
                // that is reachable.
                for (Edge<Reachable> e : p.preds()) {
                    Peer<Reachable> o = e.getTarget();
                    if (e.getKey() instanceof ExceptionEdgeKey) {
                        // it's an exception edge key, and we don't care about it
                    }
                    else {
                        // it's not an exception edge key
                        Reachable r = o.outItem(e.getKey());
                        if (r != null && r.isReachable) {
                            // yes! The exit node is reachable on a non-exception edge
                            this.isExitReachable = true;
                            return;
                        }
                    }
                }

            }
            else {
                // nothing to do here.
            }
        }

        @Override
        protected void check(FlowGraph<Reachable> graph, Term n, boolean entry,
                Reachable inItem, Map<EdgeKey, Reachable> outItems)
                throws SemanticException {
            throw new UnsupportedOperationException();
        }

    }
}
