"""Submit jobs to Sun Grid Engine."""
# pylint: disable=invalid-name
from __future__ import absolute_import

import os
import subprocess
from . import tracker

def submit(args):
    """Job submission script for SGE."""
    if args.jobname is None:
        args.jobname = ('dmlc%d.' % args.num_workers) + args.command[0].split('/')[-1]
    if args.sge_log_dir is None:
        args.sge_log_dir = args.jobname + '.log'

    if os.path.exists(args.sge_log_dir):
        if not os.path.isdir(args.sge_log_dir):
            raise RuntimeError('specified --sge-log-dir %s is not a dir' % args.sge_log_dir)
    else:
        os.mkdir(args.sge_log_dir)

    runscript = '%s/rundmlc.sh' % args.logdir
    fo = open(runscript, 'w')
    fo.write('source ~/.bashrc\n')
    fo.write('export DMLC_TASK_ID=${SGE_TASK_ID}\n')
    fo.write('export DMLC_JOB_CLUSTER=sge\n')
    fo.write('\"$@\"\n')
    fo.close()

    def sge_submit(nworker, nserver, pass_envs):
        """Internal submission function."""
        env_arg = ','.join('%s=\"%s\"' % (k, str(v)) for k, v in pass_envs.items())
        cmd = 'qsub -cwd -t 1-%d -S /bin/bash' % (nworker + nserver)
        if args.queue != 'default':
            cmd += '-q %s' % args.queue
        cmd += ' -N %s ' % args.jobname
        cmd += ' -e %s -o %s' % (args.logdir, args.logdir)
        cmd += ' -pe orte %d' % (args.vcores)
        cmd += ' -v %s,PATH=${PATH}:.' % env_arg
        cmd += ' %s %s' % (runscript, ' '.join(args.command))
        print(cmd)
        subprocess.check_call(cmd, shell=True)
        print('Waiting for the jobs to get up...')

    # call submit, with nslave, the commands to run each job and submit function
    tracker.submit(args.num_workers, args.num_servers,
                   fun_submit=sge_submit,
                   pscmd=' '.join(args.command))