package org.qore.jni.compiler;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import java.nio.charset.StandardCharsets;

import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLClassLoader;
import java.net.URLDecoder;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

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.StandardJavaFileManager;
import javax.tools.StandardLocation;
import javax.tools.ToolProvider;

import org.qore.jni.QoreURLClassLoader;
import org.qore.jni.QoreJavaFileObject;

/**
 * Compile a String or other {@link CharSequence}, returning a Java
 * {@link Class} instance that may be instantiated as well as the raw bytecode. This class is a Facade
 * around {@link JavaCompiler} for a narrower use case, but a bit easier to use.
 * <p/>
 * To compile a String containing source for a Java class which implements
 * MyInterface:
 * <p/>
 * <pre>
 * ClassLoader classLoader = MyClass.class.getClassLoader(); // optional; null is also OK
 * List&lt;Diagnostic&gt; diagnostics = new ArrayList&lt;Diagnostic&gt;(); // optional; null is also OK
 * QoreJavaCompiler&lt;Object&gt; compiler = new QoreJavaCompiler&lt;MyInterface&gt;(classLoader,
 *       null);
 * try {
 *    Class&lt;MyInterface&gt; newClass = compiler.compile(&quot;com.mypackage.NewClass&quot;,
 *          stringContainingSourceForNewClass, diagnostics, MyInterface).cls;
 *    MyInterface instance = newClass.newInstance();
 *    instance.someOperation(someArgs);
 * } catch (QoreJavaCompilerException e) {
 *    handle(e);
 * } catch (IllegalAccessException e) {
 *    handle(e);
 * }
 * </pre>
 * <p/>
 * The source can be in a String, {@link StringBuffer}, or your own class which
 * implements {@link CharSequence}. If you implement your own, it must be
 * thread safe (preferably, immutable.)
 *
 * @author <a href="mailto:David.Biesack@sas.com">David J. Biesack</a>, adapted for %Qore by
 * <a href="mailto:david@qore.org">David Nichols</a>
 */
public class QoreJavaCompiler<T> {
    // Compiler requires source files with a ".java" extension:
    static final String JAVA_EXTENSION = ".java";

    private final QoreURLClassLoader 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 a new Qore classloader.
     *
     * @throws IllegalStateException if the Java compiler cannot be loaded.
     */
    public QoreJavaCompiler() {
        this((Iterable<String>)null);
    }

