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.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.FXModelRspec;
import be.iminds.ilabt.jfed.rspec.model.javafx_impl.FXRspecNode;
import be.iminds.ilabt.jfed.util.BestNodeLoginFinder;
import be.iminds.ilabt.jfed.util.KeyUtil;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.connection.channel.direct.Session;
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 javax.annotation.Nonnull;
import javax.net.SocketFactory;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.security.KeyPair;
import java.util.List;
import java.util.stream.Collectors;

/**
 * User: twalcari
 * Date: 5/26/2016
 * Time: 16:22
 */
public class DistributeSshKeysState extends SlicedState<DistributeSshKeysState.DistributeSshKeysStateSlice> {
    private static final Logger LOG = LoggerFactory.getLogger(DistributeSshKeysState.class);

    private final List<DistributeSshKeysStateSlice> slices;

    private final Experiment experiment;
    private final GeniUserProvider geniUserProvider;
    private final BestNodeLoginFinder bestNodeLoginFinder;

    private final ProxySocketFactoryProvider proxySocketFactoryProvider;

    public DistributeSshKeysState(Job<?> job, GeniUserProvider geniUserProvider,
                                  ProxySocketFactoryProvider proxySocketFactoryProvider) {
        super("Distribute generated SSH keys");
        this.experiment = job.getExperiment();
        this.geniUserProvider = geniUserProvider;
        this.proxySocketFactoryProvider = proxySocketFactoryProvider;

        slices = experiment.getKeypairs().entrySet().stream()
                .map(entry ->
                        new DistributeSshKeysStateSlice(job, entry.getKey(), entry.getValue()))
                .collect(Collectors.toList());

        assert experiment.getSlice() != null;
        bestNodeLoginFinder = new BestNodeLoginFinder(experiment.getSlice().getManifestRspec().getStringRspec(),
                null, 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 List<DistributeSshKeysStateSlice> getSlices() {
        return slices;
    }


    public class DistributeSshKeysStateSlice extends StateSlice {
        private final DistributeSshKeypair distributeSshKeypair;
        private final KeyPair keyPair;

        private final String location;
        //private final String folder;
        //private final String filename;

        public DistributeSshKeysStateSlice(@Nonnull Job<?> job, @Nonnull DistributeSshKeypair distributeSshKeypair, @Nonnull KeyPair keyPair) {
            super(job);
            this.distributeSshKeypair = distributeSshKeypair;
            this.keyPair = keyPair;

            String location = distributeSshKeypair.getLocation();
            if (location == null)
                location = "~/.ssh/id_rsa";

            this.location = location;
        }

        @Override
        public String getName() {
            return "Distribute generated SSH key for user '" + distributeSshKeypair.getUser().getResourceName() + "'";
        }

        @Override
        public ExperimentTaskStatus statefulRun() throws JFedException, InterruptedException {
            assert experiment.getSlice() != null;
            assert experiment.getSlice().getManifestRspec() != null;
            FXModelRspec manifestModel =
                    experiment.getSlice().getManifestRspec().getModelRspec(ModelRspecType.FX);
            assert manifestModel != null;

            if (location == null) {
                return ExperimentTaskStatus.FAILED;
            }

            boolean allSuccess = true;
            boolean allFailed = true;

            for (FXRspecNode node : manifestModel.getNodes()) {
                final UploadKeyToNodeState uploadKeyToNodeState = new UploadKeyToNodeState(node);
                setAndRunState(uploadKeyToNodeState);

                if (uploadKeyToNodeState.getStatus() != ExperimentTaskStatus.SUCCESS)
                    allSuccess = false;
                if (uploadKeyToNodeState.getStatus() != ExperimentTaskStatus.FAILED)
                    allFailed = false;
            }
            return allSuccess ? ExperimentTaskStatus.SUCCESS :
                    allFailed ? ExperimentTaskStatus.FAILED : ExperimentTaskStatus.WARNING;
        }


        public class UploadKeyToNodeState extends State {
            private final FXRspecNode node;

            protected UploadKeyToNodeState(FXRspecNode node) {
                super("Upload generated key for '" + distributeSshKeypair.getUser().getResourceName() + "' to '" + node.getClientId() + "'");
                this.node = node;
            }

            @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("Skipping distribution of sshkey to {} as no login-service is available", node.getClientId());
                    return ExperimentTaskStatus.FAILED;
                }

                //connect to the node and deposit the key
                SSHClient 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());
                    ssh.authPublickey(loginService.getUsername(), new KeyPairWrapper(keyPair));

                    ExternalFileUtil externalFileUtil = new ExternalFileUtil(ssh);

                    char[] keyFile = KeyUtil.privateKeyToPem(keyPair.getPrivate());

                    String localFolder = externalFileUtil.getFolder(location);
                    String filename = externalFileUtil.getFilename(location);

                    try (Session session = ssh.startSession()) {
                        Session.Command command = session.exec("mkdir -p " + localFolder.replace(" ", "\\ "));
                        command.join();
                        if (command.getExitStatus() != 0) {
                            String errorMessage;
                            try (BufferedReader buffer = new BufferedReader(new InputStreamReader(command.getErrorStream()))) {
                                errorMessage = buffer.lines().collect(Collectors.joining(" "));
                            }
                            LOG.warn("Could not create directory to upload key into: {}", errorMessage);
                            updateMessage("Could not create directory to upload key into: " + errorMessage);
                        } else {
                            LOG.info("mkdir command for '{}' completed succesfully", localFolder);
                        }
                    }

                    ssh.newSCPFileTransfer().upload(new StringMemorySourceFile(filename, keyFile), localFolder);
                    LOG.debug("Succesfully uploaded private key to {} on {}", distributeSshKeypair.getLocation(), node.getClientId());

                    ssh.startSession().exec("chmod 600 " + localFolder + "/" + filename);

                    return ExperimentTaskStatus.SUCCESS;
                } catch (UserAuthException ex) {
                    LOG.error("Could not authenticate on node {} with generated key. Skipping.", node.getClientId(), ex);
                    updateMessage("Could not authenticate with generated key.");
                    return ExperimentTaskStatus.FAILED;
                } catch (IOException ex) {
                    LOG.error("Error while transferring private key {} to {}", distributeSshKeypair.getLocation(), node.getClientId(), ex);
                    updateMessage("Error while transferring private key: " + ex.getMessage());
                    return ExperimentTaskStatus.FAILED;
                } finally {
                    try {
                        ssh.close();
                    } catch (IOException ignored) {
                    }
                }
            }
        }
    }
}
