#!/usr/bin/env python

# Copyright 2015 Michael Broxton
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

##
# This script sets up a Spark cluster on Google Compute Engine
#
# This script originally was written by @SigmoidAnalytics, and then later
# modified by @broxtronix.

import os, sys

# Import the package from the local script directory
print os.getcwd()
sys.path.append(os.getcwd())

try:
    import spark_gce
except ImportError as e:
    print("Unable to import the spark_gce python package, this is likely due to an incorrect PYTHONPATH or missing dependencies.\nGot the following error during import:\n%s" % e)
    sys.exit()
except SyntaxError as e:
    print("Unable to import the spark_gce python package, got the following syntax error:\n%s" % e)
    sys.exit()
except Exception as e:
    print("Unable to import the spark_gce python package for unknown reasons, got the following error:\n%s" % e)
    sys.exit()

if __name__ == "__main__":

    print "Spark for Google Compute Engine v" + spark_gce.__version__
    print ""

    # Read the arguments
    from spark_gce import parse_args
    (opts, action, cluster_name, optional_arg) = parse_args()

    # Make sure gcloud is accessible.
    from spark_gce import check_gcloud
    check_gcloud(cluster_name, opts)

    # Launch the cluster
    if action == "launch":
        from spark_gce import launch_cluster
        launch_cluster(cluster_name, opts)

    elif action == "start":
        from spark_gce import start_cluster
        start_cluster(cluster_name, opts)

    elif action == "stop":
        from spark_gce import stop_cluster
        stop_cluster(cluster_name, opts)

    elif action == "stop-slaves":
        print "\n\nStopping slave nodes, but leaving the master node running. NOTE: you must run 'spark-gce stop' before restarting the cluster!\n\n"
        from spark_gce import stop_cluster
        stop_cluster(cluster_name, opts, slaves_only = True)

    elif action == "destroy":
        from spark_gce import destroy_cluster
        destroy_cluster(cluster_name, opts)

    elif action == "login" or action == "ssh":
        from spark_gce import ssh_cluster
        ssh_cluster(cluster_name, opts)

    elif action == "mosh":
        from spark_gce import mosh_cluster
        mosh_cluster(cluster_name, opts)

    elif action == "sshfs":
        from spark_gce import sshfs_cluster
        sshfs_cluster(cluster_name, opts, optional_arg)

    else:
        print >> stderr, "Invalid action: %s" % action
        sys.exit(1)