    /**
     * Construct a new instance which delegates to a new Qore classloader.
     *
     * @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 QoreJavaCompiler(Iterable<String> options) {
        this(QoreURLClassLoader.getCurrent(), options);
    }

    /**
     * Construct a new instance which delegates to a new Qore classloader.
     *
     * @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 QoreJavaCompiler(String[] options) {
        this(QoreURLClassLoader.getCurrent(), options == null ? (Iterable<String>)null : Arrays.asList(options));
    }

    /**
     * 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 QoreJavaCompiler(QoreURLClassLoader 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 QoreURLClassLoader(loader.getPtr(), loader);
        //System.out.printf("compiler classLoader: %s (loader: %s pgm: %x)\n", classLoader, loader, loader.getPtr());
        diagnostics = null;
        final QoreJavaFileManager fileManager = new QoreJavaFileManager(classLoader,
            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<String>();
        try {
            List<File> pathlist = new ArrayList<File>();
            if (options != null) { // save a copy of input options
                boolean cp_next = false;
                for (String option : options) {
                    this.options.add(option);
                }
            }

            // add locations from the parent class loader
            for (URL url : loader.getURLs()) {
                String path = URLDecoder.decode(url.getPath(), StandardCharsets.UTF_8);
                File file = new File(path);
                //System.out.printf("set location url: %s\n", file);
                pathlist.add(file);
            }

            if (pathlist.size() > 0) {
                fileManager.setLocation(StandardLocation.CLASS_PATH, pathlist);
            }
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    /**
     * 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>
     * is 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 the Class and file data generated by compiling the source
     *
     * @throws QoreJavaCompilerException 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 CompilerOutput<T> compile(final String qualifiedClassName,
            final String javaSource,
            final Class<?>... types) throws QoreJavaCompilerException, ClassCastException {
        return compile(qualifiedClassName, javaSource, null, types);
    }

    /**
     * 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>
     * is 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 the Class and file data generated by compiling the source
     *
     * @throws QoreJavaCompilerException 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 CompilerOutput<T> compile(final String qualifiedClassName,
            final CharSequence javaSource,
            final Class<?>... types) throws QoreJavaCompilerException, ClassCastException {
        return compile(qualifiedClassName, javaSource, null, types);
    }

    /**
     * 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 diagnosticsList    Any diagnostics generated by compiling the source are added to
     *                           this collector.
     * @param types              zero or more Class objects representing classes or interfaces
     *                           that the resulting class must be assignable (castable) to.
     *
     * @return the Class and file data generated by compiling the source
     *
     * @throws QoreJavaCompilerException 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 CompilerOutput<T> compile(final String qualifiedClassName,
            final String javaSource,
            final DiagnosticCollector<JavaFileObject> diagnosticsList,
            final Class<?>... types) throws QoreJavaCompilerException, ClassCastException {
        return compile(qualifiedClassName, (CharSequence)javaSource, diagnosticsList, types);
    }

    /**
     * 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 diagnosticsList    Any diagnostics generated by compiling the source are added to
     *                           this collector.
     * @param types              zero or more Class objects representing classes or interfaces
     *                           that the resulting class must be assignable (castable) to.
     *
     * @return the Class and file data generated by compiling the source
     *
     * @throws QoreJavaCompilerException 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 CompilerOutput<T> compile(final String qualifiedClassName,
            final CharSequence javaSource,
            final DiagnosticCollector<JavaFileObject> diagnosticsList,
            final Class<?>... types) throws QoreJavaCompilerException, ClassCastException {
        diagnostics = diagnosticsList;
        Map<String, CharSequence> classes = new HashMap<String, CharSequence>(1);
        classes.put(qualifiedClassName, javaSource);
        Map<String, CompilerOutput<T>> compiled = compile(classes, diagnosticsList);
        CompilerOutput<T> newClassOutput = compiled.get(qualifiedClassName);
        if (newClassOutput == null) {
            throw new QoreJavaCompilerException(
                String.format("Compilation succeeded, but requested class '%s' is " +
                    "not present; byte code for the following classes was produced: %s", qualifiedClassName,
                    compiled.keySet()),
                classes.keySet(), diagnostics
            );
        }
        return castable(newClassOutput, types);
    }

    /**
     * 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.)
     * @param diagnosticsList Any diagnostics generated by compiling the source are added to
     *                        this list.
     *
     * @return A mapping of qualified class names to their corresponding Class and file data.
     * The map has keys for all output classes produced; not just the classes prsent in the input <var>classes</var>;
     * values are the corresponding Class and file data objects.
     *
     * @throws QoreJavaCompilerException if the source cannot be compiled
     */
    public synchronized Map<String, CompilerOutput<T>> compile(
            final Map<String, CharSequence> classes,
            final DiagnosticCollector<JavaFileObject> diagnosticsList)
            throws QoreJavaCompilerException {
        List<JavaFileObject> sources = new ArrayList<JavaFileObject>();
        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);
                //System.out.printf("compile pkg: '%s' source: '%s'\n", packageName, className + JAVA_EXTENSION);
                // 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);
            }
        }
        try {
            // 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 QoreJavaCompilerException("Compilation failed.", classes.keySet(), diagnostics);
            }
            try {
                // For each class name in the output map, get compiled class and file data and put it in the output map
                ArrayList<JavaFileObject> outputObjects = javaFileManager.getOutputFiles();
                Map<String, CompilerOutput<T>> compiled = new HashMap<String, CompilerOutput<T>>();
                for (JavaFileObject f : outputObjects) {
                    compiled.put(f.getName(), new CompilerOutput<T>(
                        //((JavaFileObjectImpl)f).getByteCode(),
                        loadClass(f.getName()),
                        f
                    ));
                }
                return compiled;
            } catch (ClassNotFoundException e) {
                throw new QoreJavaCompilerException(classes.keySet(), e, diagnostics);
            } catch (IllegalArgumentException e) {
                throw new QoreJavaCompilerException(classes.keySet(), e, diagnostics);
            } catch (SecurityException e) {
                throw new QoreJavaCompilerException(classes.keySet(), e, diagnostics);
            }
        } finally {
            classLoader.clearCompilationCache();
        }
    }

    /**
     * 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, DiagnosticCollector, 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);
    }

    /**
     * Add a path to the classpath
     */
    public void addClassPath(String path) {
        //System.out.printf("compiler.addClassPath() %s\n", path);
        classLoader.addPath(path);
    }

    public void injectClass(String binName, byte[] byteCode) {
        //System.out.printf("inject class %s: %d bytes\n", binName, byteCode.length);
        classLoader.addPendingClass(binName, byteCode);
    }

    /**
     * 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 CompilerOutput<T> castable(CompilerOutput<T> newClassOutput, Class<?>... types)
            throws ClassCastException {
        for (Class<?> type : types)
            if (!type.isAssignableFrom(newClassOutput.cls)) {
                throw new ClassCastException(type.getName());
            }
        return newClassOutput;
    }

    /**
     * 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);
        }
    }

    /**
     * @return This compiler's class loader.
     */
    public ClassLoader getClassLoader() {
        return javaFileManager.getClassLoader();
    }
}

/**
 * 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 implements QoreJavaFileObject {
    // 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(QoreJavaCompiler.toURI(baseName + ".java"), 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(QoreJavaCompiler.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
     */
    @Override
    public byte[] getByteCode() {
        return byteCode.toByteArray();
    }
}

/**
 * A JavaFileManager which manages Java source and classes. This FileManager
 * delegates to the JavaFileManager and the QoreURLClassLoader 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 QoreURLClassLoader classLoader;

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

    // List of output files generated; classes generated from input source
    private final ArrayList<JavaFileObject> outputObjects = new ArrayList<JavaFileObject>();

    /**
     * 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, QoreURLClassLoader classLoader) {
        super(fileManager);
        this.classLoader = classLoader;
    }

    /**
     * Clears the output file list
     */
    public void clearOutputFiles() {
        outputObjects.clear();
    }

    /**
     * @return a list of output files
     */
    public ArrayList<JavaFileObject> getOutputFiles() {
        return outputObjects;
    }

    /**
     * @return the class loader which this file manager delegates to
     */
    public ClassLoader getClassLoader() {
        return 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 QoreJavaCompiler.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 {
        JavaFileObjectImpl file = new JavaFileObjectImpl(qualifiedName, kind);
        classLoader.add(qualifiedName, file);
        // add to output list
        outputObjects.add(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<JavaFileObject>();
        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);
            }
            classLoader.files().forEach((f) -> {
                files.add((JavaFileObject)f);
            });
        } 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;
    }
}
