package be.iminds.ilabt.jfed.experiment;

import be.iminds.ilabt.jfed.highlevel.model.Slice;
import be.iminds.ilabt.jfed.highlevel.util.ProxySocketFactoryProvider;
import be.iminds.ilabt.jfed.lowlevel.GeniUser;
import be.iminds.ilabt.jfed.lowlevel.connection.JFedConnection;
import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshKeyInfo;
import be.iminds.ilabt.jfed.lowlevel.ssh_key_info.SshKeyInfoFactory;
import be.iminds.ilabt.jfed.preferences.ProxyPreferencesManager;
import be.iminds.ilabt.jfed.rspec.model.LoginService;
import be.iminds.ilabt.jfed.rspec.model.ModelRspecType;
import be.iminds.ilabt.jfed.rspec.model.controller.ExperimentCommand;
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.ExecuteOnNotNull;
import be.iminds.ilabt.jfed.util.ProxyServiceUtil;
import be.iminds.ilabt.jfed.util.SshProxySocketFactory;
import com.jcraft.jsch.agentproxy.*;
import com.jcraft.jsch.agentproxy.sshj.AuthAgent;
import javafx.application.Platform;
import javafx.beans.InvalidationListener;
import javafx.beans.binding.ObjectBinding;
import javafx.beans.property.*;
import javafx.beans.value.ObservableBooleanValue;
import javafx.beans.value.ObservableObjectValue;
import javafx.collections.FXCollections;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import net.schmizz.sshj.userauth.UserAuthException;
import net.schmizz.sshj.userauth.keyprovider.KeyPairWrapper;
import net.schmizz.sshj.userauth.method.AuthMethod;
import net.schmizz.sshj.userauth.method.AuthPublickey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.net.SocketFactory;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.util.*;

/**
 * SshConnectionPool looks at the FXRspecModel and provides access to connections and to properties reflecting
 * if connections to nodes are available or can be made available
 */
public class SshConnectionPool implements ExperimentCommand.SshConnectionProvider {
    private final static Logger LOG = LoggerFactory.getLogger(SshConnectionPool.class);

    private final AgentProxy agentProxy;

    private final ProxySocketFactoryProvider proxySocketFactoryProvider;
    private final ProxyServiceUtil proxyServiceUtil;
    private final ProxyPreferencesManager proxyPreferencesManager;

    //Mapping nodeID's to their SSH-connection and other info
    private final Map<String, NodeInfo> nodeInfoByUniqueId = new HashMap<>();

    private final Experiment experiment;
    private final ObjectProperty<FXModelRspec> modelRspecValue = new SimpleObjectProperty<>();

    private final ReadOnlyListWrapper<String> nodeIds =
            new ReadOnlyListWrapper<>(FXCollections.<String>observableArrayList());

    private final BooleanProperty allConnectable = new SimpleBooleanProperty(false);
    private final BooleanProperty allConnected = new SimpleBooleanProperty(false);
    private final BooleanProperty anyConnectable = new SimpleBooleanProperty(false);
    private final BooleanProperty anyConnected = new SimpleBooleanProperty(false);

    private final GeniUser sshUser;

    public SshConnectionPool(Experiment experiment, GeniUser sshUser,
                             ProxySocketFactoryProvider proxySocketFactoryProvider,
                             ProxyPreferencesManager proxyPreferencesManager,
                             ProxyServiceUtil proxyServiceUtil) {
        this.experiment = experiment;
        this.sshUser = sshUser;
        this.proxySocketFactoryProvider = proxySocketFactoryProvider;
        this.proxyPreferencesManager = proxyPreferencesManager;
        this.proxyServiceUtil = proxyServiceUtil;

        this.modelRspecValue.addListener((observable) -> updateModel());

        new ExecuteOnNotNull<Slice>(experiment.sliceProperty()) {
            @Override
            public void run() {
                modelRspecValue.bind(new ObjectBinding<FXModelRspec>() {
                    {
                        assert experiment.getSlice() != null;
                        bind(experiment.getSlice().manifestRspecProperty());
                    }

                    @Override
                    protected FXModelRspec computeValue() {
                        assert experiment.getSlice() != null;
                        if (experiment.getSlice().getManifestRspec() != null)
                            return experiment.getSlice().getManifestRspec().getModelRspec(ModelRspecType.FX);
                        else
                            return null;
                    }
                });
                updateModel();
            }
        };

        AgentProxy agentProxy = null;
        try {
            Connector agentConnector = ConnectorFactory.getDefault().createConnector();
            agentProxy = new AgentProxy(agentConnector);
        } catch (AgentProxyException e) {
            LOG.error("Could not initialize agent connector");
        }
        this.agentProxy = agentProxy;
    }

