bundles/papyrus/org.eclipse.papyrus.ease.lang.python.jupyter/src/org/eclipse/papyrus/ease/lang/python/jupyter/internal/JupyterProxy.java

/*****************************************************************************
 * Copyright (c) 2019 CEA LIST and others.
 * 
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License 2.0
 * which accompanies this distribution, and is available at
 * https://www.eclipse.org/legal/epl-2.0/
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 * Contributors:
 *   CEA LIST - Initial API and implementation
 *   
 *****************************************************************************/
package org.eclipse.papyrus.ease.lang.python.jupyter.internal;

import java.awt.Desktop;
import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.MalformedURLException;
import java.net.ServerSocket;
import java.net.SocketAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import org.eclipse.core.resources.IFile;
import org.eclipse.core.runtime.CoreException;
import org.eclipse.core.runtime.FileLocator;
import org.eclipse.core.runtime.Platform;
import org.eclipse.core.runtime.URIUtil;
import org.eclipse.core.variables.IStringVariableManager;
import org.eclipse.core.variables.VariablesPlugin;
import org.eclipse.ease.ScriptEngineException;
import org.eclipse.ease.lang.python.py4j.internal.IPythonSideEngine;
import org.eclipse.papyrus.ease.lang.python.jupyter.Activator;
import org.eclipse.swt.widgets.Display;
import org.eclipse.ui.PartInitException;
import org.eclipse.ui.PlatformUI;
import org.eclipse.ui.browser.IWebBrowser;

import py4j.ClientServer;
import py4j.ClientServer.ClientServerBuilder;
import py4j.JavaServer;

public class JupyterProxy {

	private static final int PYTHON_SHUTDOWN_TIMEOUT_SECONDS = 10;
	private static final int PYTHON_STARTUP_TIMEOUT_SECONDS = 60;

	/**
	 * Path within this plug-in to the main python file.
	 */
	private static final String PYSRC_EASE_PY4J_MAIN_DIR = "/pysrc/";

	/**
	 * The ID of the py4j sources plug-in, needs to match the name of the dependent
	 * plug-in.
	 */
	private static final String PY4J_PYTHON_BUNDLE_ID = "py4j-python";

	private static final String EASE_PY4J_BUNDLE_ID = "org.eclipse.ease.lang.python.py4j";

	private static final String JUPYTER_PATH = "JUPYTER_PATH";


	private static final String PYTHON_PATH = "PYTHONPATH";

	private static final String EASE_PYTHON_COMMON_BUNDLE_ID = "org.eclipse.ease.lang.python";

	private static final String PY4J_PATH = "PY4J_PATH";

	private static final String EASE_PYTHON_COMMON_SRC = "EASE_PYTHON_COMMON_SRC";

	private static final String EASE_PY4J_SRC = "EASE_PY4J_SRC";
	private static final String JUPYTER_TOKEN = "JUPYTER_TOKEN";
	private static final String EASE_JAVA_PORT = "EASE_JAVA_PORT";

	private int notebookPort = -1;
	private String notebookToken = UUID.randomUUID().toString();
	private String notebookDirectory = System.getProperty("user.home");

	private boolean startBrowser = false;
	private boolean debug = false;

	private CountDownLatch pythonStartupComplete;
	private ClientServer gatewayServer;
	private IPythonSideEngine pythonSideEngine;
	private Process jupyterProcess;
	private String notebookName;

	private String pythonInterpreter = getPythonInterpreterExec();
	
	
	private List<Integer> cellToRun= new ArrayList<>();

	public JupyterProxy() {

	}

	public int getNotebookPort() {
		return notebookPort;
	}

	public JupyterProxy setNotebookPort(int notebookPort) {
		this.notebookPort = notebookPort;
		return this;
	}

	public String getNotebookToken() {
		return notebookToken;
	}

	public JupyterProxy setNotebookToken(String notebookToken) {
		this.notebookToken = notebookToken;
		return this;
	}

