package be.iminds.ilabt.jfed.experimenter_gui.slice.jobs;

import be.iminds.ilabt.jfed.connectivity_tester.ConnectivityTest;
import be.iminds.ilabt.jfed.connectivity_tester.HostAndPortTest;
import be.iminds.ilabt.jfed.experiment.Experiment;
import be.iminds.ilabt.jfed.experiment.ExperimentPart;
import be.iminds.ilabt.jfed.experiment.SfaExperimentPart;
import be.iminds.ilabt.jfed.experiment.tasks.ExperimentTaskStatus;
import be.iminds.ilabt.jfed.fedmon.webapi.service.json.Proxy;
import be.iminds.ilabt.jfed.highlevel.controller.TaskThread;
import be.iminds.ilabt.jfed.highlevel.jobs.*;
import be.iminds.ilabt.jfed.highlevel.model.Sliver;
import be.iminds.ilabt.jfed.highlevel.tasks.HighLevelTaskFactory;
import be.iminds.ilabt.jfed.lowlevel.GeniUser;
import be.iminds.ilabt.jfed.lowlevel.GeniUserProvider;
import be.iminds.ilabt.jfed.lowlevel.JFedException;
import be.iminds.ilabt.jfed.lowlevel.api_wrapper.StatusDetails;
import be.iminds.ilabt.jfed.lowlevel.connection.JFedConnection;
import be.iminds.ilabt.jfed.preferences.CorePreferenceKey;
import be.iminds.ilabt.jfed.preferences.JFedPreferences;
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.javafx_impl.FXModelRspec;
import be.iminds.ilabt.jfed.rspec.model.javafx_impl.FXRspecNode;
import be.iminds.ilabt.jfed.util.ProxyServiceUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.*;
import java.util.stream.Collectors;

/**
 * User: twalcari
 * Date: 2/3/2015
 * Time: 17:58
 */
public class TestConnectivityJob extends AbstractJob<Boolean> {
    private static final Logger LOG = LoggerFactory.getLogger(TestConnectivityJob.class);
    private static final int WAIT_BETWEEN_TESTS = 10; //in seconds
    private static final int MAX_TESTS = 12;

    private final GeniUserProvider geniUserProvider;
    private final JFedPreferences jFedPreferences;
    private final ProxyPreferencesManager proxyPreferencesManager;
    private final ProxyServiceUtil proxyServiceUtil;

    private final List<FXRspecNode> failedNodes = Collections.synchronizedList(new ArrayList<>());

    TestConnectivityJob(Experiment experiment, HighLevelTaskFactory hltf, TaskThread tt,
                        GeniUserProvider geniUserProvider, JFedPreferences jFedPreferences,
                        ProxyPreferencesManager proxyPreferencesManager, ProxyServiceUtil proxyServiceUtil) {
        super("Test connectivity to resources", experiment, hltf, tt);
        this.geniUserProvider = geniUserProvider;
        this.jFedPreferences = jFedPreferences;
        this.proxyPreferencesManager = proxyPreferencesManager;
        this.proxyServiceUtil = proxyServiceUtil;
    }


    @Override
    public Boolean execute() throws Exception {
        LOG.trace("Starting connectivity tests to existing resources in experiment {}", experiment.getName());

        final TestConnectivityToExperimentPartsState tcs = new TestConnectivityToExperimentPartsState();
        setAndRunState(tcs);

        return tcs.getStatus() == ExperimentTaskStatus.SUCCESS;
    }

    public class TestConnectivityToExperimentPartsState extends SlicedState<TestConnectivityToExperimentPartStateSlice> {

        private final List<TestConnectivityToExperimentPartStateSlice> slices;

        protected TestConnectivityToExperimentPartsState() {
            super("Test connectivity to experiment parts");

            slices = experiment.getParts().stream()
                    .filter(part -> part instanceof SfaExperimentPart) //only SFA parts are accesible for the moment
                    .map(TestConnectivityToExperimentPartStateSlice::new)
                    .collect(Collectors.toList());
        }

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

