/*******************************************************************************
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/
package org.apache.sling.scripting.sightly.compiler.java.utils;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import javax.tools.Diagnostic;
import javax.tools.DiagnosticCollector;
import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaCompiler;
import javax.tools.JavaCompiler.CompilationTask;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.JavaFileObject.Kind;
import javax.tools.SimpleJavaFileObject;
import javax.tools.StandardLocation;
import javax.tools.ToolProvider;

/**
 * This compiler is a slightly simplified version of the CharSequenceCompiler described at
 * http://www.ibm.com/developerworks/java/library/j-jcomp/index.html.
 */
public class CharSequenceJavaCompiler<T> {
    // Compiler requires source files with a ".java" extension:
    static final String JAVA_EXTENSION = ".java";

    private final ClassLoaderImpl classLoader;

    // The compiler instance that this facade uses.
    private final JavaCompiler compiler;

    // The compiler options (such as "-target" "1.5").
    private final List<String> options;

    // collect compiler diagnostics in this instance.
    private DiagnosticCollector<JavaFileObject> diagnostics;

    // The FileManager which will store source and class "files".
    private final FileManagerImpl javaFileManager;

    /**
     * Construct a new instance which delegates to the named class loader.
     *
     * @param loader
     *           the application ClassLoader. The compiler will look through to
     *           this // class loader for dependent classes
     * @param options
     *           The compiler options (such as "-target" "1.5"). See the usage
     *           for javac
     * @throws IllegalStateException
     *            if the Java compiler cannot be loaded.
     */
    public CharSequenceJavaCompiler(ClassLoader loader, Iterable<String> options) {
        compiler = ToolProvider.getSystemJavaCompiler();
        if (compiler == null) {
            throw new IllegalStateException("Cannot find the system Java compiler. "
                    + "Check that your class path includes tools.jar");
        }
        classLoader = new ClassLoaderImpl(loader);
        diagnostics = new DiagnosticCollector<>();
        final JavaFileManager fileManager = compiler.getStandardFileManager(diagnostics,
                null, null);
        // create our FileManager which chains to the default file manager
        // and our ClassLoader
        javaFileManager = new FileManagerImpl(fileManager, classLoader);
        this.options = new ArrayList<>();
        if (options != null) { // make a save copy of input options
            for (String option : options) {
                this.options.add(option);
            }
        }
    }

    /**
     * Compile Java source in <var>javaSource</name> and return the resulting
     * class.
     * <p>
     * Thread safety: this method is thread safe if the <var>javaSource</var>
     * and <var>diagnosticsList</var> are isolated to this thread.
     *
     * @param qualifiedClassName
     *           The fully qualified class name.
     * @param javaSource
     *           Complete java source, including a package statement and a class,
     *           interface, or annotation declaration.
     * @param types
     *           zero or more Class objects representing classes or interfaces
     *           that the resulting class must be assignable (castable) to.
     * @return a Class which is generated by compiling the source
     * @throws CharSequenceJavaCompilerException
     *            if the source cannot be compiled - for example, if it contains
     *            syntax or semantic errors or if dependent classes cannot be
     *            found.
     * @throws ClassCastException
     *            if the generated class is not assignable to all the optional
     *            <var>types</var>.
     */
    public synchronized Class<T> compile(final String qualifiedClassName,
                                         final CharSequence javaSource,
                                         final Class<?>... types) throws ClassCastException {
        Map<String, CharSequence> classes = new HashMap<>(1);
        classes.put(qualifiedClassName, javaSource);
        try {
            Map<String, Class<T>> compiled = compile(classes);
            Class<T> newClass = compiled.get(qualifiedClassName);
            return castable(newClass, types);
        } catch (CharSequenceJavaCompilerException e) {
            StringBuilder stringBuilder = new StringBuilder();
            for (Diagnostic diagnostic : e.getDiagnostics().getDiagnostics()) {
                stringBuilder.append(diagnostic.toString()).append(System.lineSeparator());
            }
            throw new RuntimeException(stringBuilder.toString());
        }
    }

