/*
 *  Copyright 2015 Michael Pymm
 *
 *  Licensed 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 com.twosigma.beaker.kdb;

import com.twosigma.beaker.core.rest.StreamGobbler;
import org.apache.commons.lang3.SystemUtils;

import java.io.File;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Wraps a kdb process.
 */
public final class KdbProcess extends Thread {
    // Environment variable names.
    private static final String QHOME                = "QHOME";
    private static final String QLIC                 = "QLIC";
    private static final String BEAKER_CORE_PASSWORD = "beaker_core_password";
    private static final String BEAKER_CORE_PORT     = "beaker_core_port";
    private static final String SESSION_ID           = "session_id";

    // kdb commands.
    public static final String KILL_CMD = "\\\\\n";

    // Kdb configuration.
    private final String qhome = System.getenv(QHOME);
    private final String qbin;

    // Session id.
    private final String sessionId;

    // The port to have kdb listen on.
    private final int kdbPort;

    /**
     * Create a new KdbProcess.
     *
     * @param sessionId  the session id (for namespace access).
     * @param kdbPort    the port for kdb to listen on.
     */
    public KdbProcess(String sessionId, int kdbPort) throws Exception {
        super("kdb-" + sessionId + ":" + kdbPort);
        this.sessionId = sessionId;
        this.kdbPort = kdbPort;

        // Try to find the q binary.
        if (qhome == null) {
            throw new Exception("QHOME is not set");
        }

        {
            // Get OS-specific candidates.
            List<String> binaries = new ArrayList<>();
            if (SystemUtils.IS_OS_WINDOWS) {
                binaries.add("q.exe");
            } else if (SystemUtils.IS_OS_MAC_OSX) {
                binaries.add("m64/q");
                binaries.add("m32/q");
            } else if (SystemUtils.IS_OS_LINUX) {
                binaries.add("l64/q");
                binaries.add("l32/q");
            } else {
                throw new Exception("Unsupported operating system");
            }

            String bin = null;
            for (String s : binaries) {
                String f = qhome + File.separator + s;
                if (new File(f).canExecute()) {
                    bin = f;
                    break;
                }
            }

            if (bin == null) {
                throw new Exception("Cannot find q binary");
            } else {
                qbin = bin;
            }
        }

        // Beaker calls destroy() on us, so add a shutdown hook to
        // cleanly shut down kdb at the end.
        Runtime.getRuntime().addShutdownHook(new Thread("kdb-killer") {
            @Override
            public void run() {
                KdbProcess.this.interrupt();
            }
        });
    }

    @Override
    public void run() {
        try {
            runImpl();
        } catch (Exception e) {
            // TODO
            e.printStackTrace();
        }
    }

    private void runImpl() throws Exception {
        // Guess at QLIC if it's not set.
        String qlic = System.getenv(QLIC);
        if (qlic == null) {
            qlic = qhome + File.separator + "k4.lic";
        }

        // Start kdb.
        Process kdbProcess = Runtime.getRuntime().exec(new String[] {
            qbin,
            "-p",
            Integer.toString(kdbPort)
        }, new String[] {
            QHOME + "=" + qhome,
            QLIC  + "=" + qlic,
            BEAKER_CORE_PASSWORD + "=" + System.getenv(BEAKER_CORE_PASSWORD),
            BEAKER_CORE_PORT + "=" + System.getenv(BEAKER_CORE_PORT),
            SESSION_ID + "=" + sessionId
        });

        // Wrap stdin/stdout.
        OutputStreamWriter stdin = new OutputStreamWriter(kdbProcess.getOutputStream());
        StreamGobbler      stdout = new StreamGobbler(kdbProcess.getInputStream(), "stdout");
        stdout.start();

        long exitTime = Long.MAX_VALUE;
        while (kdbProcess.isAlive()) {
            try {
                // Wait for a second for kdb to exit.
                if (kdbProcess.waitFor(1, TimeUnit.SECONDS)) {
                    break;
                }

                // If we've passed the exit time, forcibly kill the process.
                if (System.currentTimeMillis() > exitTime) {
                    kdbProcess.destroyForcibly();
                    break;
                }
            } catch (InterruptedException e) {
                // Interrupted - try and kill kdb.
                stdin.write(KILL_CMD);
                stdin.flush();

                // Forcibly destroy the process in two seconds if necessary.
                exitTime = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(2);
            }
        }
    }
}

@@ -759,43 +759,8 @@
 le;%0A
-import java.io.OutputStreamWriter;%0A
 impo
@@ -809,46 +809,8 @@
 ist;
-%0Aimport java.util.concurrent.TimeUnit;
 %0A%0A/*
@@ -1276,82 +1276,8 @@
 %22;%0A%0A
-    // kdb commands.%0A    public static final String KILL_CMD = %22%5C%5C%5C%5C%5Cn%22;%0A%0A
     
@@ -4189,22 +4189,16 @@
 rap 
-stdin/
 stdout.%0A
     
@@ -4197,97 +4197,8 @@
 ut.%0A
-        OutputStreamWriter stdin = new OutputStreamWriter(kdbProcess.getOutputStream());%0A
     
@@ -4215,21 +4215,16 @@
 Gobbler 
-     
 stdout =
@@ -4319,397 +4319,135 @@
     
-long exitTime = Long.MAX_VALUE;%0A        while (kdbProcess.isAlive()) %7B%0A            try %7B%0A                // Wait for a second for kdb to exit.%0A                if (kdbProcess.waitFor(1, TimeUnit.SECONDS)) %7B%0A                    break;%0A                %7D%0A%0A                // If we've passed the exit time, forcibly kill the process.%0A                if (System.currentTimeMillis() %3E exitTime) %7B
+// Wait until kdb exits or we're interrupted.%0A        while (true) %7B%0A            try %7B%0A                // Wait for kdb to exit.
 %0A   
@@ -4459,20 +4459,16 @@
         
-    
 kdbProce
@@ -4474,23 +4474,15 @@
 ess.
-destroyForcibly
+waitFor
 ();%0A
@@ -4501,37 +4501,15 @@
     
-  
-  
 break;%0A
-                %7D%0A
     
@@ -4588,257 +4588,64 @@
 d - 
-try and kill kdb.%0A                stdin.write(KILL_CMD);%0A                stdin.flush();%0A%0A                // Forcibly destroy the process in two seconds if necessary.%0A                exitTime = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(2
+destroy the process.%0A                kdbProcess.destroy(
 );%0A 

Use java 7 Process methods.