package Zookeeper;
/*
 *  MIT License
 *
 *  Copyright (c) 2019 Michael Pogrebinsky - Distributed Systems & Cloud Computing with Java
 *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy
 *  of this software and associated documentation files (the "Software"), to deal
 *  in the Software without restriction, including without limitation the rights
 *  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 *  copies of the Software, and to permit persons to whom the Software is
 *  furnished to do so, subject to the following conditions:
 *
 *  The above copyright notice and this permission notice shall be included in all
 *  copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 *  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 *  SOFTWARE.
 */
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.zookeeper.*;
import java.io.IOException;
import java.util.*;
import static org.apache.commons.io.IOUtils.close;
public class Autohealer implements Watcher {
    private static final Logger logger = LoggerFactory.getLogger(Autohealer.class);
    private static final String AUTOHEALER_ZNODES_PATH = "/workers";
    private ServiceRegistry serviceRegistry;
    private final String pathToProgram;
    private final int numberOfWorkers;
    private ZooKeeper zooKeeper;
    private Map<String, Integer> nodesCountMap = new HashMap<>();
    public Autohealer(int numberOfWorkers, String pathToProgram, ServiceRegistry serviceRegistry, ZooKeeper zooKeeper) {
        this.numberOfWorkers = numberOfWorkers;
        this.pathToProgram = pathToProgram;
        this.serviceRegistry = serviceRegistry;
        this.zooKeeper = zooKeeper;
    }
    public void startWatchingWorkers() throws KeeperException, InterruptedException {
        if (zooKeeper.exists(AUTOHEALER_ZNODES_PATH, this) == null) {
            zooKeeper.create(AUTOHEALER_ZNODES_PATH, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        launchWorkersIfNecessary();
    }
    public  void  run() throws InterruptedException {
        synchronized (this) {
            try {
                wait();
            } catch (InterruptedException e) {
                System.out.println(e.getMessage());
            }
        }
    }
    public synchronized void closeZookeeper() throws InterruptedException {
        synchronized (this) {
            try {
                close();
            }  catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }
    @Override
    public void process(WatchedEvent event) {
        switch (event.getType()) {
            case None:
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    System.out.println("Successfully connected to Zookeeper");
                    logger.info("Successfully connected to Zookeeper");

                } else {
                    System.out.println("Disconnected from Zookeeper event");
                    logger.warn("Disconnected from Zookeeper event");

                    zooKeeper.notifyAll();
                }
                break;
            case NodeChildrenChanged:
                System.out.println("change");
                if(event.getPath().equals(AUTOHEALER_ZNODES_PATH)){
                    try {
                        launchWorkersIfNecessary();
                    } catch (KeeperException | InterruptedException e) {
                        System.out.println("Failed to launch workers: " + e.getMessage());
                        logger.error("Failed to launch workers: " + e.getMessage());

                    }
                }
                break;
        }
    }
    private synchronized void launchWorkersIfNecessary() throws KeeperException, InterruptedException {
        List<String> workerZnodes = zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, this);
        int existingWorkers = workerZnodes.size();
        nodesCountMap.clear();
        for (String child : workerZnodes) {
            String znodePath = AUTOHEALER_ZNODES_PATH + "/" + child;
            byte[] data = zooKeeper.getData(znodePath, false, null);
            String workerAddress = new String(data);
            nodesCountMap.put(workerAddress, nodesCountMap.getOrDefault(workerAddress, 0) + 1);
        }
        List<String> allServiceAddresses = serviceRegistry.getAllServiceAddresses();

        for (String address : allServiceAddresses) {
            if (!nodesCountMap.containsKey(address)) {
                nodesCountMap.put(address, 0);
            }
        }
        System.out.println("existingWorkers: " + existingWorkers);
        logger.info("Existing workers: {}", existingWorkers);

        int numNodes = allServiceAddresses.size();

        int averageWorkersPerNode = 0;
        int remainingWorkers = 0;
        if(numNodes != 0){
            averageWorkersPerNode= numberOfWorkers / numNodes;
            remainingWorkers = numberOfWorkers % numNodes;
        }

        Map<String, Integer> workersToAdd = new HashMap<>();
        for (String address : allServiceAddresses) {
            int currentWorkers = nodesCountMap.getOrDefault(address, 0);
            int workersNeeded = Math.max(0, averageWorkersPerNode - currentWorkers);
            workersToAdd.put(address, workersNeeded);
        }
        List<String> sortedNodes = new ArrayList<>(allServiceAddresses);
        sortedNodes.sort(Comparator.comparingInt(node -> nodesCountMap.getOrDefault(node, 0)));
        int nodeIndex = 0;
        for (int i = 0; i < remainingWorkers; i++) {
            String nodeAddress = sortedNodes.get(nodeIndex);
            workersToAdd.put(nodeAddress, workersToAdd.get(nodeAddress) + 1);
            nodeIndex = (nodeIndex + 1) % numNodes;
        }

        int workersToLaunch = 0;
        if (existingWorkers < numberOfWorkers) {
            workersToLaunch = numberOfWorkers - existingWorkers;
        }
        System.out.println("workersToLaunch: " + workersToLaunch);
        logger.info("Workers to launch: {}", workersToLaunch);
        System.out.println("Current Nodes Status On Cluster: ");
        for (Map.Entry<String, Integer> entry : nodesCountMap.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }
        logger.info("Current Nodes Status On Cluster:");
        for (Map.Entry<String, Integer> entry : nodesCountMap.entrySet()) {
            logger.info("{}: {}", entry.getKey(), entry.getValue());
        }
        for (String address : sortedNodes) {
            int myWorkersToLaunch = workersToAdd.get(address);
            for (int i = 0; i < myWorkersToLaunch && workersToLaunch> 0; i++) {
                try {
                    startNewWorker(address, pathToProgram);
                    workersToLaunch--;
                    nodesCountMap.put(address, nodesCountMap.getOrDefault(address, 0) + 1);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }
        System.out.println("Nodes Status After Launching Process: ");
        for (Map.Entry<String, Integer> entry : nodesCountMap.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }
        logger.info("Nodes Status After Launching Process:");
        for (Map.Entry<String, Integer> entry : nodesCountMap.entrySet()) {
            logger.info("{}: {}", entry.getKey(), entry.getValue());
        }
        Thread.sleep(30000);
    }
    /**
     * Helper method to start a single worker
     * @throws IOException
     */
    private void startNewWorker(String workerAddress, String workerPath) throws IOException, InterruptedException {
        String[] addressParts = workerAddress.split("@");
        String username = addressParts[0];
        String ipAddress = addressParts[1];
        String sshCommand = "ssh " + username + "@" + ipAddress + " \"java -jar /home/" + username + "/registration-and-discovery-zookeeper/out/artifacts/Service_Registration_and_Discovery_jar2/Service-Registration-and-Discovery.jar " + username + "@" + ipAddress + "\"";
        System.out.println(String.format("Launching worker instance on: %s ", username + "@" + ipAddress));
        logger.info("Launching worker instance on: {}", username + "@" + ipAddress);
        Runtime.getRuntime().exec(sshCommand);
    }
}