    /**
     * Compile multiple Java source strings and return a Map containing the
     * resulting classes.
     * <p>
     * Thread safety: this method is thread safe if the <var>classes</var> and
     * <var>diagnosticsList</var> are isolated to this thread.
     *
     * @param classes
     *           A Map whose keys are qualified class names and whose values are
     *           the Java source strings containing the definition of the class.
     *           A map value may be null, indicating that compiled class is
     *           expected, although no source exists for it (it may be a
     *           non-public class contained in one of the other strings.)
     * @return A mapping of qualified class names to their corresponding classes.
     *         The map has the same keys as the input <var>classes</var>; the
     *         values are the corresponding Class objects.
     * @throws CharSequenceJavaCompilerException
     *            if the source cannot be compiled
     */
    public synchronized Map<String, Class<T>> compile(final Map<String, CharSequence> classes) throws CharSequenceJavaCompilerException {
        List<JavaFileObject> sources = new ArrayList<>();
        for (Entry<String, CharSequence> entry : classes.entrySet()) {
            String qualifiedClassName = entry.getKey();
            CharSequence javaSource = entry.getValue();
            if (javaSource != null) {
                final int dotPos = qualifiedClassName.lastIndexOf('.');
                final String className = dotPos == -1 ? qualifiedClassName
                        : qualifiedClassName.substring(dotPos + 1);
                final String packageName = dotPos == -1 ? "" : qualifiedClassName
                        .substring(0, dotPos);
                final JavaFileObjectImpl source = new JavaFileObjectImpl(className,
                        javaSource);
                sources.add(source);
                // Store the source file in the FileManager via package/class
                // name.
                // For source files, we add a .java extension
                javaFileManager.putFileForInput(StandardLocation.SOURCE_PATH, packageName,
                        className + JAVA_EXTENSION, source);
            }
        }
        // Get a CompliationTask from the compiler and compile the sources
        final CompilationTask task = compiler.getTask(null, javaFileManager, diagnostics,
                options, null, sources);
        final Boolean result = task.call();
        if (result == null || !result) {
            throw new CharSequenceJavaCompilerException("Compilation failed.", classes
                    .keySet(), diagnostics);
        }
        try {
            // For each class name in the input map, get its compiled
            // class and put it in the output map
            Map<String, Class<T>> compiled = new HashMap<>();
            for (String qualifiedClassName : classes.keySet()) {
                final Class<T> newClass = loadClass(qualifiedClassName);
                compiled.put(qualifiedClassName, newClass);
            }
            return compiled;
        } catch (ClassNotFoundException e) {
            throw new CharSequenceJavaCompilerException(classes.keySet(), e, diagnostics);
        }
    }

    /**
     * Load a class that was generated by this instance or accessible from its
     * parent class loader. Use this method if you need access to additional
     * classes compiled by
     * {@link #compile(String, CharSequence, Class...) compile()},
     * for example if the primary class contained nested classes or additional
     * non-public classes.
     *
     * @param qualifiedClassName
     *           the name of the compiled class you wish to load
     * @return a Class instance named by <var>qualifiedClassName</var>
     * @throws ClassNotFoundException
     *            if no such class is found.
     */
    @SuppressWarnings("unchecked")
    public Class<T> loadClass(final String qualifiedClassName)
            throws ClassNotFoundException {
        return (Class<T>) classLoader.loadClass(qualifiedClassName);
    }

    /**
     * Check that the <var>newClass</var> is a subtype of all the type
     * parameters and throw a ClassCastException if not.
     *
     * @param types
     *           zero of more classes or interfaces that the <var>newClass</var>
     *           must be castable to.
     * @return <var>newClass</var> if it is castable to all the types
     * @throws ClassCastException
     *            if <var>newClass</var> is not castable to all the types.
     */
    private Class<T> castable(Class<T> newClass, Class<?>... types)
            throws ClassCastException {
        for (Class<?> type : types)
            if (!type.isAssignableFrom(newClass)) {
                throw new ClassCastException(type.getName());
            }
        return newClass;
    }

    /**
     * Converts a String to a URI.
     *
     * @param name
     *           a file name
     * @return a URI
     */
    static URI toURI(String name) {
        try {
            return new URI(name);
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }
}

/**
 * A JavaFileManager which manages Java source and classes. This FileManager
 * delegates to the JavaFileManager and the ClassLoaderImpl provided in the
 * constructor. The sources are all in memory CharSequence instances and the
 * classes are all in memory byte arrays.
 */
final class FileManagerImpl extends ForwardingJavaFileManager<JavaFileManager> {
    // the delegating class loader (passed to the constructor)
    private final ClassLoaderImpl classLoader;

    // Internal map of filename URIs to JavaFileObjects.
    private final Map<URI, JavaFileObject> fileObjects = new HashMap<>();

