package be.iminds.ilabt.jfed.highlevel.jobs.states;

import be.iminds.ilabt.jfed.experiment.Experiment;
import be.iminds.ilabt.jfed.experiment.tasks.ExperimentTaskStatus;
import be.iminds.ilabt.jfed.highlevel.jobs.Job;
import be.iminds.ilabt.jfed.highlevel.jobs.SlicedState;
import be.iminds.ilabt.jfed.highlevel.jobs.State;
import be.iminds.ilabt.jfed.highlevel.jobs.StateSlice;
import be.iminds.ilabt.jfed.highlevel.util.ExternalFileUtil;
import be.iminds.ilabt.jfed.highlevel.util.ProxySocketFactoryProvider;
import be.iminds.ilabt.jfed.highlevel.util.StringMemorySourceFile;
import be.iminds.ilabt.jfed.lowlevel.GeniUserProvider;
import be.iminds.ilabt.jfed.lowlevel.JFedException;
import be.iminds.ilabt.jfed.rspec.model.AnsibleService;
import be.iminds.ilabt.jfed.rspec.model.BasicStringRspec;
import be.iminds.ilabt.jfed.rspec.model.DistributeSshKeypair;
import be.iminds.ilabt.jfed.rspec.model.ModelRspecType;
import be.iminds.ilabt.jfed.rspec.model.javafx_impl.FXExecuteService;
import be.iminds.ilabt.jfed.rspec.model.javafx_impl.FXModelRspec;
import be.iminds.ilabt.jfed.rspec.model.javafx_impl.FXRspecNode;
import be.iminds.ilabt.jfed.util.BestNodeLoginFinder;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import javafx.application.Platform;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.connection.ConnectionException;
import net.schmizz.sshj.connection.channel.direct.Session;
import net.schmizz.sshj.transport.TransportException;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import net.schmizz.sshj.userauth.UserAuthException;
import net.schmizz.sshj.userauth.keyprovider.KeyPairWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.event.Level;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.net.SocketFactory;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.KeyPair;
import java.util.*;

/**
 * User: twalcari
 * Date: 5/27/2016
 * Time: 11:51
 */
public class ExecuteAnsibleServicesState extends SlicedState<ExecuteAnsibleServicesState.ExecuteAnsibleServiceStateSlice> {
    private static final Logger LOG = LoggerFactory.getLogger(ExecuteAnsibleServicesState.class);
    public static final int MAX_EXECUTE_FINISHED_FILE_DETECTION_TRIES = 60;
    public static final int TIME_BETWEEN_FINISHED_FILE_DETECTION_TRIES_MS = 10_000;


    private final Job<?> job;
    private final Experiment experiment;
    private final GeniUserProvider geniUserProvider;
    private final ProxySocketFactoryProvider proxySocketFactoryProvider;

    @Nonnull
    private final KeyPair loginKeyPair;


    private final List<ExecuteAnsibleServiceStateSlice> slices;

    private final BestNodeLoginFinder bestNodeLoginFinder;

    protected ExecuteAnsibleServicesState(Job<?> job, GeniUserProvider geniUserProvider, ProxySocketFactoryProvider proxySocketFactoryProvider) {
        super("Execute ansible services");
        this.job = job;
        this.experiment = job.getExperiment();
        this.geniUserProvider = geniUserProvider;
        this.proxySocketFactoryProvider = proxySocketFactoryProvider;

        assert experiment.getSlice() != null;

        this.slices = new ArrayList<>();

        FXModelRspec model = experiment.getSlice().getRequestRspec().getModelRspec(ModelRspecType.FX);
        assert model != null;

        for (FXRspecNode node : model.getNodes()) {
            for (AnsibleService ansibleService : node.getAnsibleServices()) {
                slices.add(new ExecuteAnsibleServiceStateSlice(job, node, ansibleService));
            }
        }

        assert experiment.getKeypairs() != null;

        String distributedSshKeysUsername = null;
        KeyPair loginKeyPair = null;

        for (Map.Entry<DistributeSshKeypair, KeyPair> e : experiment.getKeypairs().entrySet()) {
            if(e.getValue() != null){
                //distributedSshKeysUsername may be null, but that's ok, that just means the logged in users account needs to be used.
                distributedSshKeysUsername = e.getKey().getUser() == null ? null : e.getKey().getUser().getResourceName();

                loginKeyPair = e.getValue();
                break;
            }
        }

        if(loginKeyPair == null)
            loginKeyPair = new KeyPair(geniUserProvider.getLoggedInGeniUser().getPublicKey(),
                    geniUserProvider.getLoggedInGeniUser().getPrivateKey());

        this.loginKeyPair = loginKeyPair;
        bestNodeLoginFinder = new BestNodeLoginFinder(experiment.getSlice().getManifestRspec().getStringRspec(),
                distributedSshKeysUsername, this.geniUserProvider.getLoggedInGeniUser(),
                new BestNodeLoginFinder.Feedback() {
                    @Override
                    public void info(String s) {
                        LOG.info("DistributeSshKeysStateSlice: " + s);
                    }

                    @Override
                    public void error(String s) {
                        LOG.error("DistributeSshKeysStateSlice: " + s);
                    }

                    @Override
                    public void debug(String s) {
                        LOG.debug("DistributeSshKeysStateSlice: " + s);
                    }
                });
    }