	public String getNotebookDirectory() {
		return notebookDirectory;
	}

	public JupyterProxy setNotebookDirectory(String notebookDirectory) {
		this.notebookDirectory = notebookDirectory;
		return this;
	}

	public JupyterProxy setStartBrowser(boolean startBrowser) {
		this.startBrowser = startBrowser;
		return this;
	}

	public ClientServer getGatewayServer() {
		return gatewayServer;
	}

	public boolean isDebug() {
		return debug;
	}

	public JupyterProxy setDebug(boolean debug) {
		this.debug = debug;
		return this;
	}

	public IPythonSideEngine getPythonSideEngine() {
		return pythonSideEngine;
	}

	public Process getJupyterProcess() throws MalformedURLException, IOException, URISyntaxException, CoreException,
			InterruptedException, ScriptEngineException {
		if (jupyterProcess == null) {
			startJupyterProcess();
		}
		return jupyterProcess;
	}

	public void startJupyterProcess()
			throws InterruptedException, ScriptEngineException, IOException, URISyntaxException {

		notebookPort = getFreePort(8888);

		gatewayServer = new ClientServerBuilder(this).javaPort(0).pythonPort(0).build();
		gatewayServer.startServer(true);
		final int javaListeningPort = ((JavaServer) gatewayServer.getJavaServer()).getListeningPort();

		final ProcessBuilder pb = new ProcessBuilder();

		// we use Py4J python interpreter, should include Jupyter
		pb.command().add(pythonInterpreter);
		pb.command().add("-m");
		pb.command().add("notebook");
		pb.command().add("--notebook-dir=" + notebookDirectory);
		String fullPath = notebookDirectory + File.separator + "static";
		fullPath = fullPath.replace("\\", "\\\\");
		pb.command().add("--NotebookApp.extra_static_paths=['" + fullPath + "']");
		if (notebookPort > -1) {
			pb.command().add("--port=" + notebookPort);
		}
		if (!startBrowser) {
			pb.command().add("--no-browser");
		}

		if (debug) {
			pb.command().add("--debug");
		}

		// JUPYTER_PATH is modified to add our new kernel description
		// see http://jupyter-client.readthedocs.io/en/stable/kernels.html#kernel-specs
		pb.environment().put(JUPYTER_PATH,
				appendToCurrentEnvPath(JUPYTER_PATH, getJupyterPathDirectory().getAbsolutePath()));

		
	

		// we also add the kernel directory in PYTHON_PATH in order to allow the
		// first-level python to load our custom python kernel
		pb.environment().put(PYTHON_PATH,
				appendToCurrentEnvPath(PYTHON_PATH, getJupyterKernelDirectory().getAbsolutePath()));

		// we now add several environment variable that will be read during our easepy4j
		// kernel startup to establish connection
		// with Py4J java side initiated by EASE
		pb.environment().put(EASE_JAVA_PORT, String.valueOf(javaListeningPort));

		pb.environment().put(JUPYTER_TOKEN, notebookToken);

		pb.environment().put(PY4J_PATH, getPy4jPythonSrc().getAbsolutePath());
		pb.environment().put(EASE_PYTHON_COMMON_SRC, getEASEPythonCommonSrc().getAbsolutePath());
		pb.environment().put(EASE_PY4J_SRC, getPy4jEaseMainDir().getAbsolutePath());

		jupyterProcess = pb.start();

		pythonStartupComplete = new CountDownLatch(1);

	}

	

	private String getPythonInterpreterExec() {
		String interpreter = Platform.getPreferencesService().getString("org.eclipse.ease.lang.python.py4j",
				"org.eclipse.ease.lang.python.py4j.INTERPRETER", "python", null);
		final IStringVariableManager variableManager = VariablesPlugin.getDefault().getStringVariableManager();
		try {
			interpreter = variableManager.performStringSubstitution(interpreter);
		} catch (CoreException e) {
			interpreter = "python";
		}
		return interpreter;
	}

