#!/usr/bin/env python

"""
Usage: prog <workflow name> <job id> <callback address> <spark services workflow script> <python location>

If the callback address does not start with an 'http' it is assumed to be a configuration path and no callback will be used

Actions:
    1.  Provides spark callback address to launching server (via server callback)
    2.  Initializes spark cluster
    3.  Runs job on server (blocking)
    4.  Examines outputs for errors and informs callback

Assumptions: The environment should be properly set by the driver.
"""

import os

##### CUSTOM CONFIG (global parameters -- change as needed!!) ####

USER = os.environ['USER']

# Location of spark distribution
SPARK_HOME = "/usr/local/spark-current"

# spark configuration path (disable default python)
CONF_DIR = "/groups/flyem/proj/cluster/miniconda/envs/flyem/DVIDSparkServices/conf"

# directory location of log file
LOG_DIR = os.path.expanduser('~') + "/sparkjoblogs"

# DVIDSparkServices will check this variable and use to override Python's tempfile.tempdir
DVIDSPARK_WORKFLOW_TMPDIR = "/scratch/" + USER

##### CUSTOM CONFIGURATION ###########
    
import os
import socket
import sys
import subprocess
import time
import requests
import tempfile
import json
from StringIO import StringIO

master_output = StringIO()
successful = False

json_header = {'content-type': 'app/json'}
start = time.time()

hascallback = True
if not sys.argv[3].startswith("http"):
    hascallback = False
    
try:
    # start workflow script location
    WORKFLOW_LAUNCHER = sys.argv[4] 
    
    # Add directories to PATH
    PATH_DIRS = SPARK_HOME + "/bin:" + SPARK_HOME + "/sbin"

    # set path
    curr_path = os.environ["PATH"]
    os.environ["PATH"] = PATH_DIRS + ":" + curr_path

    # set spark path
    os.environ["SPARK_HOME"] = SPARK_HOME

    # set configuration directory
    os.environ["SPARK_CONF_DIR"] = CONF_DIR

    # set exact python to be used
    os.environ["PYSPARK_PYTHON"] = sys.argv[5] 

    # DVIDSparkServices will check this variable and use to override Python's tempfile.tempdir
    os.environ["DVIDSPARK_WORKFLOW_TMPDIR"] = DVIDSPARK_WORKFLOW_TMPDIR

    # ******** Launch Spark Cluster *******
    HOSTNAME = socket.gethostname()

    master_output.write( subprocess.check_output([SPARK_HOME + '/sbin/start-master.sh']) + "\n" )
    
    # this sets default master -- does not need to be specified on commandline
    os.environ["MASTER"] = "spark://" + HOSTNAME + ":7077"

    master_output.write( os.environ["MASTER"] + "\n" )

    # ******** Start Job ********
    configfile = sys.argv[3]

    if hascallback:
        # write-back callback address
        status = {}
        status["sparkAddr"] = HOSTNAME
        status["job_status"] = "Running"
        status_str = json.dumps(status)

        requests.post(sys.argv[3], data=status_str, headers=json_header)

        configfile = configfile + "/config"

    # wait a few seconds for the cluster to stabilise (is this necessary?)
    time.sleep(10)

    # call workflow and wait
    try:
        job_output = subprocess.check_output([SPARK_HOME + '/bin/spark-submit', WORKFLOW_LAUNCHER, sys.argv[1], '-c', configfile])
        master_output.write( job_output + "\n" )
        successful = True
    finally:
        time.sleep(20)

except subprocess.CalledProcessError as ex:
    import traceback
    traceback.print_exc(file=master_output)
    master_output.write(ex.output)    
    master_output.write("Subprocess return code: {}".format(ex.returncode) )    
except:
    import traceback
    traceback.print_exc(file=master_output)
finally:
    # record time
    finish = time.time()
    master_output.write( "\n" + "Total Time: " + str(finish-start) + "\n" )
    
    # write status and message
    status = {}
    status["sparkAddr"] = ""
    if not successful:
        status["job_status"] = "Error"
    else:
        status["job_status"] = "Finished"
    status["job_message"] = master_output.getvalue()
    status_str = json.dumps(status)
    
    if hascallback:
        requests.post(sys.argv[3], data=status_str, headers=json_header)
    
    # write to file
    with open(LOG_DIR + "/" + sys.argv[2]+'.log', 'w') as fout:
        fout.write(master_output.getvalue() + "\n")
        fout.write( "Launch script done: {}\n".format( {True: "successful", False: "UNSUCCESSFUL"}[successful] ) )