    @Override
    public Collection<ExecuteAnsibleServiceStateSlice> getSlices() {
        return slices;
    }

    public class ExecuteAnsibleServiceStateSlice extends StateSlice {

        protected final FXRspecNode node;
        protected final AnsibleService ansibleService;

        private SSHClient ssh = null;
        private ExternalFileUtil externalFileUtil;

        public ExecuteAnsibleServiceStateSlice(@Nonnull Job<?> job, FXRspecNode node, AnsibleService ansibleService) {
            super(job);
            this.node = node;
            this.ansibleService = ansibleService;
        }

        @Override
        public String getName() {
            return "Execute ansible-playbook " + ansibleService.getExecutePlaybook() + " on " + node.getClientId();
        }

        @Override
        public ExperimentTaskStatus statefulRun() throws JFedException, InterruptedException {
            final SetupSshConnectionState setupSshConnectionState = new SetupSshConnectionState();
            setAndRunState(setupSshConnectionState);

            if (setupSshConnectionState.getStatus() != ExperimentTaskStatus.SUCCESS)
                return ExperimentTaskStatus.FAILED;

            //check if there are finished-files to wait for
            if (node.getExecuteServices().stream().anyMatch(es -> es.getFinishedFile() != null)) {
                setAndRunState(new WaitForExecuteServicesState());
            } else {
                LOG.info("Did not find any execute-services with finished-file attribute. We assume that all files that are needed are present!");
            }

            if (ansibleService.getInventory() != null)
                setAndRunState(new UploadInventoryState());
            else LOG.debug("No inventory-path specified. Skipping upload inventory");

            setAndRunState(new ExecuteAnsibleServiceState());

            try {
                if (ssh != null)
                    ssh.close();
            } catch (IOException ignored) {
            }

            return null;
        }

        public class SetupSshConnectionState extends State {

            protected SetupSshConnectionState() {
                super("Setting up SSH Connection to '" + node.getClientId() + "'");
            }

            @Nonnull
            @Override
            protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {
                //get login-service associated to the user we are distributing the key to
                BasicStringRspec.LoginService loginService = bestNodeLoginFinder.findBestLogin(node.getUniqueId());

                if(loginService == null){
                    LOG.warn("Cannot setup connection to {} as no login-service is available", node.getClientId());
                    return ExperimentTaskStatus.FAILED;
                }

                ssh = new SSHClient();
                ssh.addHostKeyVerifier(new PromiscuousVerifier());

                try {
                    SocketFactory socketFactory = proxySocketFactoryProvider.createProxySocketFactory(loginService.getHostname(), loginService.getPort());
                    if (socketFactory != null) {
                        LOG.debug("Using proxy for distributing keys");
                        ssh.setSocketFactory(socketFactory);
                    } else {
                        LOG.debug("No proxy needed for distributing keys");
                    }
                } catch (IOException ex) {
                    LOG.error("Error while creating ProxySocketFactory. Trying without now.", ex);
                }

                try {
                    ssh.connect(loginService.getHostname(), loginService.getPort());
                } catch (IOException e) {
                    updateMessage("Error while connecting to " + loginService.getHostname() + ": " + e.getMessage());
                    LOG.error("Error while connecting to {}: {}", loginService.getHostname(), e.getMessage(), e);
                    return ExperimentTaskStatus.FAILED;
                }
                try {
                    //we must use the distributed ssh keys to login: in speaksfor case, there is no other way to login
                    ssh.authPublickey(loginService.getUsername(), new KeyPairWrapper(loginKeyPair));

                    externalFileUtil = new ExternalFileUtil(ssh);
                } catch (UserAuthException e) {
                    updateMessage("Could not authenticate: " + e.getMessage());
                    LOG.error("Authentication failed on {}: {}", node.getClientId(), e.getMessage(), e);
                    return ExperimentTaskStatus.FAILED;
                } catch (TransportException e) {
                    updateMessage("Error while authenticating: " + e.getMessage());
                    LOG.error("Error while authenticating on {}: {}", node.getClientId(), e.getMessage(), e);
                    return ExperimentTaskStatus.FAILED;
                }

                return ExperimentTaskStatus.SUCCESS;

            }
        }

