/* 
 * Copyright (c) 2015 The Regents of the University of California.
 * All rights reserved.
 *
 * '$Author: crawl $'
 * '$Date: 2018-02-07 18:53:35 +0000 (Wed, 07 Feb 2018) $' 
 * '$Revision: 34661 $'
 * 
 * Permission is hereby granted, without written agreement and without
 * license or royalty fees, to use, copy, modify, and distribute this
 * software and its documentation for any purpose, provided that the above
 * copyright notice and the following two paragraphs appear in all copies
 * of this software.
 *
 * IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO ANY PARTY
 * FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES
 * ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF
 * THE UNIVERSITY OF CALIFORNIA HAS BEEN ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 * THE UNIVERSITY OF CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE SOFTWARE
 * PROVIDED HEREUNDER IS ON AN "AS IS" BASIS, AND THE UNIVERSITY OF
 * CALIFORNIA HAS NO OBLIGATION TO PROVIDE MAINTENANCE, SUPPORT, UPDATES,
 * ENHANCEMENTS, OR MODIFICATIONS.
 *
 */
package org.kepler.spark.actor;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.UnknownHostException;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.kepler.build.modules.Module;
import org.kepler.build.modules.ModuleTree;
import org.kepler.configuration.ConfigurationProperty;
import org.kepler.ddp.Utilities;
import org.kepler.spark.director.SparkEngine;

import ptolemy.data.IntToken;
import ptolemy.data.StringToken;
import ptolemy.data.expr.Parameter;
import ptolemy.data.expr.StringParameter;
import ptolemy.data.type.BaseType;
import ptolemy.kernel.util.Attribute;
import ptolemy.kernel.util.IllegalActionException;
import ptolemy.kernel.util.NameDuplicationException;
import ptolemy.kernel.util.NamedObj;
import ptolemy.util.MessageHandler;
import ptolemy.vergil.icon.BoxedValueIcon;

/** An attribute that contains connection information to a Spark server.
 *  The value specified in connectionName is used by Spark actors to
 *  refer to an instance of this attribute and use its Spark server.
 * 
 * @author Daniel Crawl
 * @version $Id: SparkConnection.java 34661 2018-02-07 18:53:35Z crawl $
 */
public class SparkConnection extends Attribute {

    /** Create a new SparkConnection with the specified name. */
    public SparkConnection(NamedObj container, String name)
            throws IllegalActionException, NameDuplicationException {
        super(container, name);
        
        connectionName = new StringParameter(this, "connectionName");
        connectionName.setToken("default");
        
        serverName = new StringParameter(this, "serverName");
        serverName.setToken("");
        
        numLocalWorkers = new Parameter(this, "numLocalWorkers");
        numLocalWorkers.setTypeEquals(BaseType.INT);
        numLocalWorkers.setToken(new IntToken(DEFAULT_NUM_LOCAL_WORKERS));
        
        driverMemory = new StringParameter(this, "driverMemory");
        driverMemory.setToken(DEFAULT_DRIVER_MEMORY);
        
        // set the icon to a BoxedValueIcon to display the profile name
        BoxedValueIcon icon = new BoxedValueIcon(this, "_icon");
        icon.displayWidth.setExpression("25");
        icon.attributeName.setExpression("connectionName");
        icon.boxColor.setExpression("{0.0, 1.0, 1.0, 1.0}");
    }

    @Override
    public void attributeChanged(Attribute attribute) throws IllegalActionException {
     
        if(attribute == numLocalWorkers) {
            int val = ((IntToken)numLocalWorkers.getToken()).intValue();
            synchronized(_contextLock) {
                if(val != _numLocalWorkers) {
                    _numLocalWorkers = val;
                }
            }
        } else if(attribute == driverMemory) {
            String driverMemoryStr = ((StringToken)driverMemory.getToken()).stringValue();
            synchronized(_contextLock) {
                if(!driverMemoryStr.equals(_driverMemory)) {
                    _driverMemory = driverMemoryStr;
                }
            }
        } else if(attribute == serverName) {
            String serverNameStr = ((StringToken)serverName.getToken()).stringValue();
            synchronized(_contextLock) {
                if(!serverNameStr.equals(_serverName)) {
                    _serverName = serverNameStr;
                }
            }
        } else {
            super.attributeChanged(attribute);
        }
    }
    
    /** Clone this object into the specified workspace. */
    /*
    @Override
    public Object clone(Workspace workspace) throws CloneNotSupportedException {
        SparkConnection newObject = (SparkConnection) super.clone(workspace);
        return newObject;
    }
    */
    
