import org.apache.zookeeper.*;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Collections;
import java.util.List;

public class Autohealer implements Watcher {

    private static final String ZOOKEEPER_ADDRESS = "192.168.144.134:2181";
    private static final int SESSION_TIMEOUT = 3000;

    // مسارات Zookeeper
    private static final String WORKERS_PATH = "/workers";
    private static final String NODES_PATH = "/nodes";
    private static final String LOG_FILE = "cluster_events.log";

    private final int numberOfWorkers;
    private ZooKeeper zooKeeper;
    private int nodeIndex = 0; //  (Round Robin)

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

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

    public void startWatching() throws KeeperException, InterruptedException {
        ensurePathExists(WORKERS_PATH);
        ensurePathExists(NODES_PATH);
        logEvent("Autohealer started. Monitoring cluster for " + numberOfWorkers + " workers.");
        checkHealthAndRebalance();
    }

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

    @Override
    public void process(WatchedEvent event) {
        switch (event.getType()) {
            case None:
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    System.out.println("Successfully connected to Zookeeper");
                }
                break;
            case NodeChildrenChanged:
                // اكتشاف تغيير في العمال أو العقد الفيزيائية
                logEvent("Change detected in cluster (Nodes or Workers changed).");
                checkHealthAndRebalance();
                break;
        }
    }

    private void checkHealthAndRebalance() {
        try {

            List<String> currentWorkers = zooKeeper.getChildren(WORKERS_PATH, this);
            List<String> availableNodes = zooKeeper.getChildren(NODES_PATH, this);

            int currentCount = currentWorkers.size();
            int nodesCount = availableNodes.size();

            System.out.println("Status: Workers = " + currentCount + ", Nodes = " + nodesCount);

            if (currentCount < numberOfWorkers) {
                if (nodesCount == 0) {
                    logEvent("CRITICAL: All physical nodes are down! Cannot heal cluster.");
                    return;
                }

                logEvent("Healing required: " + currentCount + "/" + numberOfWorkers + " workers active.");
                int workersToLaunch = numberOfWorkers - currentCount;
                distributeWorkers(workersToLaunch, availableNodes);
            }
        } catch (KeeperException | InterruptedException e) {
            logEvent("Error during health check: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private void distributeWorkers(int count, List<String> nodes) throws KeeperException, InterruptedException {
        Collections.sort(nodes);

        for (int i = 0; i < count; i++) {

            String targetNode = nodes.get(nodeIndex % nodes.size());
            nodeIndex++;

            sendLaunchCommand(targetNode);
        }
    }

    private void sendLaunchCommand(String nodeName) throws KeeperException, InterruptedException {
        String taskPath = NODES_PATH + "/" + nodeName + "/tasks";
        ensurePathExists(taskPath);

        String taskNode = zooKeeper.create(taskPath + "/task_",
                new byte[]{},
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL_SEQUENTIAL);

        logEvent("Scheduled new worker on node: " + nodeName + " (Task ID: " + taskNode.substring(taskNode.lastIndexOf("/") + 1) + ")");
    }


    private void logEvent(String message) {
        String timestamp = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
        String logMessage = "[" + timestamp + "] " + message;

        System.out.println(logMessage);

        try (FileWriter fw = new FileWriter(LOG_FILE, true);
             PrintWriter out = new PrintWriter(fw)) {
            out.println(logMessage);
        } catch (IOException e) {
            System.err.println("Could not write to log file: " + e.getMessage());
        }
    }

    public void run() throws InterruptedException {
        synchronized (zooKeeper) {
            while (zooKeeper.getState().isAlive()) {
                zooKeeper.wait(5000);
            }
        }
    }

    public void close() throws InterruptedException {
        zooKeeper.close();
    }
}