import org.apache.zookeeper.*;
import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.util.List;

public class NodeAgent implements Watcher {
    private static final String ZOOKEEPER_ADDRESS = "192.168.144.134:2181";
    private static final int SESSION_TIMEOUT = 3000;
    private static final String NODES_ROOT = "/nodes";
    private static final String WORKER_JAR_PATH = "C:\\Users\\J.N\\IdeaProjects\\auto-healer-zookeper\\transientworker\\target\\transientworker-1.0-SNAPSHOT.jar";

    private ZooKeeper zooKeeper;
    private String myTasksPath;

    public void connectToZookeeper() throws IOException {
        this.zooKeeper = new ZooKeeper(ZOOKEEPER_ADDRESS, SESSION_TIMEOUT, this);
    }

    public void registerNodeAndWatchTasks() throws KeeperException, InterruptedException, IOException {
        ensurePath(NODES_ROOT, CreateMode.PERSISTENT);

        String ipAddress = InetAddress.getLocalHost().getHostAddress();
        String nodePath = NODES_ROOT + "/node_" + ipAddress;


        ensurePath(nodePath, CreateMode.PERSISTENT);


        String activePath = nodePath + "/active";
        if (zooKeeper.exists(activePath, false) == null) {
            zooKeeper.create(activePath, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
        }


        this.myTasksPath = nodePath + "/tasks";
        ensurePath(myTasksPath, CreateMode.PERSISTENT);

        System.out.println("Node Registered: " + nodePath + ". Watching tasks at: " + myTasksPath);
        watchMyTasks();
    }


    private void ensurePath(String path, CreateMode mode) throws KeeperException, InterruptedException {
        if (zooKeeper.exists(path, false) == null) {
            zooKeeper.create(path, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, mode);
        }
    }

    private void watchMyTasks() throws KeeperException, InterruptedException {
        List<String> tasks = zooKeeper.getChildren(myTasksPath, this);
        System.out.println("Current tasks count: " + tasks.size());
    }

    @Override
    public void process(WatchedEvent event) {
        if (event.getType() == Event.EventType.NodeChildrenChanged && event.getPath().equals(myTasksPath)) {
            handleNewTask();
        } else if (event.getState() == Event.KeeperState.SyncConnected) {
            System.out.println("Connected to Zookeeper");
        }
    }

    private void handleNewTask() {
        try {
            startLocalWorker();
            watchMyTasks(); // إعادة المراقبة
        } catch (Exception e) { e.printStackTrace(); }
    }

    private void startLocalWorker() throws IOException {
        File jarFile = new File(WORKER_JAR_PATH);
        ProcessBuilder pb = new ProcessBuilder("java", "-jar", jarFile.getAbsolutePath());
        pb.inheritIO();
        pb.directory(jarFile.getParentFile());
        System.out.println("New worker launched locally!");
        pb.start();
    }

    public static void main(String[] args) throws Exception {
        NodeAgent agent = new NodeAgent();
        agent.connectToZookeeper();
        agent.registerNodeAndWatchTasks();
        synchronized (agent.zooKeeper) { agent.zooKeeper.wait(); }
    }
}