    /**
     * Construct a new FileManager which forwards to the <var>fileManager</var>
     * for source and to the <var>classLoader</var> for classes
     *
     * @param fileManager
     *           another FileManager that this instance delegates to for
     *           additional source.
     * @param classLoader
     *           a ClassLoader which contains dependent classes that the compiled
     *           classes will require when compiling them.
     */
    public FileManagerImpl(JavaFileManager fileManager, ClassLoaderImpl classLoader) {
        super(fileManager);
        this.classLoader = classLoader;
    }

    /**
     * For a given file <var>location</var>, return a FileObject from which the
     * compiler can obtain source or byte code.
     *
     * @param location
     *           an abstract file location
     * @param packageName
     *           the package name for the file
     * @param relativeName
     *           the file's relative name
     * @return a FileObject from this or the delegated FileManager
     * @see javax.tools.ForwardingJavaFileManager#getFileForInput(javax.tools.JavaFileManager.Location,
     *      java.lang.String, java.lang.String)
     */
    @Override
    public FileObject getFileForInput(Location location, String packageName,
                                      String relativeName) throws IOException {
        FileObject o = fileObjects.get(uri(location, packageName, relativeName));
        if (o != null)
            return o;
        return super.getFileForInput(location, packageName, relativeName);
    }

    /**
     * Store a file that may be retrieved later with
     * {@link #getFileForInput(javax.tools.JavaFileManager.Location, String, String)}
     *
     * @param location
     *           the file location
     * @param packageName
     *           the Java class' package name
     * @param relativeName
     *           the relative name
     * @param file
     *           the file object to store for later retrieval
     */
    public void putFileForInput(StandardLocation location, String packageName,
                                String relativeName, JavaFileObject file) {
        fileObjects.put(uri(location, packageName, relativeName), file);
    }

    /**
     * Convert a location and class name to a URI
     */
    private URI uri(Location location, String packageName, String relativeName) {
        return CharSequenceJavaCompiler.toURI(location.getName() + '/' + packageName + '/'
                + relativeName);
    }

    /**
     * Create a JavaFileImpl for an output class file and store it in the
     * classloader.
     *
     * @see javax.tools.ForwardingJavaFileManager#getJavaFileForOutput(javax.tools.JavaFileManager.Location,
     *      java.lang.String, javax.tools.JavaFileObject.Kind,
     *      javax.tools.FileObject)
     */
    @Override
    public JavaFileObject getJavaFileForOutput(Location location, String qualifiedName,
                                               Kind kind, FileObject outputFile) throws IOException {
        JavaFileObject file = new JavaFileObjectImpl(qualifiedName, kind);
        classLoader.add(qualifiedName, file);
        return file;
    }

    @Override
    public ClassLoader getClassLoader(JavaFileManager.Location location) {
        return classLoader;
    }

    @Override
    public String inferBinaryName(Location loc, JavaFileObject file) {
        String result;
        // For our JavaFileImpl instances, return the file's name, else
        // simply run the default implementation
        if (file instanceof JavaFileObjectImpl)
            result = file.getName();
        else
            result = super.inferBinaryName(loc, file);
        return result;
    }

    @Override
    public Iterable<JavaFileObject> list(Location location, String packageName,
                                         Set<Kind> kinds, boolean recurse) throws IOException {
        Iterable<JavaFileObject> result = super.list(location, packageName, kinds,
                recurse);
        ArrayList<JavaFileObject> files = new ArrayList<>();
        if (location == StandardLocation.CLASS_PATH
                && kinds.contains(JavaFileObject.Kind.CLASS)) {
            for (JavaFileObject file : fileObjects.values()) {
                if (file.getKind() == Kind.CLASS && file.getName().startsWith(packageName))
                    files.add(file);
            }
            files.addAll(classLoader.files());
        } else if (location == StandardLocation.SOURCE_PATH
                && kinds.contains(JavaFileObject.Kind.SOURCE)) {
            for (JavaFileObject file : fileObjects.values()) {
                if (file.getKind() == Kind.SOURCE && file.getName().startsWith(packageName))
                    files.add(file);
            }
        }
        for (JavaFileObject file : result) {
            files.add(file);
        }
        return files;
    }
}

/**
 * A JavaFileObject which contains either the source text or the compiler
 * generated class. This class is used in two cases.
 * <ol>
 * <li>This instance uses it to store the source which is passed to the
 * compiler. This uses the
 * {@link JavaFileObjectImpl#JavaFileObjectImpl(String, CharSequence)}
 * constructor.
 * <li>The Java compiler also creates instances (indirectly through the
 * FileManagerImplFileManager) when it wants to create a JavaFileObject for the
 * .class output. This uses the
 * {@link JavaFileObjectImpl#JavaFileObjectImpl(String, JavaFileObject.Kind)}
 * constructor.
 * </ol>
 * This class does not attempt to reuse instances (there does not seem to be a
 * need, as it would require adding a Map for the purpose, and this would also
 * prevent garbage collection of class byte code.)
 */
final class JavaFileObjectImpl extends SimpleJavaFileObject {
    // If kind == CLASS, this stores byte code from openOutputStream
    private ByteArrayOutputStream byteCode;