    /** Get the name used for this attribute. */
    public String getConnectionName() throws IllegalActionException {
        return ((StringToken)connectionName.getToken()).stringValue();
    }

    /** Get the context using the parameters of this SparkConnection
     *  (serverName, driverMemory, numLocalWorkers, etc.).
     */
    public JavaSparkContext getContext() throws IllegalActionException {
        synchronized(_contextLock) {
            return getContext(_serverName, _numLocalWorkers, _driverMemory,
                getConnectionName());
        }
    }
    
    /** Get the context for a specific master.
     *  @param masterNameAndPortStr The master host and port. If empty,
     *  then environment variables MASTER, SPARK_MASTER_IP, and SPARK_LOCAL_IP
     *  are checked for the master. If these are empty, then try to connect
     *  to a master running on localhost. If the connection fails, then a
     *  local master is used (runs in same JVM). 
     *  @param numLocalWorkersVal The number of workers to use for a
     *  local master.
     *  @param driverMemoryStr The amount of memory in the driver.
     *  @param connectionNameStr The connecting name string.
     */
    public static JavaSparkContext getContext(String masterNameAndPortStr,
        int numLocalWorkersVal, String driverMemoryStr,
        String connectionNameStr) throws IllegalActionException {
        
        synchronized(_contextLock) {
            
            // see if we already created the context and any of the
            // parameters are different.
            if(_context != null && (
                    _serverName == null ||
                    !_serverName.equals(masterNameAndPortStr) ||
                    _numLocalWorkers != numLocalWorkersVal ||
                    _driverMemory == null ||
                    !_driverMemory.equals(driverMemoryStr))) {
                System.out.println("Will recreate Spark context since parameter(s) have changed.");
                //System.out.println("CONTEXT COUNTER = " + _contextCounter);
                _context.stop();
                _context = null;
                if(_startedMySparkServer) {
                    _stopServerWithMySpark();
                }
            }
            
            if(_context == null) {
                
                String masterStr;
                // if the server was not specified, use the default.
                if(masterNameAndPortStr == null || 
                    masterNameAndPortStr.trim().isEmpty()) {
                    
                    // get the default address for the master
                    InetSocketAddress socketAddress = getDefaultMasterSocketAddress();
                    
                    // see if we can connect
                    boolean connected = false;
                    try(Socket socket = new Socket();) {
                        socket.connect(socketAddress, _CONNECT_TIMEOUT);
                        connected = true;
                    } catch(IOException e) {
                        //System.err.println("IOException connecting to " + socketAddress +
                            //": " + e.getMessage());
                        connected = false;
                    }

                    // if we couldn't connect, use local
                    if(connected) {
                        masterStr = "spark://" +
                                socketAddress.getHostName() + ":" +
                                socketAddress.getPort();
                    }
                    // if we are running on a cluster, start the standalone
                    // server using the myspark scripts.
                    else if(Utilities.isExecutingUnderResourceManager()) {
                        masterStr = _startServerWithMySpark();
                    } else {
                        masterStr = "local[" +
                            (numLocalWorkersVal < 1 ? "*" : numLocalWorkersVal) +
                            "]";
                    }
                } else if(masterNameAndPortStr.startsWith("local")) {
                    masterStr = "local[" +
                        (numLocalWorkersVal < 1 ? "*" : numLocalWorkersVal) +
                        "]";
                } else if(!masterNameAndPortStr.startsWith("spark://")) {
                    masterStr = "spark://" + masterNameAndPortStr;
                } else {
                    masterStr = masterNameAndPortStr;
                }
                
                if(driverMemoryStr.trim().isEmpty()) {
                    driverMemoryStr = DEFAULT_DRIVER_MEMORY;
                    System.err.println("WARNING: Spark driver memory not specified. " +
                        "Using default value of " +
                        DEFAULT_DRIVER_MEMORY);
                }
                                
                _context = _createContext(masterStr, driverMemoryStr,
                    connectionNameStr);
                
                _serverName = masterStr;
                _numLocalWorkers = numLocalWorkersVal;
                _driverMemory = driverMemoryStr;
            }

            if(_context != null) {
                _contextCounter++;
                //System.out.println("CONTEXT COUNTER incremented = " + _contextCounter);
                
                if(!_context.getConf().get("spark.master").startsWith("local")) {
                    // add the module jars, e.g., actors.jar, ptolemy.jar, etc.
                    final ModuleTree moduleTree = ModuleTree.instance();
                    for(Module module : moduleTree) {
                        
                        // Spark 2 throws an exception when more than one jar with
                        // the same name is added (even when in different directories).
                        // So we need to exclude the (empty) kepler-tasks.jar that is
                        // built for the kepler-tasks module.
                        if(!module.getName().equals("kepler-tasks")) {
                            final File moduleJar = module.getTargetJar();
                            
                            // add the module jar if it exists.
                            // some modules, e.g., outreach, do not have jars.
                            if(moduleJar.exists()) {
                                _context.addJar(moduleJar.getAbsolutePath());                
                            }
                        }
                    }
                }
            }
        }
        return _context;
    }

