#!/usr/bin/env python
import os
import sys

try:
    import thunder
except ImportError as e:
    thunder = None
    print("Unable to import Thunder, this is likely due to an incorrect PYTHONPATH or missing dependencies."
          " Got the following error during import:\n%s" % e)
    sys.exit()
except SyntaxError as e:
    thunder = None
    print("Unable to import Thunder, got the following syntax error:\n%s" % e)
    sys.exit()
except BaseException as e:
    print("Unable to import Thunder for unknown reasons, got the following error:\n%s" % e)
    sys.exit()

from thunder.utils.launch import getFilteredHelpMessage, getSparkHome, transformArguments


def getUsage(wrappedScriptName='pyspark'):
    scriptName = os.path.basename(sys.argv[0])
    return "Usage: %s [-i | -n] [%s options]\n\n" % (sys.argv[0], wrappedScriptName) + \
           "The '%s' script is a wrapper around Spark's '%s', and accepts all the same options, " % (scriptName, wrappedScriptName) + \
           "although not all are meaningful when running python scripts." + \
           "\n\nThunder options:\n" + \
           "  -i : Start thunder via an IPython shell, equivalent to running with $IPYTHON=1\n" + \
           "  -n : Start thunder via an IPython Notebook, equivalent to running with $IPYTHON_OPTS=\"notebook\"\n\n" + \
           "Options for '%s' follow.\n" % wrappedScriptName


def main():
    SPARK_HOME = getSparkHome()

    childArgs = transformArguments(sys.argv)

    sparkSubmit = os.path.join(SPARK_HOME, 'bin', 'pyspark')

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

    # check for ipython option
    if "-i" in childArgs:
        os.environ['IPYTHON'] = '1'
        childArgs = [opt for opt in childArgs if "-i" not in opt]

    # check for notebook option
    if "-n" in childArgs:
        args = [opt for opt in childArgs if "--notebook-opts" in opt]
        if len(args) > 0:
            extraOpts = ' ' + args[0][16:]
        else:
            extraOpts = ''
        os.environ['IPYTHON_OPTS'] = 'notebook' + extraOpts
        del childArgs[childArgs.index("-n")]
        childArgs = [opt for opt in childArgs if "--notebook-opts" not in opt and "-n" not in opt]

    # check for help flags, and print our own first if present
    if "-h" in childArgs or "--help" in childArgs:
        print >> sys.stderr, getFilteredHelpMessage(sparkSubmit, getUsage())
    else:
        os.system(sparkSubmit + " " + " ".join(childArgs))

if __name__ == "__main__":
    main()