    private final InvalidationListener updateAllConnectionStateListener = (observable) -> updateAllConnectionState();

    private void updateAllConnectionState() {
        boolean newAllConnectable = true;
        boolean newAllConnected = true;
        boolean newAnyConnectable = false;
        boolean newAnyConnected = false;
        for (NodeInfo nodeInfo : nodeInfoByUniqueId.values()) {
            switch (nodeInfo.state.get()) {
                case NO_INFO: {
                    newAllConnectable = false;
                    newAllConnected = false;
                    break;
                }
                case CONNECTABLE: {
                    newAnyConnectable = true;
                    newAllConnected = false;
                    break;
                }
                case CONNECTING: {
                    newAnyConnectable = true;
                    newAllConnected = false;
                    break;
                }
                case CONNECTED: {
                    newAnyConnectable = true;
                    newAnyConnected = true;
                    break;
                }
                case FAILCONNECT: {
                    newAllConnected = false;
                    break;
                }
            }
//            if (!newAllConnectable && !newAllConnected &&
//                    newAllConnectable && newAnyConnected)
//                break;
        }
        allConnectable.setValue(newAllConnectable);
        allConnected.setValue(newAllConnected);
        anyConnectable.setValue(newAnyConnectable);
        anyConnected.setValue(newAnyConnected);
    }

    private void updateModel() {
        if (modelRspecValue.get() == null)
            return;

        for (FXRspecNode n : modelRspecValue.get().getNodes()) {
            String nodeId = n.getUniqueId();
            getNodeInfo(nodeId).update(n);
        }
    }

    private void closeAll() {
        nodeInfoByUniqueId.values().forEach(NodeInfo::close);
    }

    public ReadOnlyListProperty<String> getNodeIds() {
        return nodeIds.getReadOnlyProperty();
    }

    private NodeInfo getNodeInfo(String nodeId) {
        NodeInfo ni = nodeInfoByUniqueId.get(nodeId);
        if (ni == null) {
            ni = new NodeInfo(nodeId);
            ni.state.addListener(updateAllConnectionStateListener);
            nodeInfoByUniqueId.put(nodeId, ni);
            nodeIds.get().add(nodeId);
        }
        assert Objects.equals(ni.nodeId, nodeId);
        return ni;
    }

    public ObservableObjectValue<NodeState> getNodeState(String nodeId) {
        return getNodeInfo(nodeId).state;
    }

    public ObservableBooleanValue isNodeConnectable(String nodeId) {
        return getNodeInfo(nodeId).state.isEqualTo(NodeState.NO_INFO).not();
    }

    public ObservableBooleanValue isNodeConnected(String nodeId) {
        return getNodeInfo(nodeId).state.isEqualTo(NodeState.CONNECTED);
    }

    public ObservableBooleanValue allConnectableProperty() {
        return allConnectable;
    }

    public ObservableBooleanValue allConnectedProperty() {
        return allConnected;
    }

    public BooleanProperty anyConnectableProperty() {
        return anyConnectable;
    }

    public BooleanProperty anyConnectedProperty() {
        return anyConnected;
    }

