import java.lang.reflect.Method;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.graalvm.polyglot.Context;
import org.graalvm.polyglot.Value;


/**
 * Some code taken from: https://www.soulmachine.me/blog/2015/07/22/compile-and-run-java-source-code-in-memory/
 */
public class DynamicJava {

    static final String FILE_NAME = "Solution.java";    
    static final String SOURCE =
        "public final class Solution {\n" +
        "   public static String greet(String name) {\n" +
        "       return \"Hello \" + name;\n" +
        "   }\n" +
        "}\n";


    static Map<String, byte[]> compileInTheHost(String fileName, String source) {
        final InMemoryJavaCompiler compiler = new InMemoryJavaCompiler();
        final Map<String, byte[]> classes = compiler.compile(fileName, source);
        return classes;
    }

    static Value compileInTheGuest(Context context, String fileName, String source) {
        Value bindings = context.getBindings("java");
        Value InMemoryJavaCompiler_klass = bindings.getMember(InMemoryJavaCompiler.class.getName());
        Value compiler = InMemoryJavaCompiler_klass.newInstance();
        Value classes = compiler.invokeMember("compile", fileName, source);
        return classes;
    }

    static void testInTheHost(Map<String, byte[]> classes) {
        MemoryClassLoader loader = new MemoryClassLoader(classes);
        String result;
        try {            
            Class<?> Solution_class = loader.loadClass("Solution");
            Method greet = Solution_class.getDeclaredMethod("greet", String.class);
            result = (String) greet.invoke(null, "host");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }        
        System.out.println("(testInTheHost) result: " + result);
    }

    static Value toHost(Context context, Map<String, byte[]> classes) {        
        Value bindings = context.getBindings("java");
        Value ByteArray_klass = bindings.getMember(byte[].class.getName());
        Value HashMap_klass = bindings.getMember(HashMap.class.getName());

        Value guestClasses = HashMap_klass.newInstance();
        for (Map.Entry<String, byte[]> entry : classes.entrySet()) {
            String key = entry.getKey();
            byte[] value = entry.getValue();    
            Value guestValue = ByteArray_klass.newInstance(value.length);
            for (int i = 0; i < value.length; ++i) {
                guestValue.setArrayElement(i, value[i]);
            }
            // key is automatically converted into a guest String.
            guestClasses.invokeMember("put", key, guestValue);
        }

        return guestClasses;
    }    

    static void testInTheGuest(Context context, Value classes) {
        Value bindings = context.getBindings("java");
        Value MemoryClassLoader_klass = bindings.getMember(MemoryClassLoader.class.getName());
        Value loader = MemoryClassLoader_klass.newInstance(classes);
        Value Solution = loader.invokeMember("loadClass", "Solution");
        Value result = Solution.getMember("static").invokeMember("greet", "Espresso");
        System.out.println("(testInTheGuest) result: " + result.asString());
    }
    public static void main(String[] args) {
        try (Context context = Context.newBuilder("java")
            .allowAllAccess(true)
            // To expose MemoryClassLoader to the guest.
            .option("java.Classpath", System.getProperty("java.class.path"))
            .build()) {

            // Compile and run in the host.
            Map<String, byte[]> classes_host = compileInTheHost(FILE_NAME, SOURCE);
            testInTheHost(classes_host);

            // Compile and run in the guest.
            Value classes_guest = compileInTheGuest(context, FILE_NAME, SOURCE);
            testInTheGuest(context, classes_guest);

            // Compile in the host (already done) and run in the guest.
            testInTheGuest(context, toHost(context, classes_host));
        }
    }
}

class MemoryClassLoader extends URLClassLoader {
    private final Map<String, byte[]> classBytes = new ConcurrentHashMap<String, byte[]>();

    public MemoryClassLoader(Map<String, byte[]> classBytes) {
        super(new URL[0], MemoryClassLoader.class.getClassLoader());
        this.classBytes.putAll(classBytes);
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        byte[] buf = classBytes.get(name);
        if (buf == null) {
            return super.findClass(name);
        }
        classBytes.remove(name);
        return defineClass(name, buf, 0, buf.length);
    }
}