package org.checkerframework.javacutil;

import static com.sun.tools.javac.code.Kinds.PCK;
import static com.sun.tools.javac.code.Kinds.TYP;
import static com.sun.tools.javac.code.Kinds.VAR;

import com.sun.source.util.TreePath;
import com.sun.source.util.Trees;
import com.sun.tools.javac.api.JavacScope;
import com.sun.tools.javac.code.Symbol;
import com.sun.tools.javac.code.Symbol.ClassSymbol;
import com.sun.tools.javac.code.Symbol.PackageSymbol;
import com.sun.tools.javac.code.Symbol.TypeSymbol;
import com.sun.tools.javac.code.Type;
import com.sun.tools.javac.comp.AttrContext;
import com.sun.tools.javac.comp.DeferredAttr;
import com.sun.tools.javac.comp.Env;
import com.sun.tools.javac.comp.Resolve;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.Log;
import com.sun.tools.javac.util.Name;
import com.sun.tools.javac.util.Names;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.TypeMirror;

/** A Utility class to find symbols corresponding to string references. */
public class Resolver {
    private final Resolve resolve;
    private final Names names;
    private final Trees trees;
    private final Log log;

    private static final Method FIND_METHOD;
    private static final Method FIND_VAR;
    private static final Method FIND_IDENT;
    private static final Method FIND_IDENT_IN_TYPE;
    private static final Method FIND_IDENT_IN_PACKAGE;
    private static final Method FIND_TYPE;

    private static final Class<?> ACCESSERROR;
    // Note that currently access(...) is defined in InvalidSymbolError, a superclass of AccessError
    private static final Method ACCESSERROR_ACCESS;

    static {
        try {
            FIND_METHOD =
                    Resolve.class.getDeclaredMethod(
                            "findMethod",
                            Env.class,
                            Type.class,
                            Name.class,
                            List.class,
                            List.class,
                            boolean.class,
                            boolean.class,
                            boolean.class);
            FIND_METHOD.setAccessible(true);

            FIND_VAR = Resolve.class.getDeclaredMethod("findVar", Env.class, Name.class);
            FIND_VAR.setAccessible(true);

            FIND_IDENT =
                    Resolve.class.getDeclaredMethod("findIdent", Env.class, Name.class, int.class);
            FIND_IDENT.setAccessible(true);

            FIND_IDENT_IN_TYPE =
                    Resolve.class.getDeclaredMethod(
                            "findIdentInType", Env.class, Type.class, Name.class, int.class);
            FIND_IDENT_IN_TYPE.setAccessible(true);

            FIND_IDENT_IN_PACKAGE =
                    Resolve.class.getDeclaredMethod(
                            "findIdentInPackage",
                            Env.class,
                            TypeSymbol.class,
                            Name.class,
                            int.class);
            FIND_IDENT_IN_PACKAGE.setAccessible(true);

            FIND_TYPE = Resolve.class.getDeclaredMethod("findType", Env.class, Name.class);
            FIND_TYPE.setAccessible(true);
        } catch (Exception e) {
            Error err =
                    new AssertionError(
                            "Compiler 'Resolve' class doesn't contain required 'find' method");
            err.initCause(e);
            throw err;
        }

        try {
            ACCESSERROR = Class.forName("com.sun.tools.javac.comp.Resolve$AccessError");
            ACCESSERROR_ACCESS = ACCESSERROR.getMethod("access", Name.class, TypeSymbol.class);
            ACCESSERROR_ACCESS.setAccessible(true);
        } catch (ClassNotFoundException e) {
            ErrorReporter.errorAbort(
                    "Compiler 'Resolve$AccessError' class could not be retrieved.", e);
            // Unreachable code - needed so the compiler does not warn about a possibly uninitialized final field.
            throw new AssertionError();
        } catch (NoSuchMethodException e) {
            ErrorReporter.errorAbort(
                    "Compiler 'Resolve$AccessError' class doesn't contain required 'access' method",
                    e);
            // Unreachable code - needed so the compiler does not warn about a possibly uninitialized final field.
            throw new AssertionError();
        }
    }

    public Resolver(ProcessingEnvironment env) {
        Context context = ((JavacProcessingEnvironment) env).getContext();
        this.resolve = Resolve.instance(context);
        this.names = Names.instance(context);
        this.trees = Trees.instance(env);
        this.log = Log.instance(context);
    }

