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

import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.PuTTYFilesKeyInfo;
import be.iminds.ilabt.jfed.preferences.GuiPreferenceKey;
import be.iminds.ilabt.jfed.preferences.JFedPreferences;
import be.iminds.ilabt.jfed.util.KeyUtil;
import com.jcraft.jsch.agentproxy.*;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.inject.Inject;
import javax.inject.Singleton;
import java.io.File;
import java.io.IOException;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.util.*;
import java.util.stream.Collectors;

@Singleton
public class PageantHelper {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(PageantHelper.class);

    private final JFedPreferences jFedPreferences;

    @Inject
    public PageantHelper(JFedPreferences jFedPreferences) {
        this.jFedPreferences = jFedPreferences;
    }

    @Nullable
    private static AgentProxy getAgentProxy() {
        try {
            Connector agentConnector = ConnectorFactory.getDefault().createConnector();
            return new AgentProxy(agentConnector);
        } catch (AgentProxyException e) {
            LOG.error("Could not initialize agent connector");
            return null;
        }
    }

    public synchronized void registerKeys(Collection<? extends PuTTYFilesKeyInfo> sshKeyInfos, SshKeyInfoUnlocker unlocker) {
        if (sshKeyInfos == null || sshKeyInfos.isEmpty()) return;

        AgentProxy agentProxy = getAgentProxy();

        if (agentProxy == null || !agentProxy.isRunning()) {
            //pageant is not running. Try to start it!

            String pageantLocation =
                    jFedPreferences.getFile(GuiPreferenceKey.PREF_PUTTY_DIRECTORY).getPath() + File.separator + "pageant.exe";
            try {
                executeCommand(Collections.singletonList(pageantLocation));
            } catch (IOException e) {
                LOG.error("Error while executing pageant");
            }
            try {
                Thread.sleep(750); //allow pageant to start
            } catch (InterruptedException ignored) {
            }

            agentProxy = getAgentProxy();
        }


        if (agentProxy != null) {
            final AgentProxy finalAgentProxy = agentProxy;

            List<PublicKey> existingPubkeys = Arrays.stream(finalAgentProxy.getIdentities())
                    .map(Identity::getBlob)
                    .map(KeyUtil::blobToPublicKey)
                    .collect(Collectors.toList());

            sshKeyInfos.forEach(keyInfo -> {
                if (keyInfo == null) {
                    LOG.warn("Skipping a null keyInfo");
                    return;
                }

                boolean alreadyLoaded =
                        existingPubkeys.stream().anyMatch(existingKey -> Objects.equals(keyInfo.getPublicKey(), existingKey));

                if (alreadyLoaded) {
                    LOG.debug("Skipping already loaded key {}", keyInfo.getComment());
                    return;
                }

                if (keyInfo.isLocked()) {
                    //try to unlock the ssh key info
                    boolean success = unlocker.unlock(keyInfo);

                    if (!success)
                        return;
                }

                if (!(keyInfo.getPrivateKey() instanceof RSAPrivateCrtKey))
                    LOG.warn("Skipping a private key that isn't of type RSAPrivateCrtKey");
                else {
                    LOG.trace("Registering {} to pageant", keyInfo.getPublicKey());
                    finalAgentProxy.addIdentity(KeyUtil.keyToSshAgent((RSAPrivateCrtKey) keyInfo.getPrivateKey(),
                            keyInfo.getComment() != null ? keyInfo.getComment() : ""));
                }
            });
        } else {
            LOG.error("Could not create an agent proxy. Is Pageant running?");
        }
    }

    private static void executeCommand(List<String> command) throws IOException {
        Runtime.getRuntime().exec(command.toArray(new String[command.size()]));
    }
}
