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 ch.qos.logback.classic.Logger;
import org.apache.zookeeper.*;
import org.slf4j.LoggerFactory;

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

public class Autohealer implements Watcher {

    private static final String REGISTRY_ZNODE = "/service_registry";

    // Parent Znode where each worker stores an ephemeral child to indicate it is alive
    private static final String AUTOHEALER_ZNODES_PATH = "/workers";

    // Path to the worker jar
    private final String pathToProgram;

    // The number of worker instances we need to maintain at all times
    private final int numberOfWorkers;
    private ZooKeeper zooKeeper;
    private static final Logger logger = (Logger) LoggerFactory.getLogger(Autohealer.class);


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

    public void startWatchingWorkers() throws KeeperException, InterruptedException {
        System.out.println("start watching");
        if (zooKeeper.exists(AUTOHEALER_ZNODES_PATH, false) == null) {
            zooKeeper.create(AUTOHEALER_ZNODES_PATH, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        launchWorkersIfNecessary();
    }

    public synchronized void launchWorkersIfNecessary() throws InterruptedException {
        Thread.sleep(15000);
        logger.info("start launch");
        try {
            // Check if the parent znode exists, create it if not
            if (zooKeeper.exists(AUTOHEALER_ZNODES_PATH, false) == null) {
                zooKeeper.create(AUTOHEALER_ZNODES_PATH, new byte[]{}, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            // Get the current number of workers
            List<String> Jobs = zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, this);
            for (String job : Jobs) {
                zooKeeper.exists(AUTOHEALER_ZNODES_PATH + "/" + job, this);
            }
            int currentJobs = Jobs.size();
            logger.info("Number of jobs is {}" , numberOfWorkers);
            if (ServiceRegistry.getWorkers() == null)
                return;
            for (Worker worker : ServiceRegistry.getWorkers()) {
                worker.setnumOfJobsZero();
            }
            for (String job : Jobs) {
                String[] parts = job.split("'");
                // Extract workerNode using Substring
                Worker worker = Worker.findByNodeName(parts[1], ServiceRegistry.getWorkers());
                worker.increaseOne();
            }
            Worker.SortNodes(ServiceRegistry.getWorkers());
            while (true) {
                List<Worker> workers = ServiceRegistry.getWorkers();
                for (Worker worker : workers) {
                    logger.info("Total number of job {}",Worker.getTotalNumOfJob(workers));
                    Thread.sleep(2000);
                    if (Worker.getTotalNumOfJob(workers) < numberOfWorkers) {
                        startJobOnNode(worker);
                        currentJobs++;
                    } else break;
                }
                if (Worker.getTotalNumOfJob(workers) >= numberOfWorkers) break;
            }
        } catch (KeeperException | InterruptedException | IOException e) {
            e.printStackTrace();
        }
    }

    private synchronized void startJobOnNode(Worker worker) throws IOException, InterruptedException, KeeperException {
        logger.info("start job at {} " , worker.getIpAddress());
        // Start a new worker instance on the specified node
        File file = new File(pathToProgram);
        String remoteDirectory = "/home/" + worker.getUserName() + "/";
        String remoteJarFilePath = remoteDirectory + file.getName();
        String userName = worker.getUserName();
        String nodeName = worker.getNodeName();
        String ip = worker.getIpAddress();

        // SCP command to copy the JAR file to the remote machine
        String scpCommand = "scp " + pathToProgram + " " +
                userName + "@" + ip + ":" + remoteDirectory;

        // SSH command to execute the JAR file remotely
        String sshCommand = "ssh " + userName + "@" + ip +
                " java -jar " + remoteJarFilePath + " " + nodeName;

        logger.info("first command: " + scpCommand);
        // Execute SCP command to copy the JAR file
        Process scpProcess = Runtime.getRuntime().exec(scpCommand);
        scpProcess.waitFor();
        // Check if SCP was successful before executing the JAR remotely
        if (scpProcess.exitValue() == 0) {
            logger.info("second command: " + sshCommand);
            // Execute SSH command to run the JAR file remotely
            Process sshProcess = Runtime.getRuntime().exec(sshCommand);
            worker.increaseOne();
            logger.info("The number of job on " + worker.getNodeName() + " is " + worker.getJobsNum());
        }
    }

    private void restartJobOnSameNode(Worker worker) throws IOException, InterruptedException, KeeperException {


    }

    @Override
    public void process(WatchedEvent event) {
        switch (event.getType()) {
            case NodeDeleted:
                // Handle changes in the child nodes of the watched znode (e.g., workers added or removed)
                logger.info("Job nodes deleted in " + event.getPath());
                String fullPath = event.getPath();
                String[] parts = fullPath.split("'");
                // Extract workerNode using Substring
                String workerNode = parts[1];

                if (event.getType() == Event.EventType.NodeDeleted)
                try {
                    launchWorkersIfNecessary();
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                break;
            default:
                List<String> jobs = null;
                try {
                    jobs = zooKeeper.getChildren(AUTOHEALER_ZNODES_PATH, false);
                    for (String job : jobs) {
                        zooKeeper.exists(AUTOHEALER_ZNODES_PATH + "/" + job, this);
                    }
                } catch (KeeperException e) {
                    throw new RuntimeException(e);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                break;
        }
    }
}