    /**
     * @param nodeId       node to get connection to
     * @param blockingMode
     * @return
     */
    public SSHClient getConnection(String nodeId, ExperimentCommand.BlockingMode blockingMode) {
        NodeInfo ni = getNodeInfo(nodeId);

        //make sure that we don't return stale connections
        if (ni.connection != null && !ni.connection.isConnected()) {
            ni.close();
        }

        switch (blockingMode) {
            case NON_BLOCKING: {
                switch (ni.state.get()) {
                    case NO_INFO:
                    case FAILCONNECT:
                        return null;
                    case CONNECTABLE: {
                        ni.startConnect();
                        return null;
                    }
                    case CONNECTING:
                        return null;
                    case CONNECTED:
                        return ni.connection;
                    default:
                        throw new RuntimeException("unsupported case " + ni.state.get());
                }
            }
            case WAIT_FOR_SINGLE: {
                switch (ni.state.get()) {
                    case NO_INFO:
                    case FAILCONNECT:
                        return null;
                    case CONNECTABLE: {
                        ni.startConnect();
                        ni.waitForConnect();
                        return ni.connection;
                    }
                    case CONNECTING: {
                        ni.waitForConnect();
                        return ni.connection;
                    }
                    case CONNECTED:
                        return ni.connection;
                    default:
                        throw new RuntimeException("unsupported case " + ni.state.get());
                }
            }
            case WAIT_FOR_ALL: {
                throw new RuntimeException("WAIT_FOR_ALL not yet implemented");
            }
            default:
                throw new RuntimeException("unsupported blockingMode " + blockingMode);
        }
    }


    /**
     * Creates an SSH-connection for every node. Every connection is
     * authenticated with the info provided by loginServices.
     * <p>
     * The connections are initialised in another thread, so this call is non-blocking.
     */
    public void initAllConnection() {
        for (FXRspecNode node : modelRspecValue.get().getNodes()) {
            final String nodeID = node.getUniqueId();
            NodeInfo nodeInfo = getNodeInfo(nodeID);
            NodeState state = nodeInfo.state.get();
            if (!Objects.equals(state, NodeState.NO_INFO) && !Objects.equals(state, NodeState.CONNECTED))
                nodeInfo.startConnect();
        }
    }

    @Nullable
    private LoginService chooseLoginService(@Nullable List<LoginService> loginServices) {
        if (loginServices == null || loginServices.isEmpty())
            return null;

        LoginService chosenLoginService = loginServices.get(0);
        for (int i = 1; i < loginServices.size(); i++) {
            LoginService otherLoginService = loginServices.get(i);
            if (Objects.equals(otherLoginService.getUsername(), sshUser.getUserUrn().getEncodedResourceName()))
                chosenLoginService = otherLoginService;
        }

        if (chosenLoginService == null || chosenLoginService.getHostname() == null || chosenLoginService.getIntPort() <= 0) {
            LOG.error(" invalid chosenLoginService skipped: " + chosenLoginService);
            return null;
        }

        return chosenLoginService;
    }

    private class NodeInfo {
        final String nodeId;

        SSHClient connection;
        Thread connectionCreatorThread;

        final ObjectProperty<NodeState> state = new SimpleObjectProperty<>(NodeState.NO_INFO);

        private LoginService loginService;
        private JFedConnection.SshProxyInfo proxyService;

        private NodeInfo(String nodeId) {
            this.nodeId = nodeId;
        }

        public void update(FXRspecNode rspecNode) {
            List<LoginService> loginServices = rspecNode.getLoginServices();

            if (!loginServices.isEmpty()) {
                loginService = chooseLoginService(rspecNode.getLoginServices());
                if (loginService != null)
                    proxyService = proxyServiceUtil.findTestbedProxy(rspecNode, loginService);
            } else {
                loginService = null;
                proxyService = null;
            }

            if (loginService == null) {
                close();
                state.setValue(NodeState.NO_INFO);
            } else {
                if (Objects.equals(state.get(), NodeState.NO_INFO))
                    state.set(NodeState.CONNECTABLE);
            }
        }

        public void close() {
            if (connection != null) {
                try {
                    connection.close();
                } catch (IOException ignored) {
                }
                connection = null;
            }
            if (loginService != null)
                state.set(NodeState.CONNECTABLE);
            else
                state.setValue(NodeState.NO_INFO);
        }

        public void startConnect() {
            if (state.get() == NodeState.NO_INFO)
                return;


            //The thread will automatically die when the runnable returns
            connectionCreatorThread = new Thread(this::blockingConnect, "ConnectionThread#" + nodeId);
            connectionCreatorThread.start();
        }

        public void waitForConnect() {
            if (connectionCreatorThread != null)
                try {
                    connectionCreatorThread.join();
                    connectionCreatorThread = null;
                } catch (InterruptedException e) {
                    LOG.debug("connectionCreatorThread.join() interrupted", e);
                }
        }