	private void waitForJupyterStartup() throws InterruptedException, ScriptEngineException, IOException {
		boolean scanning = true;
		SocketAddress address = new InetSocketAddress("localhost", notebookPort);

		int attempts = 0;
		while (scanning && attempts < 200) {
			try {
				attempts++;
				SocketChannel.open(address);
				scanning = false;
			} catch (IOException e) {
				Thread.sleep(100);// 0.1 second
			}
		}
		if (scanning) {
			throw new ScriptEngineException("Jupyter notebook did not start within 20 seconds");
		}
	}

	public void waitForKernelStartup() throws InterruptedException, ScriptEngineException, IOException {

		waitForJupyterStartup();

		// web browser will start the kernel if a notebook has been selected to be run
		// as EASE script
		// else user will have to manually start the kernel in the browser
		openWebBrowser();

		// We wait here for the main notebook python process initialization
		if (!pythonStartupComplete.await(PYTHON_STARTUP_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
			throw new ScriptEngineException(
					"Kernel process did not start within " + PYTHON_STARTUP_TIMEOUT_SECONDS + " seconds");
		}

	}

	public void pythonStartupComplete(final int pythonPort, final IPythonSideEngine pythonSideEngine) {
		final JavaServer javaServer = (JavaServer) gatewayServer.getJavaServer();
		javaServer.resetCallbackClient(javaServer.getCallbackClient().getAddress(), pythonPort);

		this.pythonSideEngine = pythonSideEngine;

		pythonStartupComplete.countDown();

	}
	
	public List<Integer> getCellsToRun() {
		return cellToRun;
	}
	
	public void addAutoCells(List<Integer> cellsToRun) {
		this.cellToRun.addAll(cellsToRun);
		
	}

	private void openWebBrowser() {
		if (PlatformUI.isWorkbenchRunning()) {
			Display.getDefault().asyncExec(new Runnable() {

				@Override
				public void run() {

					try {
						// even if we wait for python startup, still have to wait for server startup to
						// avoid to have to refresh the page
						TimeUnit.SECONDS.sleep(3);
						IWebBrowser browser = PlatformUI.getWorkbench().getBrowserSupport().createBrowser(null);
						browser.openURL(getNoteBookURL());

					} catch (MalformedURLException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();

					} catch (InterruptedException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					} catch (PartInitException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
				}

			});

		} else {
			if (Desktop.isDesktopSupported() && Desktop.getDesktop().isSupported(Desktop.Action.BROWSE)) {

				try {
					URI uri = getNoteBookURL().toURI();
					Desktop.getDesktop().browse(uri);
				} catch (IOException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (URISyntaxException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
		}

	}

	private URL getNoteBookURL() throws MalformedURLException {
		String urlString = "http://localhost:" + notebookPort;
		if (notebookName != null) {
			urlString += "/notebooks/" + notebookName;
		}
		urlString += "?token=" + notebookToken;

		return new URL(urlString);
	}

	private File getEASEPythonCommonSrc() throws IOException {

		final File easePythonCommonBundleFile = FileLocator
				.getBundleFile(Platform.getBundle(EASE_PYTHON_COMMON_BUNDLE_ID));
		final File easePythonCommonSrc = new File(easePythonCommonBundleFile, "/pysrc");

		if (!easePythonCommonSrc.exists() || !easePythonCommonSrc.isDirectory()) {
			throw new IOException(
					"Failed to find EASE commong python directory, expected it here: " + easePythonCommonSrc);
		}
		return easePythonCommonSrc;
	}

	private File getPy4jPythonSrc() throws IOException {
		final File py4jPythonBundleFile = FileLocator.getBundleFile(Platform.getBundle(PY4J_PYTHON_BUNDLE_ID));
		final File py4jPythonSrc = new File(py4jPythonBundleFile, "/src");
		// final File py4j = new File(py4jPythonSrc, "py4j");
		if (!py4jPythonSrc.exists() || !py4jPythonSrc.isDirectory()) {
			throw new IOException("Failed to find py4j python directory, expected it here: " + py4jPythonSrc);
		}
		return py4jPythonSrc;
	}

	private File getPy4jEaseMainDir() throws MalformedURLException, IOException, URISyntaxException {
		final URL url = new URL("platform:/plugin/" + EASE_PY4J_BUNDLE_ID + PYSRC_EASE_PY4J_MAIN_DIR);
		final URL fileURL = FileLocator.toFileURL(url);
		final File py4jEaseMain = new File(URIUtil.toURI(fileURL).normalize());
		if (!py4jEaseMain.exists() && !py4jEaseMain.isDirectory()) {
			throw new IOException("Failed to find " + PYSRC_EASE_PY4J_MAIN_DIR + ", expected it here: " + py4jEaseMain);
		}
		return py4jEaseMain;
	}

	private File getJupyterPathDirectory() throws IOException, URISyntaxException {
		final URL url = new URL("platform:/plugin/" + Activator.PLUGIN_ID + "/jupyter");
		final URL fileURL = FileLocator.toFileURL(url);
		final File rootDir = new File(URIUtil.toURI(fileURL).normalize());
		if (!rootDir.exists() || !rootDir.isDirectory()) {
			throw new IOException(
					"Failed to find ease-py4j Jupyter kernel root directory, expected it here: " + rootDir);
		}
		return rootDir;
	}

	private File getJupyterKernelDirectory() throws IOException, URISyntaxException {
		final File jupyterPath = getJupyterPathDirectory();
		final File kernelDirectory = new File(jupyterPath, "/kernels/ease_py4j_kernel");
		if (!kernelDirectory.exists() || !kernelDirectory.isDirectory()) {
			throw new IOException(
					"Failed to find ease-py4j Jupyter kernel  directory, expected it here: " + kernelDirectory);
		}
		return kernelDirectory;
	}

	private String appendToCurrentEnvPath(String variableName, String value) {
		String currentPath = System.getenv(variableName);

		if (currentPath == null) {
			currentPath = value;
		} else {
			currentPath = value + File.pathSeparator + currentPath;
		}
		return currentPath;
	}

	/**
	 * Tries to find a new free port close to the given start port.
	 * <p>
	 * Not that there is a (unlikely but) possible race condition if the returned
	 * port is being used between calling this and actually starting a server.
	 * Handle blocked servers appropriately.
	 * 
	 * @param startPort Lowest desired port number.
	 * @return Free port (hopefully) close to startPort.
	 */
	public static int getFreePort(int startPort) {
		ServerSocket socket = null;
		while (true) {
			try {
				// Try to create server for port number
				socket = new ServerSocket(startPort);

				// Make reusable to timeout state problems
				socket.setReuseAddress(true);

				// No exception -> port is free
				return startPort;
			} catch (IOException ignore) {
				// IOException most (!) likely because of used port
				startPort++;
			} finally {
				if (socket != null) {
					try {
						socket.close();
					} catch (IOException ignore) {
						// Should not occur
					}
				}
			}
		}
	}

	public JupyterProxy setNotebook(IFile scriptFile) {
		notebookDirectory = scriptFile.getParent().getLocation().toOSString();
		notebookName = scriptFile.getName();
		return this;
	}

	public void shutDown() {

		ProcessBuilder stopProcess = new ProcessBuilder().command(pythonInterpreter, "-m", "notebook", "stop",
				Integer.toString(notebookPort));

		try {
			Process process = stopProcess.start();
			process.waitFor(PYTHON_SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS);

		} catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		} catch (InterruptedException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

		if (gatewayServer != null) {
			gatewayServer.shutdown();
		}

	}



}