package com.zaxxer.jna;

import com.sun.jna.JNIEnv;
import com.sun.jna.Library;
import com.sun.jna.Native;
import com.sun.jna.NativeLibrary;

import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.TimeUnit;

public class ProcessTest {
    @SuppressWarnings("unused")
    private enum LaunchMechanism {
        // order IS important!
        FORK,
        POSIX_SPAWN,
        VFORK
    }

    public static class JavaNativeProcess {

        static {
            Map<String, Object> options = new HashMap<>();
            options.put(Library.OPTION_ALLOW_OBJECTS, Boolean.TRUE);

            Native.register(NativeLibrary.getProcess(options));
        }

        public native void Java_java_lang_UNIXProcess_init(JNIEnv jniEnv, Object clazz);

        /**
         * JNIEXPORT jint JNICALL
         * Java_java_lang_UNIXProcess_forkAndExec(JNIEnv *env,
         *                                        jobject process,
         *                                        jint mode,
         *                                        jbyteArray helperpath,
         *                                        jbyteArray prog,
         *                                        jbyteArray argBlock, jint argc,
         *                                        jbyteArray envBlock, jint envc,
         *                                        jbyteArray dir,
         *                                        jintArray std_fds,
         *                                        jboolean redirectErrorStream)
         *
         * @return the PID of the process
         */
        public native int Java_java_lang_UNIXProcess_forkAndExec(
                JNIEnv jniEnv,
                int mode,
                byte[] helperpath,
                byte[] prog,
                byte[] argBlock, int argc,
                byte[] envBlock, int envc,
                byte[] dir,
                int[] fds,
                boolean redirectErrorStream
        );
    }

    public static void main(String[] cargs) throws InterruptedException {

        JavaNativeProcess jnp = new JavaNativeProcess();

        jnp.Java_java_lang_UNIXProcess_init(JNIEnv.CURRENT, ProcessTest.class);

        LaunchMechanism launchMechanism = LaunchMechanism.POSIX_SPAWN;
        if (System.getProperty("os.name").contains("Linux")) {
            launchMechanism = LaunchMechanism.VFORK;
        }

        String[] cmdarray = {"/bin/cat", "/Users/brettw/pgadmin.log"};

        // See https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/java/lang/ProcessImpl.java#L71-L83
        byte[][] args = new byte[cmdarray.length-1][];
        int size = args.length; // For added NUL bytes
        for (int i = 0; i < args.length; i++) {
            args[i] = cmdarray[i+1].getBytes();
            size += args[i].length;
        }
        byte[] argBlock = new byte[size];
        int i = 0;
        for (byte[] arg : args) {
            System.arraycopy(arg, 0, argBlock, i, arg.length);
            i += arg.length + 1;
            // No need to write NUL bytes explicitly
        }

        // See https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/java/lang/ProcessImpl.java#L86
        TreeMap<String, String> environment = new TreeMap<>(System.getenv());
        byte[] envBlock = toEnvironmentBlock(environment);

        // See https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/java/lang/ProcessImpl.java#L96
        int[] std_fds = new int[] { -1, -1, -1 };

        // See https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/classes/java/lang/UNIXProcess.java#L247
        // Native source code: https://github.com/JetBrains/jdk8u_jdk/blob/master/src/solaris/native/java/lang/UNIXProcess_md.c#L566
        int pid = jnp.Java_java_lang_UNIXProcess_forkAndExec(
                JNIEnv.CURRENT,
                launchMechanism.ordinal() + 1,
                toCString(System.getProperty("java.home") + "/lib/jspawnhelper"),
                toCString(cmdarray[0]),
                argBlock, args.length,
                envBlock, environment.size(),
                null,
                std_fds,
                false
        );

        System.out.println("Child PID: " + pid);
        TimeUnit.SECONDS.sleep(3);
    }

    private static byte[] toCString(String s) {
        if (s == null)
            return null;
        byte[] bytes = s.getBytes();
        byte[] result = new byte[bytes.length + 1];
        System.arraycopy(bytes, 0,
                result, 0,
                bytes.length);
        result[result.length-1] = (byte)0;
        return result;
    }


    private static byte[] toEnvironmentBlock(TreeMap<String, String> environment) {
        int count = environment.size() * 2;
        for (Map.Entry<String, String> entry : environment.entrySet()) {
            count += entry.getKey().getBytes().length;
            count += entry.getValue().getBytes().length;
        }

        byte[] block = new byte[count];

        int i = 0;
        for (Map.Entry<String, String> entry : environment.entrySet()) {
            byte[] key   = entry.getKey  ().getBytes();
            byte[] value = entry.getValue().getBytes();
            System.arraycopy(key, 0, block, i, key.length);
            i+=key.length;
            block[i++] = (byte) '=';
            System.arraycopy(value, 0, block, i, value.length);
            i+=value.length + 1;
            // No need to write NUL byte explicitly
            //block[i++] = (byte) '\u0000';
        }

        return block;
    }
}