    // if kind == SOURCE, this contains the source text
    private final CharSequence source;

    /**
     * Construct a new instance which stores source
     *
     * @param baseName
     *           the base name
     * @param source
     *           the source code
     */
    JavaFileObjectImpl(final String baseName, final CharSequence source) {
        super(CharSequenceJavaCompiler.toURI(baseName + CharSequenceJavaCompiler.JAVA_EXTENSION),
                Kind.SOURCE);
        this.source = source;
    }

    /**
     * Construct a new instance
     *
     * @param name
     *           the file name
     * @param kind
     *           the kind of file
     */
    JavaFileObjectImpl(final String name, final Kind kind) {
        super(CharSequenceJavaCompiler.toURI(name), kind);
        source = null;
    }

    /**
     * Return the source code content
     *
     * @see javax.tools.SimpleJavaFileObject#getCharContent(boolean)
     */
    @Override
    public CharSequence getCharContent(final boolean ignoreEncodingErrors)
            throws UnsupportedOperationException {
        if (source == null)
            throw new UnsupportedOperationException("getCharContent()");
        return source;
    }

    /**
     * Return an input stream for reading the byte code
     *
     * @see javax.tools.SimpleJavaFileObject#openInputStream()
     */
    @Override
    public InputStream openInputStream() {
        return new ByteArrayInputStream(getByteCode());
    }

    /**
     * Return an output stream for writing the bytecode
     *
     * @see javax.tools.SimpleJavaFileObject#openOutputStream()
     */
    @Override
    public OutputStream openOutputStream() {
        byteCode = new ByteArrayOutputStream();
        return byteCode;
    }

    /**
     * @return the byte code generated by the compiler
     */
    byte[] getByteCode() {
        return byteCode.toByteArray();
    }
}

/**
 * A custom ClassLoader which maps class names to JavaFileObjectImpl instances.
 */
final class ClassLoaderImpl extends ClassLoader {
    private final Map<String, JavaFileObject> classes = new HashMap<>();

    ClassLoaderImpl(final ClassLoader parentClassLoader) {
        super(parentClassLoader);
    }

    /**
     * @return An collection of JavaFileObject instances for the classes in the
     *         class loader.
     */
    Collection<JavaFileObject> files() {
        return Collections.unmodifiableCollection(classes.values());
    }

    @Override
    protected Class<?> findClass(final String qualifiedClassName)
            throws ClassNotFoundException {
        JavaFileObject file = classes.get(qualifiedClassName);
        if (file != null) {
            byte[] bytes = ((JavaFileObjectImpl) file).getByteCode();
            return defineClass(qualifiedClassName, bytes, 0, bytes.length);
        }
        try {
            return Class.forName(qualifiedClassName);
        } catch (ClassNotFoundException nf) {
            // Ignore and fall through
        }
        return super.findClass(qualifiedClassName);
    }

    /**
     * Add a class name/JavaFileObject mapping
     *
     * @param qualifiedClassName
     *           the name
     * @param javaFile
     *           the file associated with the name
     */
    void add(final String qualifiedClassName, final JavaFileObject javaFile) {
        classes.put(qualifiedClassName, javaFile);
    }

    @Override
    protected synchronized Class<?> loadClass(final String name, final boolean resolve)
            throws ClassNotFoundException {
        return super.loadClass(name, resolve);
    }

    @Override
    public InputStream getResourceAsStream(final String name) {
        if (name.endsWith(".class")) {
            String qualifiedClassName = name.substring(0,
                    name.length() - ".class".length()).replace('/', '.');
            JavaFileObjectImpl file = (JavaFileObjectImpl) classes.get(qualifiedClassName);
            if (file != null) {
                return new ByteArrayInputStream(file.getByteCode());
            }
        }
        return super.getResourceAsStream(name);
    }
}