    /**
     * Determine the environment for the given path.
     *
     * @param path the tree path to the local scope
     * @return the corresponding attribution environment
     */
    public Env<AttrContext> getEnvForPath(TreePath path) {
        TreePath iter = path;
        JavacScope scope = null;
        while (scope == null && iter != null) {
            try {
                scope = (JavacScope) trees.getScope(iter);
            } catch (Throwable t) {
                // Work around Issue #1059 by skipping through the TreePath until something
                // doesn't crash. This probably returns the class scope, so users might not
                // get the variables they expect. But that is better than crashing.
                iter = iter.getParentPath();
            }
        }
        if (scope != null) {
            return scope.getEnv();
        } else {
            ErrorReporter.errorAbort(
                    "Could not determine any possible scope for path: " + path.getLeaf());
            return null;
        }
    }

    /**
     * Finds the package with name {@code name}.
     *
     * @param name the name of the package
     * @param path the tree path to the local scope
     * @return the {@code PackageSymbol} for the package if it is found, {@code null} otherwise
     */
    public PackageSymbol findPackage(String name, TreePath path) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);
            Element res =
                    wrapInvocationOnResolveInstance(FIND_IDENT, env, names.fromString(name), PCK);
            // findIdent will return a PackageSymbol even for a symbol that is not a package,
            // such as a.b.c.MyClass.myStaticField. "exists()" must be called on it to ensure
            // that it exists.
            if (res.getKind() == ElementKind.PACKAGE) {
                PackageSymbol ps = (PackageSymbol) res;
                return ps.exists() ? ps : null;
            } else {
                return null;
            }
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /**
     * Finds the field with name {@code name} in a given type.
     *
     * <p>The method adheres to all the rules of Java's scoping (while also considering the imports)
     * for name resolution.
     *
     * @param name the name of the field
     * @param type the type of the receiver (i.e., the type in which to look for the field).
     * @param path the tree path to the local scope
     * @return the element for the field
     */
    public VariableElement findField(String name, TypeMirror type, TreePath path) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);
            Element res =
                    wrapInvocationOnResolveInstance(
                            FIND_IDENT_IN_TYPE, env, type, names.fromString(name), VAR);
            if (res.getKind() == ElementKind.FIELD) {
                return (VariableElement) res;
            } else if (res.getKind() == ElementKind.OTHER && ACCESSERROR.isInstance(res)) {
                // Return the inaccessible field that was found
                return (VariableElement) wrapInvocation(res, ACCESSERROR_ACCESS, null, null);
            } else {
                // Most likely didn't find the field and the Element is a SymbolNotFoundError
                return null;
            }
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /**
     * Finds the local variable with name {@code name} in the given scope.
     *
     * @param name the name of the local variable
     * @param path the tree path to the local scope
     * @return the element for the local variable
     */
    public VariableElement findLocalVariableOrParameterOrField(String name, TreePath path) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);
            Element res = wrapInvocationOnResolveInstance(FIND_VAR, env, names.fromString(name));
            if (res.getKind() == ElementKind.LOCAL_VARIABLE
                    || res.getKind() == ElementKind.PARAMETER
                    || res.getKind() == ElementKind.FIELD) {
                return (VariableElement) res;
            } else {
                // Most likely didn't find the variable and the Element is a SymbolNotFoundError
                return null;
            }
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /**
     * Finds the class literal with name {@code name}.
     *
     * <p>The method adheres to all the rules of Java's scoping (while also considering the imports)
     * for name resolution.
     *
     * @param name the name of the class
     * @param path the tree path to the local scope
     * @return the element for the class
     */
    public Element findClass(String name, TreePath path) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);
            return wrapInvocationOnResolveInstance(FIND_TYPE, env, names.fromString(name));
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /**
     * Finds the class with name {@code name} in a given package.
     *
     * @param name the name of the class
     * @param pck the PackageSymbol for the package
     * @param path the tree path to the local scope
     * @return the {@code ClassSymbol} for the class if it is found, {@code null} otherwise
     */
    public ClassSymbol findClassInPackage(String name, PackageSymbol pck, TreePath path) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);
            Element res =
                    wrapInvocationOnResolveInstance(
                            FIND_IDENT_IN_PACKAGE, env, pck, names.fromString(name), TYP);
            if (res.getKind() == ElementKind.CLASS) {
                return (ClassSymbol) res;
            } else {
                return null;
            }
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /**
     * Finds the method element for a given name and list of expected parameter types.
     *
     * <p>The method adheres to all the rules of Java's scoping (while also considering the imports)
     * for name resolution.
     *
     * @param methodName name of the method to find
     * @param receiverType type of the receiver of the method
     * @param path tree path
     * @return the method element (if found)
     */
    public Element findMethod(
            String methodName,
            TypeMirror receiverType,
            TreePath path,
            java.util.List<TypeMirror> argumentTypes) {
        Log.DiagnosticHandler discardDiagnosticHandler = new Log.DiscardDiagnosticHandler(log);
        try {
            Env<AttrContext> env = getEnvForPath(path);

            Type site = (Type) receiverType;
            Name name = names.fromString(methodName);
            List<Type> argtypes = List.nil();
            for (TypeMirror a : argumentTypes) {
                argtypes = argtypes.append((Type) a);
            }
            List<Type> typeargtypes = List.nil();
            boolean allowBoxing = true;
            boolean useVarargs = false;
            boolean operator = true;

            try {
                // For some reason we have to set our own method context, which is rather ugly.
                // TODO: find a nicer way to do this.
                Object methodContext = buildMethodContext();
                Object oldContext = getField(resolve, "currentResolutionContext");
                setField(resolve, "currentResolutionContext", methodContext);
                Element result =
                        wrapInvocationOnResolveInstance(
                                FIND_METHOD,
                                env,
                                site,
                                name,
                                argtypes,
                                typeargtypes,
                                allowBoxing,
                                useVarargs,
                                operator);
                setField(resolve, "currentResolutionContext", oldContext);
                return result;
            } catch (Throwable t) {
                Error err = new AssertionError("Unexpected Reflection error");
                err.initCause(t);
                throw err;
            }
        } finally {
            log.popDiagnosticHandler(discardDiagnosticHandler);
        }
    }

    /** Build an instance of {@code Resolve$MethodResolutionContext}. */
    protected Object buildMethodContext()
            throws ClassNotFoundException, InstantiationException, IllegalAccessException,
                    InvocationTargetException, NoSuchFieldException {
        // Class is not accessible, instantiate reflectively.
        Class<?> methCtxClss =
                Class.forName("com.sun.tools.javac.comp.Resolve$MethodResolutionContext");
        Constructor<?> constructor = methCtxClss.getDeclaredConstructors()[0];
        constructor.setAccessible(true);
        Object methodContext = constructor.newInstance(resolve);
        // we need to also initialize the fields attrMode and step
        setField(methodContext, "attrMode", DeferredAttr.AttrMode.CHECK);
        @SuppressWarnings("rawtypes")
        List<?> phases = (List) getField(resolve, "methodResolutionSteps");
        setField(methodContext, "step", phases.get(1));
        return methodContext;
    }

    /** Reflectively set a field. */
    private void setField(Object receiver, String fieldName, Object value)
            throws NoSuchFieldException, IllegalAccessException {
        Field f = receiver.getClass().getDeclaredField(fieldName);
        f.setAccessible(true);
        f.set(receiver, value);
    }

    /** Reflectively get the value of a field. */
    private Object getField(Object receiver, String fieldName)
            throws NoSuchFieldException, IllegalAccessException {
        Field f = receiver.getClass().getDeclaredField(fieldName);
        f.setAccessible(true);
        return f.get(receiver);
    }

    private Symbol wrapInvocationOnResolveInstance(Method method, Object... args) {
        return wrapInvocation(resolve, method, args);
    }

    private Symbol wrapInvocation(Object receiver, Method method, Object... args) {
        try {
            return (Symbol) method.invoke(receiver, args);
        } catch (IllegalAccessException e) {
            Error err = new AssertionError("Unexpected Reflection error");
            err.initCause(e);
            throw err;
        } catch (IllegalArgumentException e) {
            Error err = new AssertionError("Unexpected Reflection error");
            err.initCause(e);
            throw err;
        } catch (InvocationTargetException e) {
            Error err = new AssertionError("Unexpected Reflection error");
            err.initCause(e);
            throw err;
        }
    }
}