        public class WaitForExecuteServicesState extends State {

            protected WaitForExecuteServicesState() {
                super("Waiting for execute-services to finish on " + node.getClientId());
            }

            @Nonnull
            @Override
            protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {
                try {
                    for (FXExecuteService executeService : node.getExecuteServices()) {
                        if (executeService.getFinishedFile() == null)
                            continue;

                        boolean fileFound = false;
                        int tries = 0;

                        while (!fileFound && tries < MAX_EXECUTE_FINISHED_FILE_DETECTION_TRIES) {
                            Session.Command command =
                                    ssh.startSession().exec("test -f " + executeService.getFinishedFile());
                            command.join();

                            fileFound = command.getExitStatus() == 0;

                            tries++;
                            if (!fileFound) {
                                updateMessage("Looked for '" + executeService.getFinishedFile() + "' " + tries + " times.");
                                Thread.sleep(TIME_BETWEEN_FINISHED_FILE_DETECTION_TRIES_MS);
                            }
                        }

                        if (!fileFound) {
                            updateMessage("Timed out while waiting for finished-file '" + executeService.getFinishedFile() + "' to appear");
                            return ExperimentTaskStatus.FAILED;
                        } else {
                            updateMessage("Found finished-file '" + executeService.getFinishedFile() + "'");
                        }
                    }

                    updateMessage("Found all finished-files");
                    return ExperimentTaskStatus.SUCCESS;


                } catch (IOException ex) {
                    LOG.error("Error while waiting for finished-files on {}", node.getClientId(), ex);
                    updateMessage("Error while waiting for finished-files on {}" + ex.getMessage());
                    return ExperimentTaskStatus.FAILED;
                }
            }
        }

        public class UploadInventoryState extends State {

            private final FXModelRspec request;
            private final FXModelRspec manifest;

            private final BestNodeLoginFinder bestNodeLoginFinder;

            protected UploadInventoryState() {
                super("Uploading inventory on " + node.getClientId());

                //initialize request and manifest-model
                assert experiment.getNewRequestRspecSource() != null;
                request = experiment.getNewRequestRspecSource().getModelRspec(ModelRspecType.FX);

                if (request == null) {
                    LOG.error("Cannot create inventory file because parsing request rspec failed");
                    throw new RuntimeException("Cannot create inventory file because parsing request rspec failed");
                }


                assert experiment.getSlice() != null;
                manifest = experiment.getSlice().getManifestRspec().getModelRspec(ModelRspecType.FX);

                if (manifest == null) {
                    LOG.error("Cannot create inventory file because parsing manifest rspec failed");
                    throw new RuntimeException("Cannot create inventory file because parsing manifest rspec failed");
                }

                bestNodeLoginFinder = new BestNodeLoginFinder(experiment.getSlice().getManifestRspec().getStringRspec(),
                        null,
                        geniUserProvider.getLoggedInGeniUser(),
                        new BestNodeLoginFinder.Feedback() {
                            @Override
                            public void info(String s) {
                                LOG.info(s);
                            }

                            @Override
                            public void error(String s) {
                                LOG.error(s);
                            }

                            @Override
                            public void debug(String s) {
                                LOG.debug(s);
                            }
                        });
            }