    /** Get the default context. */
    public static JavaSparkContext getDefaultContext() throws IllegalActionException {
        return getContext(_serverName, _numLocalWorkers,
            _driverMemory, "default");    
    }

    /** Get the default Spark master socket address. */
    public static InetSocketAddress getDefaultMasterSocketAddress() throws IllegalActionException {
        
        String masterHostAndPortStr = null;
        InetSocketAddress socketAddress;
        
        // first try environment variables: MASTER, SPARK_MASTER_IP, SPARK_LOCAL_IP
        String envMaster = System.getenv("MASTER");
        if(envMaster != null && envMaster.startsWith("spark://")) {
            masterHostAndPortStr = envMaster;
            System.out.println("Using $MASTER: " + masterHostAndPortStr);
            socketAddress = getMasterAndPortAddress(masterHostAndPortStr);
        } else {
            String hostname = null;
            String envSparkMasterIP = System.getenv("SPARK_MASTER_IP");
            if(envSparkMasterIP != null) {
                hostname = envSparkMasterIP;
                System.out.println("Using $SPARK_MASTER_IP: " + envSparkMasterIP);
            } else {
                String envSparkLocalIP = System.getenv("SPARK_LOCAL_IP");
                if(envSparkLocalIP != null) {
                    hostname = envSparkLocalIP;
                    System.out.println("Using $SPARK_LOCAL_IP: " + envSparkLocalIP);
                } else {
                    // finally, use localhost
                    try {
                        hostname = InetAddress.getLocalHost().getHostName();
                    } catch (UnknownHostException e) {
                        throw new IllegalActionException("Error determining host name:" +
                                e.getMessage());
                    }
                }
            }

            if(hostname == null || hostname.isEmpty()) {
                throw new IllegalActionException("Unable to determine host name.");
            }
                
            //System.out.println("Using default port for master: " + DEFAULT_MASTER_PORT);
            socketAddress = new InetSocketAddress(hostname, DEFAULT_MASTER_PORT);
        }     
        
        return socketAddress;
    }
    
    /** Parse a string containing the master host and port and return the
     *  socket address.
     */
    public static InetSocketAddress getMasterAndPortAddress(String masterHostAndPortStr) throws IllegalActionException {
        boolean isSparkURL = masterHostAndPortStr.startsWith("spark://");
        String hostname = null;
        int port;
        String[] parts = masterHostAndPortStr.split(":");
        if(parts.length != 2 &&
                (!isSparkURL || parts.length != 3)) {
            throw new IllegalActionException("masterHostAndPort must have both " +
                    " master host name and port.");
        }
        if(isSparkURL) {
            hostname = parts[1].substring("//".length());
            try {
                port = Integer.parseInt(parts[2]);
            } catch(NumberFormatException e) {
                throw new IllegalActionException("Unable to parse port: " + parts[1] +
                        ": " + e.getMessage());
            }
        } else {
            hostname = parts[0];
            try {
                port = Integer.valueOf(parts[1]).intValue();
            } catch(NumberFormatException e) {
                throw new IllegalActionException("Unable to parse port: " + parts[1] +
                        ": " + e.getMessage());
            }
        }
        //System.out.println("h = " + hostname + " p = " + port);
        return new InetSocketAddress(hostname, port);
    }

    /** Parse the output from starting the server using start-myspark.sh 
     *  @return The master URL if successful, otherwise null. 
     */
    public static String parseOutputFromStartingServer(InputStream input) throws IOException, IllegalActionException {
        String masterHostAndPortStr = null;
        try(BufferedReader reader = new BufferedReader(new InputStreamReader(input));) {
            String line = null;
            while ((line = reader.readLine()) != null) {
                if(line.startsWith("MASTER=spark://")) {
                    String[] parts = line.split("=");
                    masterHostAndPortStr = parts[1];
                    System.out.println("Spark master host and port is " + masterHostAndPortStr);
                }
            }
        }
        return masterHostAndPortStr;
    }
    
