package com.thinkaurelius.titan;

import com.google.common.base.Joiner;
import com.thinkaurelius.titan.diskstorage.configuration.ModifiableConfiguration;
import com.thinkaurelius.titan.diskstorage.configuration.WriteConfiguration;
import com.thinkaurelius.titan.graphdb.configuration.GraphDatabaseConfiguration;

import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.RandomAccessFile;

public class HBaseStorageSetup {

    private static Process HBASE = null;
    // amount of seconds to wait before assuming that HBase shutdown
    private static final int SHUTDOWN_TIMEOUT_SEC = 20;

    // hbase config for testing
    private static final String HBASE_CONFIG_DIR = "./conf";

    // default pid file location
    private static final String HBASE_PID_FILE = "/tmp/hbase-" + System.getProperty("user.name") + "-master.pid";

    private static final Logger log = LoggerFactory.getLogger(HBaseStorageSetup.class);

    static {
        try {
            log.info("Deleting old test directories (if any).");

            // please keep in sync with HBASE_CONFIG_DIR/hbase-site.xml, reading HBase XML config is huge pain.
            File hbaseRoot = new File("./target/hbase-root");
            File zookeeperDataDir = new File("./target/zk-data");

            if (hbaseRoot.exists())
                FileUtils.deleteDirectory(hbaseRoot);

            if (zookeeperDataDir.exists())
                FileUtils.deleteDirectory(zookeeperDataDir);
        } catch (IOException e) {
            log.error("Failed to delete old HBase test directories: '" + e.getMessage() + "', ignoring.");
        }

        Runtime.getRuntime().addShutdownHook(new Thread() {
            public void run() {
                log.info("All done. Shutting done HBase.");

                try {
                    HBaseStorageSetup.shutdownHBase();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }

    public static ModifiableConfiguration getHBaseConfiguration() {
        ModifiableConfiguration config = GraphDatabaseConfiguration.buildConfiguration();
        config.set(GraphDatabaseConfiguration.STORAGE_BACKEND, "hbase");
        return config;
    }

    public static WriteConfiguration getHBaseGraphConfiguration() {
        return getHBaseConfiguration().getConfiguration();
    }

    public static void startHBase() throws IOException {
        if (HBASE != null) {
            log.info("HBase already started");
            return; // already started, nothing to do
        }

        try {
            log.info("Starting HBase");

            // start HBase instance with environment set
            String cmd = String.format("./bin/hbase-daemon.sh --config %s start master", HBASE_CONFIG_DIR);
            log.info("Executing {}", cmd);
            HBASE = Runtime.getRuntime().exec(cmd);

            try {
                HBASE.waitFor(); // wait for script to return
                log.info("HBase forked");
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }

            assert HBASE.exitValue() >= 0; // check if we have started successfully
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static void shutdownHBase() throws IOException {
        if (HBASE == null)
            return; // HBase hasn't been started yet

        // First try graceful shutdown through the script...
        String cmdParts[] = new String[]{ "./bin/hbase-daemon.sh", "--config", HBASE_CONFIG_DIR, "stop", "master" };
        log.info("Executing {}", Joiner.on(" ").join(cmdParts));
        ProcessBuilder pb = new ProcessBuilder(cmdParts);
        pb.redirectErrorStream(true);
        Process stopMaster = pb.start();
        StreamLogger sl = new StreamLogger(stopMaster.getInputStream());
        sl.setDaemon(true);
        sl.start();

        // ...but send SIGKILL if that times out
        HBaseKiller killer = new HBaseKiller();
        killer.start();
        try {
            stopMaster.waitFor();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        killer.abort();
        killer.interrupt();
        try {
            killer.join();
        } catch (InterruptedException e) {
            log.warn("Failed to join HBase process killer thread", e);
        }

        // StreamLogger is a daemon thread, so failing to stop it isn't so bad
        try {
            sl.join(1000L);
        } catch (InterruptedException e) {
            log.warn("Failed to stop HBase output logger thread", e);
        }
    }

    private static class HBaseKiller extends Thread {

        private volatile boolean proceed = true;

        @Override
        public void run() {
            try {
                log.info("Waiting {} seconds for HBase to shutdown...", SHUTDOWN_TIMEOUT_SEC);
                Thread.sleep(SHUTDOWN_TIMEOUT_SEC * 1000L);
            } catch (InterruptedException e) {
                log.info("HBase killer thread interrupted");
            }

            if (proceed) {
                try {
                    readPidfileAndKill();
                } catch (Exception e) {
                    log.error("Failed to kill HBase", e);
                }
            } else {
                log.info("HBase shutdown cleanly, not attempting to kill its process");
            }
        }

        public void abort() {
            proceed = false;
        }

        private void readPidfileAndKill() throws IOException, InterruptedException {
            File pid = new File(HBASE_PID_FILE);

            if (pid.exists()) {
                RandomAccessFile pidFile = new RandomAccessFile(pid, "r");

                StringBuilder b = new StringBuilder();

                while (pidFile.getFilePointer() < (pidFile.length() - 1)) // we don't need newline
                   b.append((char) pidFile.readByte());

                String cmd = "kill -KILL " + b.toString();
                log.info("Executing {}", cmd);
                Process kill = Runtime.getRuntime().exec(cmd);
                kill.waitFor();

                pidFile.close();
                pid.delete(); // delete pid file like nothing happened

                return;
            }
        }
    }

    /*
     * This could be retired in favor of ProcessBuilder.Redirect when we move to
     * source level 1.7.
     */
    private static class StreamLogger extends Thread {

        private final BufferedReader reader;
        private static final Logger log =
                LoggerFactory.getLogger(StreamLogger.class);

        private StreamLogger(InputStream is) {
            this.reader = new BufferedReader(new InputStreamReader(is));
        }

        @Override
        public void run() {
            String line;
            try {
                while (null != (line = reader.readLine())) {
                    log.info("> {}", line);
                    if (Thread.currentThread().isInterrupted()) {
                        break;
                    }
                }

                log.info("End of stream.");
            } catch (IOException e) {
                log.error("Unexpected IOException while reading stream {}", reader, e);
            }
        }
    }
}

@@ -304,16 +304,277 @@
 uration;
+%0Aimport com.thinkaurelius.titan.graphdb.database.idassigner.IDPartitionMode;%0Aimport com.thinkaurelius.titan.graphdb.database.idassigner.placement.PartitionIDRange;%0Aimport com.thinkaurelius.titan.graphdb.database.idassigner.placement.SimpleBulkPlacementStrategy;
 %0A%0Aimport
@@ -2693,24 +2693,185 @@
 , %22hbase%22);%0A
+        config.set(GraphDatabaseConfiguration.IDS_PARTITION, IDPartitionMode.ENABLED);%0A        config.set(SimpleBulkPlacementStrategy.CONCURRENT_PARTITIONS, 1);%0A
         retu

Fix bogus HBaseGraphMetrics test failure