        public List<FXRspecNode> getFailedNodes() {
            return failedNodes;
        }
    }


    public class TestConnectivityToExperimentPartStateSlice extends ExperimentPartStateSlice {
        private TestConnectivityToExperimentPartStateSlice(ExperimentPart experimentPart) {
            super(TestConnectivityJob.this, experimentPart);
        }

        @Override
        public String getName() {
            return "Testing connectivity to resources in " + experimentPart.getName();
        }

        @Override
        public ExperimentTaskStatus statefulRun() throws JFedException, InterruptedException {
            FindTestableResourcesState findTestableResourcesState = new FindTestableResourcesState(experimentPart);

            setAndRunState(findTestableResourcesState);

            if (findTestableResourcesState.getStatus() == ExperimentTaskStatus.SUCCESS &&
                    !findTestableResourcesState.getNodesToTest().isEmpty()) {
                final TestConnectivityToExperimentPartState tcs =
                        new TestConnectivityToExperimentPartState(experimentPart, findTestableResourcesState.getNodesToTest());
                setAndRunState(tcs);

                tcs.nodesConnectivity.entrySet().stream()
                        .filter(entry -> entry.getValue() != ConnectivityTest.Status.SUCCEEDED)
                        .map(Map.Entry::getKey)
                        .forEach(failedNodes::add);


                return tcs.getStatus();
            } else {
                LOG.debug("Skipping: no testable resources found");
                return ExperimentTaskStatus.SUCCESS;
            }
        }
    }

    /**
     * The connectivity-test can begin if, for each sliver of this experimentPart if:
     * - the status of the sliver is READY
     * - the manifest of the sliver is available
     * <p/>
     * If the authority has one or more nodes, it must have at least one loginService available
     * (we assume that if one loginService is available, it will be available for all loginServices in that sliver)
     */
    public class FindTestableResourcesState extends State {

        private final ExperimentPart experimentPart;

        private Collection<FXRspecNode> nodesToTest = null;

        private FindTestableResourcesState(ExperimentPart experimentPart) {
            super(String.format("Determining testable resources on %s", experimentPart.getName()));
            this.experimentPart = experimentPart;
        }

        @Nonnull
        @Override
        protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {
            //use a map to make sure that only one rspecnode per id is collected
            Map<String, FXRspecNode> nodesOfAuthorityWithLoginService = new HashMap<>();
            Set<String> nodeIdsWithoutLoginService = new HashSet<>();

            if (!(experimentPart instanceof SfaExperimentPart)) {
                LOG.debug("Currently no support for connectivity testing of {}", experimentPart.getClass().getName());
                nodesToTest = Collections.emptyList();

                return ExperimentTaskStatus.WARNING;
            }

            SfaExperimentPart sfaExperimentPart = (SfaExperimentPart) experimentPart;

            assert sfaExperimentPart.getSlivers() != null;
            for (Sliver sliver : sfaExperimentPart.getSlivers()) {
                //skip fake slivers
                if(sliver.isFake())
                    continue;

                //check if sliver is ready
                if (sliver.getStatus().getGlobalStatus() != StatusDetails.SliverStatus.READY) {
                    if (LOG.isTraceEnabled())
                        LOG.trace("Cannot start connectivity test to {} as sliver {} is not READY", experimentPart.getName(),
                                sliver.getUrn());

                    return ExperimentTaskStatus.WARNING;
                }

                //check if manifest is available
                if (sliver.getManifestRspec() == null) {
                    if (LOG.isTraceEnabled())
                        LOG.trace("Cannot start connectivity test to {} as manifest is not available for sliver {}",
                                experimentPart.getName(), sliver.getUrn());

                    return ExperimentTaskStatus.WARNING;
                } else {

                    //only add node if it has a loginService, and hasn't been added yet to the map
                    FXModelRspec modelRspec = sliver.getManifestRspec().getModelRspec(ModelRspecType.FX);
                    assert modelRspec != null;
                    Collection<FXRspecNode> nodes = modelRspec.getNodesByAuthority(sfaExperimentPart.getConnectSfaAuthority());
                    assert nodes != null;
                    for (FXRspecNode node : nodes) {
                        if (!node.getLoginServices().isEmpty()) {
                            if (!nodesOfAuthorityWithLoginService.containsKey(node.getUniqueId())) {
                                nodesOfAuthorityWithLoginService.put(node.getUniqueId(), node);
                                nodeIdsWithoutLoginService.remove(node.getUniqueId()); //remove it from the nodeIdsWithoutLoginService
                            }
                        } else {
                            //check if we haven't found a loginService for this node before, and add it to nodeIdsWithoutLoginService otherwise
                            if (!nodesOfAuthorityWithLoginService.containsKey(node.getUniqueId()))
                                nodeIdsWithoutLoginService.add(node.getUniqueId());
                        }
                    }
                }
            }

            if (!nodesOfAuthorityWithLoginService.isEmpty()) {
                nodesToTest = nodesOfAuthorityWithLoginService.values();
            } else if (!nodeIdsWithoutLoginService.isEmpty()) {
                if (LOG.isTraceEnabled())
                    LOG.trace("Cannot start connectivity test for authority {} as no nodes with loginServices are available",
                            experimentPart.getName());
                nodesToTest = Collections.emptyList(); //we didn't find a loginService
            } else {
                //start the connectivity test
                if (LOG.isTraceEnabled())
                    LOG.trace("No connectivity test will be performed for authority {}, as it isn't responsible for any nodes",
                            experimentPart.getName());
                nodesToTest = Collections.emptyList();   //the authority isn't responsible for a node
            }

            return ExperimentTaskStatus.SUCCESS;
        }