            public String createInventoryFile() {
                //construct a multimap containing the nodes grouped per ansible-group


                Multimap<String, FXRspecNode> nodesPerGroup = HashMultimap.create();
                for (FXRspecNode node : request.getNodes()) {
                    if (!node.getAnsibleGroups().isEmpty()) {
                        node.getAnsibleGroups().forEach(groupName -> nodesPerGroup.put(groupName, node));
                    } else {
                        nodesPerGroup.put(null, node);
                    }
                }

                //now write the inventory file
                StringBuilder inventory = new StringBuilder();

                //first write the groupless nodes
                if (nodesPerGroup.containsKey(null)) {
                    Collection<FXRspecNode> grouplessNodes = nodesPerGroup.removeAll(null);
                    grouplessNodes.forEach(node -> writeNodeToInventory(node, inventory));
                    inventory.append('\n');
                }

                //now write all groups

                for (String groupName : nodesPerGroup.keySet()) {
                    inventory.append("[").append(groupName).append("]\n");
                    nodesPerGroup.get(groupName).forEach(node -> writeNodeToInventory(node, inventory));
                    inventory.append('\n');
                }

                return inventory.toString();
            }

            private void writeNodeToInventory(FXRspecNode node, StringBuilder inventory) {
                final BasicStringRspec.LoginService loginService =
                        bestNodeLoginFinder.findBestLogin(node.getUniqueId());

                inventory.append(node.getClientId())
                        .append("\tansible_host=").append(loginService.getHostname());

                if (loginService.getPort() != 22) {
                    inventory.append("\tansible_port=").append(loginService.getPort());
                }

                if (!Objects.equals(loginService.getUsername(), geniUserProvider.getLoggedInGeniUser().getUserUrn().getResourceName())) {
                    inventory.append("\tansible_user=").append(loginService.getUsername());
                }

                inventory.append('\n');
            }

            @Nonnull
            @Override
            protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {
                String inventoryFilename;
                if (ansibleService.getInventory() != null) {
                    inventoryFilename = ansibleService.getInventory();
                } else {
                    inventoryFilename= "/etc/ansible/hosts"; //ansible default inventory file location
                }

                try {
                    //create inventory file
                    String inventory = createInventoryFile();

                    ssh.newSCPFileTransfer().upload(
                            new StringMemorySourceFile(externalFileUtil.getFilename(inventoryFilename), inventory),
                            externalFileUtil.getFolder(inventoryFilename));

                } catch (IOException ex) {
                    LOG.error("Error while uploading inventory-file to {} on {}", inventoryFilename, node.getClientId(), ex);
                    updateMessage(String.format("Error while uploading inventory-file to %s on %s: %s", inventoryFilename, node.getClientId(), ex.getMessage()));
                    return ExperimentTaskStatus.FAILED;
                }

                return ExperimentTaskStatus.SUCCESS;
            }
        }

        public class ExecuteAnsibleServiceState extends State {

            private final StringProperty galaxyInputStream, galaxyErrorStream;
            private final StringProperty playbookInputStream, playbookErrorStream;


            protected ExecuteAnsibleServiceState() {
                super("Executing ansible-playbook '" + ansibleService.getExecutePlaybook() + "' on " + node.getClientId());

                this.galaxyInputStream = new SimpleStringProperty("");
                this.galaxyErrorStream = new SimpleStringProperty("");
                this.playbookInputStream = new SimpleStringProperty("");
                this.playbookErrorStream = new SimpleStringProperty("");

            }