    public static void releaseContext() {
        synchronized(_contextLock) {
            if(_context != null) {
                _contextCounter--;
                //System.out.println("CONTEXT COUNTER decremented = " + _contextCounter);
                if(_contextCounter == 0) {
                    // if we are using a standalone server, close the connection
                    // since otherwise we can't submit new jobs.
                    //if(!_context.getConf().get("spark.master").startsWith("local")) {
                        _context.stop();
                    //}
                    _context = null;
                }
            }
        }
    }

    /** Stop any contexts that were started. */
    public static void shutdown() {        
        synchronized(_contextLock) {
            if(_context != null) {
                _context.stop();
                _context = null;
            }
            
            if(_startedMySparkServer) {
                try {
                    _stopServerWithMySpark();
                } catch (IllegalActionException e) {
                    MessageHandler.error("Error stopping Spark server.", e);
                }
            }
        }
    }
        
    /** The name of the connection. */
    public StringParameter connectionName;
    
    /** The host name of the Spark server. Use "local" to run on the local machine. */
    public StringParameter serverName;
    
    /** The number of workers when running Spark locally. */
    public Parameter numLocalWorkers;

    /** The amount of memory to use in the Spark driver. */
    public StringParameter driverMemory;
    
    ///////////////////////////////////////////////////////////////////
    ////                         private methods                   ////

    /** Create a context. */
    private static JavaSparkContext _createContext(final String masterAndPortStr,
            String driverMemoryStr, final String connectionNameStr)
                throws IllegalActionException {
        
        final JavaSparkContext[] contexts = new JavaSparkContext[1];
        
        // try to get the spark context.
        // use a ExecutorService with a timeout to handle
        // connecting to a master run by someone else.
        /*
        ExecutorService service = Executors.newSingleThreadExecutor();
        try {
            service.submit(new Runnable() {
                @Override
                public void run() {
        */          
                    //System.out.println("Getting JavaSparkContext for " +
                            //connectionNameStr + " at " + masterAndPortStr);                

                    SparkConf conf = new SparkConf()
                            .setMaster(masterAndPortStr)
                            .setAppName("Kepler Spark " + connectionNameStr)
                            .set("spark.driver.memory", DEFAULT_DRIVER_MEMORY);
                          
                    contexts[0] = new JavaSparkContext(conf);
                    
                    //contexts[0].sc().addSparkListener(new StdoutSparkListener());

                    // if we're not using the local server, we need to wait
                    // for the scheduler to set the default parallelism to
                    // be the total number of cores. see:
                    // https://issues.apache.org/jira/browse/SPARK-809
                    //
                    
                    if(!masterAndPortStr.startsWith("local")) {                        
                        // wait 5 seconds for the executor added events
                        System.out.println("Waiting 5 seconds to receive executor(s) information.");
                        try {
                            Thread.sleep(5000);
                        } catch (InterruptedException e) {
                            throw new IllegalActionException("Interrupted while waiting for " +
                                "Spark scheduler to initialize: " + e.getMessage());
                        }
                    }
       /*
                    System.out.println("finished calling jsp ctor");
                }
            }).get(_CONNECT_TIMEOUT, TimeUnit.MILLISECONDS);                    
        } catch (InterruptedException | ExecutionException | TimeoutException e) {
            System.err.println("WARNING: Timeout connecting to master at " +
                    masterAndPortStr + ".");
        } finally {
            service.shutdownNow();
        }
        */
        
        if(contexts[0] != null) {
            System.out.println("Created Spark context for: " + 
                contexts[0].getConf().get("spark.master"));
        }
        
        return contexts[0];
    }
    
