#!/usr/bin/env python

import os
import glob
import sys

spark = os.getenv('SPARK_HOME')
if spark is None or spark == '':
    raise Exception('You must assign the environmental variable SPARK_HOME with the location of Spark')

def findPy4J(sparkpath):
    py4jglob = os.path.join(sparkpath, "python", "lib", "py4j-*-src.zip")
    try:
        return glob.iglob(py4jglob).next()
    except StopIteration:
        raise Exception("No py4j jar found at "+py4jglob+"; is SPARK_HOME set correctly?")

sys.path.append(spark + "python/")
sys.path.append(findPy4J(spark))

import thunder

if os.getenv('PYTHONPATH') is None:
    os.environ['PYTHONPATH'] = ""

os.environ['PYTHONPATH'] = os.environ['PYTHONPATH'] + ":" + spark + "/python:../"
os.environ['PYTHONPATH'] = os.environ['PYTHONPATH'] + ":" + findPy4J(spark)

# if not running locally, check whether egg exists, and if not, build it
master = os.getenv('MASTER')
ec2 = os.path.isfile('/root/spark-ec2/cluster-url')

# get calling directory
calldir = os.path.dirname(os.path.realpath(__file__))

# add python script
os.environ['PYTHONSTARTUP'] = os.path.join(os.path.dirname(os.path.realpath(thunder.__file__)), 'utils/shell.py')

if not ((master is None) or ("local" in master)) or ec2:

    # get directory
    dist = os.path.join(calldir, '../dist')

    # check for egg
    egg = glob.glob(os.path.join(dist, "thunder_python-" + str(thunder.__version__) + "*.egg"))

    # build if egg not found
    if len(egg) == 0:
        os.system(os.path.join(calldir, 'build'))
        egg = glob.glob(os.path.join(dist, "thunder_python-" + str(thunder.__version__) + "*.egg"))

    os.environ['ADD_FILES'] = egg[0]

# get location of jar and add to the submit classpath
jar = os.path.join(os.path.dirname(os.path.realpath(thunder.__file__)), 'lib/thunder_2.10-'+str(thunder.__version__)+'.jar')

if os.getenv('SPARK_SUBMIT_CLASSPATH') is None:
    os.environ['SPARK_SUBMIT_CLASSPATH'] = ""

os.environ['SPARK_SUBMIT_CLASSPATH'] = os.environ['SPARK_SUBMIT_CLASSPATH'] + ":" + jar

os.system(os.path.join(spark, 'bin/pyspark --jars ' + jar))