            @Nonnull
            @Override
            protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {

                try {
                    //first install the requirements if needed
                    if (ansibleService.getInstallRequirements() != null) {
                        int dividerPosition = ansibleService.getInstallRequirements().lastIndexOf('/');
                        String folder = ansibleService.getInstallRequirements().substring(0, dividerPosition);
                        String filename = ansibleService.getInstallRequirements().substring(dividerPosition + 1);

                        String galaxyCommand = ansibleService.getGalaxyCommand();
                        if (galaxyCommand == null)
                            galaxyCommand = "ansible-galaxy";

                        String galaxyCommandString = String.format("cd %s && %s install -r %s -p roles",
                                                                folder, galaxyCommand, filename);
                        LOG.debug("Executing galaxy command: "+galaxyCommandString);
                        Session.Command c =
                                ssh.startSession().exec(galaxyCommandString);

                        new InputStreamToStringPropertyReaderThread(c.getInputStream(), galaxyInputStream, Level.INFO, "GALAXY-OUTPUT: ").start();
                        new InputStreamToStringPropertyReaderThread(c.getErrorStream(), galaxyErrorStream, Level.WARN, "GALAXY-ERROR: ").start();

                        c.join();

                        if (c.getExitStatus() != 0) {
                            updateMessage("ansible-galaxy execution was unsuccessful. Aborting");
                            return ExperimentTaskStatus.FAILED;
                        }
                    }

                    assert ansibleService.getExecutePlaybook() != null;
                    int dividerPosition = ansibleService.getExecutePlaybook().lastIndexOf('/');
                    String folder = ansibleService.getExecutePlaybook().substring(0, dividerPosition);
                    String filename = ansibleService.getExecutePlaybook().substring(dividerPosition + 1);

                    String playbookCommand = ansibleService.getPlaybookCommand();
                    if (playbookCommand == null)
                        playbookCommand = "ansible-playbook";

                    String arguments = "";
                    arguments += filename;

                    if (ansibleService.getDebug() == Boolean.TRUE)
                        arguments += " -v";

                    if (ansibleService.getInventory() != null)
                        arguments += " -i '"+ansibleService.getInventory()+"'";

                    //this must be the LAST argument, as it is a redirect and pipe
                    if (ansibleService.getLogfile() != null)
                        arguments += " 2>&1 | tee '"+ansibleService.getLogfile()+"'";

                    String playbookCommandString = String.format("cd %s && %s %s", folder, playbookCommand, arguments);
                    LOG.debug("Executing playbook command: "+playbookCommandString);
                    Session.Command c =
                            ssh.startSession().exec(playbookCommandString);

                    new InputStreamToStringPropertyReaderThread(c.getInputStream(), playbookInputStream, Level.INFO, "PLAYBOOK-OUTPUT: ").start();
                    new InputStreamToStringPropertyReaderThread(c.getErrorStream(), playbookErrorStream, Level.WARN, "PLAYBOOK-ERROR: ").start();

                    c.join();


                    if (c.getExitStatus() != 0) {
                        updateMessage("ansible-playbook execution was unsuccessful.");
                        return ExperimentTaskStatus.FAILED;
                    }


                } catch (TransportException | ConnectionException ex) {
                    throw new JFedException("SSH-exception while executing ansible: " + ex.getMessage(), ex);
                }

                return ExperimentTaskStatus.SUCCESS;
            }

            public StringProperty getGalaxyInputStream() {
                return galaxyInputStream;
            }

            public StringProperty getGalaxyErrorStream() {
                return galaxyErrorStream;
            }

            public StringProperty getPlaybookInputStream() {
                return playbookInputStream;
            }

            public StringProperty getPlaybookErrorStream() {
                return playbookErrorStream;
            }

            public FXRspecNode getNode() {
                return node;
            }

            public AnsibleService getAnsibleService() {
                return ansibleService;
            }
        }
    }

    private static class InputStreamToStringPropertyReaderThread extends Thread {

        private final InputStream inputStream;
        private final StringProperty output;

        private final String logPrefix;
        private final LogFunction logFunction;


        private InputStreamToStringPropertyReaderThread(InputStream inputStream, StringProperty output,
                                                        @Nullable Level logLevel, String logPrefix) {
            this.inputStream = inputStream;
            this.output = output;
            this.logPrefix = logPrefix;

            if (logLevel != null) {
                switch (logLevel) {
                    case TRACE:
                        logFunction = LOG::trace;
                        break;
                    case DEBUG:
                        logFunction = LOG::debug;
                        break;
                    case INFO:
                        logFunction = LOG::info;
                        break;
                    case WARN:
                        logFunction = LOG::warn;
                        break;
                    case ERROR:
                        logFunction = LOG::error;
                        break;
                    default:
                        logFunction = null;
                }

            } else
                logFunction = null;
        }

        @Override
        public void run() {

            try {
                BufferedReader stdInput = new BufferedReader(new InputStreamReader(inputStream));

                String s;
                while ((s = stdInput.readLine()) != null) {
                    System.out.println(s);
                    String str = s;

                    if (logFunction != null) {
                        logFunction.log(logPrefix + str);
                    }

                    Platform.runLater(() -> output.setValue(output.get() + "\n" + str));

                }

                inputStream.close();
            } catch (IOException ex) {
                LOG.error("Error while reading from InputStream to StringProperty", ex);
            }

        }

        private interface LogFunction {
            void log(String text);
        }
    }

}