    /** Start the standalone server using start-myspark.sh.
     *  @return Returns the master and port string.
     */
    private static String _startServerWithMySpark() throws IllegalActionException {
        
        // set the default location of the config directory
        String workflowDirStr = System.getProperty("spark.workflowdir");
        if(workflowDirStr == null) {
            throw new IllegalActionException("System property " +
                    "spark.workflowdir not set.");
        }
        
        String configDirStr = workflowDirStr + File.separator +
            "tools" + File.separator + "conf";
        
        File configDirFile = new File(configDirStr);
        String parentDirStr = configDirFile.getParent();
        
        final String pathStr = parentDirStr + File.separator + "sbin" +
                    File.separator;
        
        ConfigurationProperty engineProperty =
            Utilities.getEngineProperty(SparkEngine.SPARK_ENGINE_NAME, null);

        ConfigurationProperty startProperty = engineProperty.getProperty("Server.Scripts.ClusterStart");
        if(startProperty == null) {
            throw new IllegalActionException("Could not find ClusterStart script in configuration file.");
        }
        
        String startScriptStr = pathStr + startProperty.getValue();
        
        System.out.println("Starting Spark server with " + startScriptStr);

        // see if the script is executable. kepler modules are zipped,
        // which does not preserve the permissions.
        File startScriptFile = new File(startScriptStr);
        if(!startScriptFile.canExecute()) {
            throw new IllegalActionException( 
                    "The script " + startScriptFile + " is not executable.\n" +
                            "You must change the permissions so that " +
                            startScriptFile.getName() + 
                            " and all the other scripts in \n" +
                            startScriptFile.getParent() + " are executable.");
        }
        
        ProcessBuilder builder = new ProcessBuilder(startScriptStr);        
        builder.redirectErrorStream(true);

        String masterAndPortStr;
        
        try {
            Process process = builder.start();
            masterAndPortStr = parseOutputFromStartingServer(process.getInputStream());
            process.waitFor();
            _startedMySparkServer = true;
        } catch(Exception e) {
            throw new IllegalActionException("Error starting Spark server " +
                    "using myspark: " + e.getMessage());
        }
        
        
        return masterAndPortStr;
    }
    
    /** Stop the standalone server using stop-myspark.sh. */
    private static void _stopServerWithMySpark() throws IllegalActionException {
        
        // set the default location of the config directory
        String workflowDirStr = System.getProperty("spark.workflowdir");
        if(workflowDirStr == null) {
            throw new IllegalActionException("System property " +
                    "spark.workflowdir not set.");
        }
        
        String configDirStr = workflowDirStr + File.separator +
            "tools" + File.separator + "conf";
        
        File configDirFile = new File(configDirStr);
        String parentDirStr = configDirFile.getParent();
        
        final String pathStr = parentDirStr + File.separator + "sbin" +
                File.separator;
    
        ConfigurationProperty engineProperty =
            Utilities.getEngineProperty(SparkEngine.SPARK_ENGINE_NAME, null);
    
        ConfigurationProperty startProperty = engineProperty.getProperty("Server.Scripts.ClusterStop");
        if(startProperty == null) {
            throw new IllegalActionException("Could not find ClusterStop script in configuration file.");
        }
        
        String stopScriptStr = pathStr + startProperty.getValue();
        
        System.out.println("Stopping Spark server with " + stopScriptStr);

        // see if the script is executable. kepler modules are zipped,
        // which does not preserve the permissions.
        File startScriptFile = new File(stopScriptStr);
        if(!startScriptFile.canExecute()) {
            throw new IllegalActionException( 
                    "The script " + startScriptFile + " is not executable.\n" +
                            "You must change the permissions so that " +
                            startScriptFile.getName() + 
                            " and all the other scripts in \n" +
                            startScriptFile.getParent() + " are executable.");
        }
        
        ProcessBuilder builder = new ProcessBuilder(stopScriptStr);        
        builder.redirectErrorStream(true);

        try {
            Process process = builder.start();
            process.waitFor();
            _startedMySparkServer = false;
        } catch(Exception e) {
            throw new IllegalActionException("Error stopping Spark server " +
                    "using myspark: " + e.getMessage());
        }
        
    }

    ///////////////////////////////////////////////////////////////////
    ////                         private fields                    ////

    /** The spark context.
     *  NOTE: Spark only allows a single SparkContext per JVM,
     *  so this is static.
     */
    private static JavaSparkContext _context;
    
    /** Lock object for _context. */
    private static final Object _contextLock = new Object();
    
    /** The default number of local workers. */
    public static final int DEFAULT_NUM_LOCAL_WORKERS = 0;
    
    /** The number of workers for the local context. */
    private static int _numLocalWorkers = DEFAULT_NUM_LOCAL_WORKERS;
    
    /** The default amount of memory for the Spark driver. */
    public static final String DEFAULT_DRIVER_MEMORY = "1g";
    
    /** The amount of memory for the Spark driver. */
    private static String _driverMemory = DEFAULT_DRIVER_MEMORY;
    
    /** Default port for the master. */
    public final static int DEFAULT_MASTER_PORT = 7077;

    /** Timeout (in ms) when seeing if server is running. */
    private final static int _CONNECT_TIMEOUT = 10*1000;
    
    /** Default server name. */
    private static String _serverName = "";
    
    /** Reference counter for the context. */
    private static int _contextCounter = 0;

    /** If true, we've started standalone server using myspark. */
    private static boolean _startedMySparkServer = false;
    
}