package accrue.analysis.interprocanalysis;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import polyglot.ast.CodeDecl;
import polyglot.ast.CodeNode;
import polyglot.ast.ConstructorDecl;
import polyglot.ast.FieldDecl;
import polyglot.ast.Initializer;
import polyglot.ast.MethodDecl;
import polyglot.frontend.ExtensionInfo;
import polyglot.main.Report;
import polyglot.types.ClassType;
import polyglot.types.InitializerInstance;
import polyglot.types.MemberInstance;
import polyglot.types.MethodInstance;
import polyglot.types.ProcedureInstance;
import polyglot.types.ReferenceType;
import polyglot.types.SemanticException;
import accrue.AccrueOptions;
import accrue.AccrueUtil;
import accrue.analysis.AnalysisTopics;

/**
 * The Registrar maintains useful information about the whole program's code.
 * It has a map from ProcedureInstances to CodeDecls, and a collection of start procedures
 * (i.e., main methods that will be the starting point for interprocedural analyses).
 * It also maintains for each class type a list of the instance initializers (including field initializers)
 * and static initializers (including field initializers). 
 */
public class Registrar {
    private Map<ProcedureInstance, CodeDecl> code =
            new LinkedHashMap<ProcedureInstance, CodeDecl>();
    private Set<ProcedureInstance> startProcs =
            new LinkedHashSet<ProcedureInstance>();
    /**
     * For each class, a list of instance initializers and instance field declarations.
     */
    private Map<ClassType, List<CodeNode>> instanceInitializers =
            new LinkedHashMap<ClassType, List<CodeNode>>();
    /**
     * For each class, a list of static initializers and static field declarations.
     */
    private Map<ClassType, List<CodeNode>> staticInitializers =
            new LinkedHashMap<ClassType, List<CodeNode>>();
    /**
     * A list of all the static initializers and field decls, in the order they are analyzed.
     */
    private List<CodeNode> allOrderedStaticInitializers = null;
    private final ExtensionInfo extInfo;

    public Registrar(ExtensionInfo extInfo) {
        this.extInfo = extInfo;
    }

    public void registerMethod(MethodDecl n) throws SemanticException {
        ProcedureInstance pi = n.procedureInstance();

        if (Report.should_report(AnalysisTopics.REGISTRAR, 1))
            Report.report(1, "Registering " + pi);

        if (n.flags().isAbstract()) {
            //throw new InternalCompilerError("Whoops! Got an abstract method!");
            // nothing to do with an abstract method...
            return;
        }

        code.put(pi, n);
        if (isStartProc(pi)) {
            startProcs.add(pi);

            // for the moment, restrict it to one start proc...
            // XXX could allow the user to select which start proc by specifying
            // a distinguished class on the command line.
            if (startProcs.size() > 1) {
                throw new SemanticException("More than one main method: "
                        + startProcs);
            }
        }
    }

    public void registerConstructor(ConstructorDecl n) {
        if (Report.should_report(AnalysisTopics.REGISTRAR, 1))
            Report.report(1, "Registering " + n.procedureInstance());

        code.put(n.procedureInstance(), n);
    }

    public void registerInitializer(Initializer n) {
        if (Report.should_report(AnalysisTopics.REGISTRAR, 1))
            Report.report(1, "Registering " + n.initializerInstance());

        registerInitializer(n, n.initializerInstance());
    }

    public void registerFieldDecl(FieldDecl n) {
        if (Report.should_report(AnalysisTopics.REGISTRAR, 1))
            Report.report(1,
                          "Registering field decl " + n.initializerInstance());
        registerInitializer(n, n.initializerInstance());
    }

    protected void registerInitializer(CodeNode n, InitializerInstance ii) {
        if (ii != null) {
            Map<ClassType, List<CodeNode>> initializerList =
                    ii.flags().isStatic() ? this.staticInitializers
                            : this.instanceInitializers;

            ClassType ct = ii.container().toClass();
            List<CodeNode> l = initializerList.get(ct);
            if (l == null) {
                l = new ArrayList<CodeNode>();
                initializerList.put(ct, l);
            }
            l.add(n);
        }
    }

    public CodeDecl getCode(ProcedureInstance pi) {
        return code.get(pi);
    }

    public boolean hasCode(ProcedureInstance pi) {
        return code.containsKey(pi);
    }

    public List<CodeNode> getInstanceInitializers(ClassType container) {
        if (container == null) throw new IllegalArgumentException();
        List<CodeNode> l = instanceInitializers.get(container);
        if (l == null) {
            return Collections.emptyList();
        }
        return l;
    }

    public Set<ClassType> getClassesWithStaticInitializers() {
        return staticInitializers.keySet();
    }

    public List<CodeNode> getStaticInitializers(ClassType container) {
        if (container == null) throw new IllegalArgumentException();
        List<CodeNode> l = staticInitializers.get(container);
        if (l == null) {
            return Collections.emptyList();
        }
        return l;
    }

    public void setAllOrderedStaticInitializers(List<CodeNode> staticInits) {
        this.allOrderedStaticInitializers = staticInits;
    }

    public List<CodeNode> getAllOrderedStaticInitializers() {
        return this.allOrderedStaticInitializers;
    }

    /**
     * Is pi a starting point for the analysis?
     * @return
     */
    public static boolean isStartProc(ProcedureInstance pi) {
        // main methods are currently the only start procs
        // Static initializers are not start procedures: they will be dealt with when main methods
        if (!AccrueUtil.isMainMethod(pi)) {
            return false;
        }

        ExtensionInfo extInfo = pi.typeSystem().extensionInfo();
        AccrueOptions oao = (AccrueOptions) extInfo.getOptions();
        if (oao.helper().startProcClass == null) {
            // we are interested in any main method
            return true;
        }
        return ((MethodInstance) pi).container()
                                    .toClass()
                                    .fullName()
                                    .equals(oao.helper().startProcClass);
    }

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

    public Collection<CodeNode> allCode() {
        ArrayList<CodeNode> c =
                new ArrayList<CodeNode>(this.code.size()
                        + instanceInitializers.size()
                        + staticInitializers.size());
        c.addAll(this.code.values());
        for (List<CodeNode> li : this.instanceInitializers.values()) {
            c.addAll(li);
        }
        for (List<CodeNode> li : this.staticInitializers.values()) {
            c.addAll(li);
        }
        return c;
    }

    public int numRegisteredClasses() {
        Set<ClassType> classes = new HashSet<ClassType>();
        for (ProcedureInstance pi : code.keySet()) {
            if (pi instanceof MemberInstance) {
                ReferenceType container = ((MemberInstance) pi).container();
                if (container.isClass()) {
                    classes.add(container.toClass());

                }
            }
        }
        return classes.size();
    }

    public int numRegisteredProcedures() {
        return code.size();
    }
}
