#!/usr/bin/env python
from __future__ import print_function

from xbow import utilities
from xbow import instances
from xbow import filesystems
from xbow import pools

import xbow
import yaml
import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-s', '--script', help='name of provisioning script')
parser.add_argument('-n', '--n_workers', help='number of workers to launch')

args = parser.parse_args()

cfg_file = os.path.join(xbow.XBOW_CONFIGDIR, "settings.yml")

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

cfg['image_id'] = utilities.get_image_id(cfg)

schedulers = instances.get_by_name(cfg['scheduler_name'])
if len(schedulers) > 1:
    print('Error - there is more than one scheduler already running with this name.')
    exit(1)

fs_id = filesystems.fs_id_from_name(cfg['shared_file_system'], 
                                   region=cfg['region'] 
                                  )
if fs_id is None:
    fs_id = filesystems.create_fs(cfg['shared_file_system'],
                                 region=cfg['region'], 
                                 efs_security_groups=cfg['efs_security_groups']
                                )
cfg['fs_id'] = fs_id

preamble_data = '''Content-Type: multipart/mixed; boundary="//"
MIME-Version: 1.0

--//
Content-Type: text/cloud-config; charset="us-ascii"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Content-Disposition: attachment; filename="cloud-config.txt"

#cloud-config
cloud_final_modules:
- [scripts-user, always]

--//
Content-Type: text/x-shellscript; charset="us-ascii"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Content-Disposition: attachment; filename="userdata.txt"
'''
scheduler_user_data='''
#!/bin/bash
mkdir -p /run/metadata/xbow
touch /run/metadata/xbow/is_scheduler
echo 'XBOW_SHARED_FILESYSTEM={fs_id}.efs.{region}.amazonaws.com:/' > /run/metadata/xbow/shared_file_system
echo 'SHARED={mount_point}' >> /etc/environment
'''.format(**cfg)

scheduler_extra_data = ''
worker_extra_data = ''
if args.script:
    with open(args.script, 'r') as f:
        for_scheduler = True
        for_worker = True
        for line in f:
            if len(line) == 0:
                pass
            elif line[0] == '#':
                pass
            elif line[0] == '[':
                if '[all]' in line:
                    for_scheduler = True
                    for_worker = True
                elif [scheduler] in line:
                    for_scheduler = True
                    for_worker = False
                elif '[workers]' in line:
                    for_scheduler = False
                    for_worker = True
                else:
                    raise ValueError('Error - unrecognised command in provisioning file')
            else:
                if for_scheduler:
                    scheduler_extra_data += line
                if for_worker:
                    worker_extra_data += line

final_data = '''

--//'''
    
if len(schedulers) == 1:
    inst = schedulers[0]
    print('Scheduler already running, now starting workers. This may take some time...')
else:
    print("Starting the scheduler node - this may take some time...")

    user_data = preamble_data + scheduler_user_data + scheduler_extra_data + final_data
    inst = instances.create(
                            cfg['scheduler_name'],
                            image_id=cfg['image_id'],
                            instance_type=cfg['scheduler_instance_type'],
                            ec2_security_groups = cfg['ec2_security_groups'],
                            user_data=user_data
                      )
    print("Scheduler ready, now starting workers. This may take some time...")
    
cfg['scheduler_ip_address'] = inst.private_ip_address

worker_user_data = '''
#!/bin/bash
mkdir -p /run/metadata/xbow
echo 'XBOW_SCHEDULER_IP_ADDRESS={scheduler_ip_address}' > /run/metadata/xbow/scheduler_ip_address
echo 'XBOW_SHARED_FILESYSTEM={fs_id}.efs.{region}.amazonaws.com:/' > /run/metadata/xbow/shared_file_system
echo 'SHARED={mount_point}' >> /etc/environment
'''.format(**cfg)

if args.n_workers:
    n_workers = int(args.n_workers)
else:
    n_workers = cfg['pool_size']

user_data = preamble_data + worker_user_data + worker_extra_data + final_data
sip = pools.create_spot_pool(cfg['worker_pool_name'],
                        count=n_workers,
                        price=cfg['price'],
                        image_id=cfg['image_id'],
                        instance_type=cfg['worker_instance_type'],
                        ec2_security_groups=cfg['ec2_security_groups'],
                        user_data=user_data
                      )
print("Instance running, now waiting until xbow-cluster is accessible...")

ci = instances.ConnectedInstance(inst)
print("All ready for use")
