package be.iminds.ilabt.jfed.ssh_terminal_tool.unix;

import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshFilesKeyInfo;
import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshKeyInfo;
import be.iminds.ilabt.jfed.util.IOUtils;
import be.iminds.ilabt.jfed.util.JFedUtils;
import be.iminds.ilabt.jfed.util.KeyUtil;
import be.iminds.ilabt.jfed.util.tmp_file_helpers.TmpFile;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.security.interfaces.RSAPublicKey;
import java.util.Arrays;
import java.util.Date;

/**
 * SshAgentHelper
 */
public class SshAgentHelper {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(SshAgentHelper.class);

    public static void registerKey(final SshFilesKeyInfo sshKeyInfo) {
        if (isKeyRegistered(sshKeyInfo)) return;

        if (sshKeyInfo.hasUnencryptedPrivateKey()) {
            TmpFile unencryptedPrivateKeyFile = sshKeyInfo.getUnencryptedPrivateKeyFile();
            unencryptedPrivateKeyFile.store();
            myCommandHelperWithinEnvironment(false, "ssh-add", unencryptedPrivateKeyFile.getFilename());
            unencryptedPrivateKeyFile.delete();
        } else {
            TmpFile encryptedPrivateKeyFile = sshKeyInfo.getPrivateKeyFile();
            encryptedPrivateKeyFile.store();
            myCommandHelperWithinEnvironment(false, "ssh-add", encryptedPrivateKeyFile.getFilename());
            encryptedPrivateKeyFile.delete();
        }

        if (!isKeyRegistered(sshKeyInfo))
            throw new RuntimeException("Key did not register successfully");
    }

    public static boolean isKeyRegistered(final SshKeyInfo sshKeyInfo) {
        String openSshFormatPubKey = KeyUtil.rsaPublicKeyToOpenSshAuthorizedKeysFormatData((RSAPublicKey) sshKeyInfo.getPublicKey());
        LOG.debug("Searching for pubkey \""+openSshFormatPubKey+"\"");

        String pubKeys = myCommandHelperWithinEnvironment(true, "ssh-add", "-L");

        if (pubKeys != null)
            for (String pubKey : pubKeys.split("\n")) {
                LOG.debug("     Checking pubkey \""+pubKey.trim()+"\"");
                if (pubKey.trim().contains(openSshFormatPubKey))
                    return true;
            }
        LOG.debug("Did not find matching pubkey in: "+pubKeys);

        return false;
    }

    private static String cachedSshAgentEnvironment = null;
    /**
     * This also starts an agent. (without caching, would start an agent for each call!)
     *
     * @return Multi line string containing bash commands to set ssh-agent environment vars for ssh command
     */
    public static String getSshAgentEnvironment() {
        if (cachedSshAgentEnvironment != null) return cachedSshAgentEnvironment;

        cachedSshAgentEnvironment = myCommandHelper(false, "ssh-agent", "-s");
        LOG.debug("cachedSshAgentEnvironment: "+cachedSshAgentEnvironment);

        cachedSshAgentEnvironment = cachedSshAgentEnvironment.replaceFirst("echo Agent pid [0-9]*;", "");
        LOG.debug("cachedSshAgentEnvironment without echo: "+cachedSshAgentEnvironment);

        return cachedSshAgentEnvironment;
    }
//    private static Map<String, String> cachedSshAgentEnvironmentMap = null;
//    public static Map<String, String> getSshAgentEnvironmentMap() {
//        if (cachedSshAgentEnvironmentMap != null) return cachedSshAgentEnvironmentMap;
//
//        cachedSshAgentEnvironmentMap = new HashMap<>();
//
//        /*
//            example:
//                SSH_AUTH_SOCK=/tmp/ssh-zLVhor4084/agent.4084; export SSH_AUTH_SOCK;
//                SSH_AGENT_PID=4085; export SSH_AGENT_PID;
//         */
//
//        String env = getSshAgentEnvironment();
//        for () {
//
//        }
//
//        return cachedSshAgentEnvironmentMap;
//    }

    public static void stopSshAgent() {
        if (cachedSshAgentEnvironment != null) {
            myCommandHelperWithinEnvironment(true, "ssh-agent", "-k");
            cachedSshAgentEnvironment = null;
        }
    }



    private static String myCommandHelperWithinEnvironment(boolean mayFail, String ... command) {
        File tmpDir = new File(JFedUtils.getUserDataDirectoryFile(), "tmp");
        if (!tmpDir.exists()) //noinspection ResultOfMethodCallIgnored
            tmpDir.mkdir();
        long nr = new Date().getTime();
        File tmpScript = new File(tmpDir, "ssh-agent-helper="+nr+".sh");

        String script = "#!/bin/bash\n\n"+getSshAgentEnvironment()+"\n";

        boolean first = true;
        for (String c : command) {
            if (first)
                script += c;
            else
                script += " '"+c+"'";
            first = false;
        }

        LOG.debug("Script: "+script);

        IOUtils.stringToFile(tmpScript, script);
        tmpScript.deleteOnExit();

        String res = myCommandHelper(mayFail, "bash", tmpScript.getPath());

        LOG.debug("Script output: "+res);

        //noinspection ResultOfMethodCallIgnored
        tmpScript.delete();

        return res;
    }
    private static String myCommandHelper(boolean mayFail, String ... command) {
        StringBuilder res = new StringBuilder();

        LOG.debug("Starting command: "+Arrays.asList(command));

        ProcessBuilder pb = new ProcessBuilder(command);
        pb.redirectErrorStream();

        Process p = null;
        try {
            p = pb.start();
            p.getOutputStream().close(); //we don't write to the command's stdin
            BufferedReader input = new BufferedReader(new InputStreamReader(p.getInputStream()));

            String line = null;
            while ((line = input.readLine()) != null) {
                res.append(line);
            }
            //note: input might close before terminal stops.
            input.close();
        } catch (IOException e) {
            LOG.error("Error executing command: "+e.getMessage(), e);
        }
        try {
            p.waitFor();
        } catch (InterruptedException e) {
            LOG.warn("SshAgentHelper.myCommandHelper thread interrupted while waiting for terminal to stop", e);
        }

        if (p.exitValue() != 0) {
            if (mayFail) {
                LOG.info("SshAgentHelper command \"" + command[0] + "\" returned non zero exit value: " + p.exitValue() + " output: " + res.toString());
                return null;
            }
            else
                throw new RuntimeException("SshAgentHelper command \""+command[0]+"\" returned non zero exit value: "+p.exitValue()+" output: "+res.toString());
 }
       
        return res.toString();
    }
}