        public SSHClient blockingConnect() {
            int attemptCount = 1;
            boolean success = false;
            safeSetState(NodeState.CONNECTING);
            //Sometimes the connection gave a timeout on the first try, so retry once if this is the case
            while (attemptCount < 3 && !success) {
                connection = new SSHClient();
                connection.addHostKeyVerifier(new PromiscuousVerifier());
                try {
                    SshKeyInfo sshKeyInfo = SshKeyInfoFactory.createGeniUserSshKeyInfo(sshUser);

                    assert sshUser.getUserAuthorityServer().getTestbed() != null;

                    LOG.debug("Creating SSH connection to " + loginService.getHostname() + ":" + loginService.getIntPort());
                    SocketFactory socketFactory = proxySocketFactoryProvider.createProxySocketFactoryForSsh(
                            proxyService != null ? null : loginService.getHostname(),
                            proxyService != null ? null : loginService.getIntPort());
                    if (socketFactory != null) {
                        LOG.debug("Using user-authority proxy for  SSH connection");
                    }

                    if (proxyService != null) {
                        LOG.info("Wrapping SSH connection into proxy defined by testbed: {}", proxyService);
                        if (proxyPreferencesManager.isDnsOverProxyRequiredForSsh())
                            socketFactory =
                                    SshProxySocketFactory.createWithForcedTarget(proxyService, socketFactory, loginService.getHostname(), loginService.getIntPort());
                        else
                            socketFactory = SshProxySocketFactory.create(proxyService, socketFactory);
                    }

                    if (socketFactory != null)
                        connection.setSocketFactory(socketFactory);

                    try {
                        connection.connect(loginService.getHostname(), loginService.getIntPort());
                        LOG.debug("Connected to server.");
                    } catch (IOException ex) {

                        LOG.debug("Connection to server seems to have failed (ConnectionInfo is null), will try to continue anyway.");
                    }


                    final List<AuthMethod> authMethods = new ArrayList<>();
                    if (agentProxy != null) {
                        authMethods.addAll(getAuthMethods(agentProxy));
                    }
                    KeyPairWrapper keypair =
                            new KeyPairWrapper(sshKeyInfo.getPublicKey(), sshKeyInfo.getPrivateKey());

                    authMethods.add(new AuthPublickey(keypair));

                    try {
                        connection.auth(loginService.getUsername(), authMethods);
                    } catch (UserAuthException ex) {
                        LOG.warn("Could not authenticate connection for FXRspecNode: {}, for username {}", nodeId, loginService.getUsername(), ex);
                        throw ex;
                    }

                    success = true;
                    safeSetState(NodeState.CONNECTED);
                } catch (AssertionError ex) {
                    LOG.error("Could not setup a connection to node {} due to AssertionError", nodeId, ex);
                    safeSetState(NodeState.FAILCONNECT);
                } catch (Exception ex) {
                    attemptCount++;
                    if (attemptCount < 3 && ex instanceof SocketTimeoutException) {
                        try {
                            Thread.sleep(100);
                        } catch (InterruptedException e) {
                            LOG.debug("InterruptedException while waiting for next connect attempt. Will cancel connecting.", e);
                            safeSetState(NodeState.FAILCONNECT);
                            return null;
                        }
                        LOG.info("Connection for FXRspecNode " + nodeId + " gave a SocketTimeoutException. Retrying once more..");
                    }
                    //Log if not a SocketTimeoutException or if it still occured after 2 tries
                    else {
                        LOG.error("Could not setup a connection to node {}", nodeId, ex);
                        safeSetState(NodeState.FAILCONNECT);
                    }
                }
            }

            return connection;
        }

        public void safeSetState(final NodeState newState) {
            Platform.runLater(() -> state.setValue(newState));
        }

    }

    private static List<AuthMethod> getAuthMethods(AgentProxy agent) throws Exception {
        Identity[] identities = agent.getIdentities();
        List<AuthMethod> result = new ArrayList<>();
        for (Identity identity : identities) {
            result.add(new AuthAgent(agent, identity));
        }
        return result;
    }

    public enum NodeState {NO_INFO, CONNECTABLE/*but unconnected*/, CONNECTING, CONNECTED, FAILCONNECT}
}