        public Collection<FXRspecNode> getNodesToTest() {
            return nodesToTest;
        }
    }

    public class TestConnectivityToExperimentPartState extends State {

        private final ExperimentPart experimentPart;
        private final Collection<FXRspecNode> nodesToTest;

        private int currentTestCount = 0;
        private final Map<FXRspecNode, ConnectivityTest.Status> nodesConnectivity = new HashMap<>();

        private TestConnectivityToExperimentPartState(ExperimentPart experimentPart, Collection<FXRspecNode> nodesToTest) {
            super(String.format("Testing connectivity to resources on %s", experimentPart.getName()));
            this.experimentPart = experimentPart;
            this.nodesToTest = nodesToTest;
        }


        @Nonnull
        @Override
        protected ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException, JFedException {
            assert (!nodesToTest.isEmpty());
            assert experimentPart instanceof SfaExperimentPart;
            assert ((SfaExperimentPart) experimentPart).getSlivers() != null;

            for (Sliver sliver : ((SfaExperimentPart) experimentPart).getSlivers()) {
                assert (sliver.getStatus().getGlobalStatus() == StatusDetails.SliverStatus.READY);
                assert (sliver.isFake() || sliver.getManifestRspec() != null);
            }


            JFedConnection.ProxyInfo proxyInfo = null;
            if (Objects.equals(jFedPreferences.getString(CorePreferenceKey.PREF_SSHPROXY_USE_FOR_SSH), ProxyPreferencesManager.ALWAYS_USE_PROXY)) {
                GeniUser geniUser = geniUserProvider.getLoggedInGeniUser();
                assert geniUser.getUserAuthorityServer() != null;
                List<Proxy> userAuthProxyInfos = geniUser.getUserAuthorityServer().getTestbed().getProxies();
                proxyInfo =
                        proxyPreferencesManager.getSshProxySettings(userAuthProxyInfos, geniUser);

                if (proxyInfo == null) {
                    LOG.warn("Connectivity tests to nodes on authority {} are skipped as SSH-proxy is active but proxy info was not found." +
                            " (this is a jFed bug)", experimentPart.getName());


                    updateMessage("Skipped because proxy for SSH-connections is active but proxy info was not found. (this is a jFed bug)");
                    return ExperimentTaskStatus.FAILED;
                }
            }


            LOG.debug("Starting nodes connectivity test for {} nodes on {}", nodesToTest.size(), experimentPart.getName());

            Iterator<FXRspecNode> it = nodesToTest.iterator();
            while (it.hasNext()) {
                FXRspecNode node = it.next();

                if (node.getLoginServices().isEmpty()) {
                    LOG.debug("Removing node {} from test because it doesn't have any loginServices available", node.getUniqueId());
                    it.remove();
                }
            }

            int successCount = 0;

            while (currentTestCount < MAX_TESTS) {
                LOG.trace("Starting connectivity test round {} to {}", currentTestCount + 1, experimentPart.getName());

                updateMessage(String.format("Could connect to %d out of %d nodes. Testing... ",
                        successCount, nodesToTest.size()));
                for (FXRspecNode node : nodesToTest) {
                    if (!nodesConnectivity.containsKey(node)
                            || nodesConnectivity.get(node) != ConnectivityTest.Status.SUCCEEDED) {

                        ConnectivityTest.Status result = testConnectToNode(node, proxyInfo);
                        if (result == ConnectivityTest.Status.SUCCEEDED) {
                            successCount++;

                            updateMessage(String.format("Could connect to %d out of %d nodes. Testing... ",
                                    successCount, nodesToTest.size()));

                            LOG.trace("Could successfully connect to node {}@{} in test round {}",
                                    node.getClientId(), experimentPart.getName(), currentTestCount + 1);
                        } else {

                            LOG.trace("Got result {} when connecting to node {}@{} in test round {}",
                                    result, node.getClientId(), experimentPart.getName(), currentTestCount + 1);
                        }

                        nodesConnectivity.put(node, result);
                    }
                }

                //updateProgress(successCount, nodesToTest.size());

                if (successCount == nodesToTest.size()) {
                    updateMessage("");
                    return ExperimentTaskStatus.SUCCESS;
                } else {
                    LOG.trace("Could connect to {} of {} nodes in round {}", successCount, nodesToTest.size(), currentTestCount + 1);
                    updateMessage("Could connect to " + successCount + " of " + nodesToTest.size() + " requested nodes.");
                }

                //wait a few seconds before continuing test
                currentTestCount++;
                int waitLeft = WAIT_BETWEEN_TESTS;

                while (currentTestCount < MAX_TESTS && waitLeft > 0) {
                    updateMessage(String.format("Could connect to %d out of %d nodes. Trying again in %d seconds... ",
                            successCount, nodesToTest.size(), waitLeft));
                    try {
                        Thread.sleep(1000);
                    } catch (InterruptedException ignored) {
                    }
                    waitLeft--;
                }
            }

            return ExperimentTaskStatus.FAILED;
        }

