#!/Library/Frameworks/Python.framework/Versions/3.6/bin/python3
from __future__ import print_function

import uuid
import glob
import os, yaml
import sys
import time
import xbow
from xbow.metering import SpotMeter
from xbow.instances import get_by_name, ConnectedInstance

def pack_and_run_remote(command):
    """
    Stage the contents of the current directory to the xbow cluster, then
    run the given command, then stage output files back, then clean up.
    """
    cfg_file = os.path.join(xbow.XBOW_CONFIGDIR, "settings.yml")

    with open(cfg_file, 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    instances = get_by_name(cfg['scheduler_name'])
    if len(instances) == 0:
        raise ValueError('Error - no such instance')
    elif len(instances) > 1:
        raise ValueError('Error - more than one instance has that name')
    instance = instances[0]
    ci = ConnectedInstance(instance)
    az = instance.placement['AvailabilityZone']
    meter = SpotMeter(cfg['worker_instance_type'], az, count = cfg['pool_size'])
    jobid = uuid.uuid4()
    mount_point=cfg['mount_point']

    ci.exec_command('sudo apt install task-spooler && tsp -S 20')
    print('remote directory will be {}/{}'.format(mount_point, jobid))
    ci.exec_command('mkdir {}/{}'.format(mount_point, jobid))
    infiles = glob.glob('*')
    print('uploading files:')
    for filename in infiles:
        if os.path.isfile(filename):
            print('    {}'.format(filename))
            ci.upload(filename, '{}/{}/{}'.format(mount_point, jobid, filename))
    ci.exec_command("cd {}/{} && tsp xflow-exec '{}'".format(mount_point, jobid, command)) 
    tsp_id = ci.output[:-1]
    print('tsp job {} submitted.'.format(tsp_id))

    HOSTFILE = open(os.path.expanduser('.xbow_ids.yml'), 'w+')
    HOSTFILE.write('jobid: {}\n'.format(jobid))
    HOSTFILE.write('tsp_id: {}\n'.format(tsp_id))
    HOSTFILE.close()

if __name__ == '__main__':
    command = ' '.join(sys.argv[1:])
    pack_and_run_remote(command)
