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

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

public class Autohealer implements Watcher {
    private static final Logger logger = LoggerFactory.getLogger(Autohealer.class);

    // Update to include all ensemble nodes
    private static final String ZOOKEEPER_ADDRESS = "127.0.0.1:2181,127.0.0.1:2182,127.0.0.1:2183";
    private static final int SESSION_TIMEOUT = 3000;
    private static final String WORKERS_PATH = "/workers";

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

    // Simulate multiple physical nodes
    private final List<String> physicalNodes = Arrays.asList("node1", "node2", "node3");

    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("Connecting to ZooKeeper ensemble: {}", ZOOKEEPER_ADDRESS);
    }

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

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

    public void run() throws InterruptedException {
        synchronized (zooKeeper) {
            zooKeeper.wait(); // block main thread
        }
    }

    public void close() throws InterruptedException {
        zooKeeper.close();
        logger.info("ZooKeeper connection closed");
    }

    @Override
    public void process(WatchedEvent event) {
        if (event.getType() == Event.EventType.None) {
            if (event.getState() == Event.KeeperState.SyncConnected) {
                logger.info("Successfully connected to ZooKeeper");
            } else {
                synchronized (zooKeeper) {
                    logger.warn("Disconnected from ZooKeeper");
                    zooKeeper.notifyAll();
                }
            }
            return;
        }

        if (event.getType() == Event.EventType.NodeChildrenChanged && event.getPath().equals(WORKERS_PATH)) {
            try {
                logger.info("Workers changed, checking cluster health...");
                checkAndLaunchWorkers();
            } catch (Exception e) {
                logger.error("Error checking/launching workers", e);
            }
        }

        // Re-set the watch
        try {
            zooKeeper.getChildren(WORKERS_PATH, this);
        } catch (Exception e) {
            logger.error("Failed to reset watch", e);
        }
    }

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

        if (toLaunch > 0) {
            logger.info("Need to launch {} new worker(s)", toLaunch);
            for (int i = 0; i < toLaunch; i++) {
                try {
                    String node = selectNodeForWorker(children);
                    startNewWorker(node);
                } catch (IOException e) {
                    logger.error("Failed to start new worker", e);
                }
            }
        } else {
            logger.debug("All workers are running, no action needed");
        }
    }

    // Simple round-robin assignment of worker to a physical node
    private String selectNodeForWorker(List<String> currentWorkers) {
        int idx = currentWorkers.size() % physicalNodes.size();
        return physicalNodes.get(idx);
    }

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