        private ConnectivityTest.Status testConnectToNode(FXRspecNode node, JFedConnection.ProxyInfo proxyInfo) {
            LoginService loginService = node.getLoginServices().get(0);

            LOG.trace("Going to try to connect to {} by using loginService {}:{}",
                    node.getUniqueId(), loginService.getHostname(), loginService.getPort());


            JFedConnection.ProxyInfo testbedProxyInfo = proxyServiceUtil.findTestbedProxy(node, loginService);
            if (testbedProxyInfo != null)
                LOG.trace("Trying to make connection to {} with testbed proxy {}", node.getUniqueId(), testbedProxyInfo);

            if (proxyInfo != null && testbedProxyInfo != null) {
                LOG.warn("jFed has proxy currently enabled. However, connectivity test will only use testbed proxy {} while trying to connect to {}!", testbedProxyInfo, node.getUniqueId());
            }

            //try to setup a connection to the node
            HostAndPortTest hostAndPortTest = new HostAndPortTest(loginService.getHostname(),
                    Integer.parseInt(loginService.getPort()),
                    null, null, testbedProxyInfo == null ? proxyInfo : testbedProxyInfo);


            ConnectivityTest.ConnectivityTestResult result = hostAndPortTest.call();
            if (result.getStatus() != ConnectivityTest.Status.SUCCEEDED) {
                return result.getStatus();
            }

            return ConnectivityTest.Status.SUCCEEDED;
        }


    }
}
