import org.apache.zookeeper.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.List;

public class Autohealer implements Watcher {

    private static final Logger logger = LoggerFactory.getLogger(Autohealer.class);

    private static final String ZOOKEEPER_ADDRESS = "127.0.0.1:2181";
    private static final int SESSION_TIMEOUT = 3000;
    private static final String AUTOHEALER_ZNODES_PATH = "/workers";

    private final String pathToWorkerJar;
    private final int numberOfWorkers;
    private ZooKeeper zooKeeper;

    public Autohealer(int numberOfWorkers, String pathToWorkerJar) {
        this.numberOfWorkers = numberOfWorkers;
        this.pathToWorkerJar = pathToWorkerJar;
    }

    public void connectToZookeeper() throws IOException {
        this.zooKeeper = new ZooKeeper(ZOOKEEPER_ADDRESS, SESSION_TIMEOUT, this);
        logger.info("Connected to ZooKeeper at {}", ZOOKEEPER_ADDRESS);
    }

    public void startWatchingWorkers() throws KeeperException, InterruptedException {
        if (zooKeeper.exists(AUTOHEALER_ZNODES_PATH, false) == null) {
            zooKeeper.create(AUTOHEALER_ZNODES_PATH, new byte[]{},
                    ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            logger.info("Created parent znode: {}", AUTOHEALER_ZNODES_PATH);
        }

        checkAndLaunchWorkers();
        zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, this);
    }

    private void checkAndLaunchWorkers() throws KeeperException, InterruptedException {
        List<String> children = zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, false);
        int runningWorkers = children.size();
        int toLaunch = numberOfWorkers - runningWorkers;

        if (toLaunch > 0) {
            logger.info("Launching {} new worker(s)", toLaunch);
            for (int i = 0; i < toLaunch; i++) {
                try {
                    startNewWorker();
                } catch (IOException e) {
                    logger.error("Failed to launch worker", e);
                }
            }
        } else {
            logger.debug("All workers are running");
        }
    }

    private void startNewWorker() throws IOException {
        File file = new File(pathToWorkerJar);
        String command = "java -jar " + file.getAbsolutePath();
        logger.info("Launching worker: {}", command);
        Runtime.getRuntime().exec(command, null, file.getParentFile());
    }

    @Override
    public void process(WatchedEvent event) {
        if (event.getType() == Event.EventType.NodeChildrenChanged &&
                AUTOHEALER_ZNODES_PATH.equals(event.getPath())) {
            try {
                logger.info("Detected worker znode change, checking workers...");
                checkAndLaunchWorkers();
            } catch (Exception e) {
                logger.error("Error while handling worker changes", e);
            }
        }

        try {
            zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, this);
        } catch (Exception e) {
            logger.error("Failed to reset watch", e);
        }
    }

    public void run() throws InterruptedException {
        synchronized (zooKeeper) {
            zooKeeper.wait();
        }
    }
}
