#! /usr/bin/env python3

# - Mitä on erä?
# - Riistaa kuten te.
# (Seen in a Tenavat cartoon that has not been found again.)

from argparse import ArgumentParser, REMAINDER
from itertools import chain
from shlex import quote
from subprocess import Popen, PIPE
from tempfile import mkstemp
import os, grp, sys

parser = ArgumentParser(description = '''

Send a single command to the batch system in Puhti, with all manner of
defaults, to be run in the current working directory, without waiting.

By default, module kieli is loaded quietly.

''')
parser.add_argument('command', help = "an executable command")
parser.add_argument('argument', metavar = '...', nargs = REMAINDER,
                    help = "its options and arguments")
parser.add_argument('--log', default = 'gamelog', metavar = 'dir',
                    help = '''
                    directory [./gamelog] where the standard input
                    and standard output of the batch job are written in
                    files named %%j-<job>.{out,err} where %%j is the
                    job number (log directory is created as needed)''')
parser.add_argument('--job', default = 'game', metavar = 'name',
                    help = '''a short name [game] for the job''')

group = parser.add_mutually_exclusive_group()
group.add_argument('--out', '-o', metavar = 'file',
                   help = '''
                   standard output from command on zero exit status;
                   on non-zero status leave standard output in a
                   sibling temporary file named file.<random>
                   ''')
group.add_argument('--accept', '-a', metavar = 'file',
                   help = '''
                   standard output from command regardless of
                   exit status
                   ''')

# time group - either specify hours or specify minutes (default one
# hour is probably good for the purpose)
group = parser.add_mutually_exclusive_group()
group.add_argument('--minutes', '-M', metavar = 'num',
                   help = 'minutes [60] to reserve')
group.add_argument('--hours', '-H', metavar = 'num',
                   help = 'hours [1] to reserve')

# memory group - man pages refer to megabytes and gigabytes but
# example at CSC says 4GB = 4096MB so they must mean 2^20 and 2^30
# rather than 10^6 and 10^9 https://en.wikipedia.org/wiki/Gibibyte
# https://research.csc.fi/taito-constructing-a-batch-job-file#3.1.2
# (defaults are tentative - less might do for the purpose?)
group = parser.add_mutually_exclusive_group()
group.add_argument('--MiB', metavar = 'num',
                   help = 'mebibytes [8192] to reserve')
group.add_argument('--GiB', metavar = 'num',
                   help = 'gibibytes [8] to reserve')

parser.add_argument('--cat', action = 'store_true',
                    help = '''
                    write the job description to standard output
                    instead of sending it to the batch queue
                    (for information only when using temp file)
                    ''')

if os.path.exists('/appl/soft/ling'):
    parser.add_argument('--cores', '-C',
                        choices = [
                            # Puhti nodes have 40 cores, trying to
                            # allocate all cores in the same node so that
                            # the communication between them is fast -
                            # TODO to allow larger multi-node numbers for
                            # heavier but core-savvy jobs like ffmpeg
                            '1', '2', '4', '5',
                            '10', '20', '40'
                        ],
                        default = '4', # not sure!
                        help = '''
                        how many Puhti cores to use,
                        all in one node [tentative default is 4]
                        (nodes have 40 cores)
                        ''')
else:
    parser.add_argument('--cores', '-C',
                        choices = [
                            # Taito nodes have 24 cores, trying to
                            # allocate all cores in the same node so that
                            # the communication between them is fast
                            '1', '2', '4', '8', '12', '24'
                        ],
                        default = '4', # not sure!
                        help = '''
                        how many Taito cores to use,
                        all in one node [tentative default is 4]
                        (nodes have 24 cores)
                        ''')

# default partition is set just before parsing the arguments;
# Taito defaults to "serial" as always,
# Puhti will default to "small" (1 node, up to 40 cores).
group = parser.add_mutually_exclusive_group()
group.add_argument('--test', dest = 'partition',
                   action = 'store_const', const = 'test',
                   help = '''
                   run in "test" partition
                   ''')
group.add_argument('--serial', dest = 'partition',
                   action = 'store_const', const = 'serial',
                   help = '''
                   run in "serial" partition (Taito default)
                   ''')
group.add_argument('--small', dest = 'partition',
                   action = 'store_const', const = 'small',
                   help = '''
                   run in "small" partition (Puhti default)
                   ''')
group.add_argument('--large', dest = 'partition',
                   action = 'store_const', const = 'large',
                   help = '''
                   run in "large" partition (Puhti)
                   ''')

# default billing group set below just before parsing
parser.add_argument('--bill', '-B', metavar = 'group',
                    help = '''
                    bill the project associated with the group
                    [defaults to clarin if user is in clarin]
                    ''')

parser.add_argument('--kieli',
                    choices = [
                        'yes',
                        'out',
                        'err',
                        'null',
                        'no'
                    ],
                    default = 'null',
                    help = '''
                    how to load modules kieli and biojava
                    (default is null, load quietly;
                    yes leaves its stdout and stderr as they are,
                    out redirects stderr to stdout,
                    err stdout to stderr,
                    null both to /dev/null, and
                    no means do not load)
                    ''')

# add options

groupnames = [ grp.getgrgid(k).gr_name for k in os.getgroups() ]
parser.set_defaults(partition = ('small' # guess this is Puhti
                                 if os.path.exists('/appl/soft/ling') else
                                 # guess this is Taito
                                 'serial'),
                    bill = ('clarin' if 'clarin' in groupnames else None))

args = parser.parse_args()
args.prog = parser.prog

if not os.path.exists('/appl/soft/ling'):
    # Taito does not seem to accept any value for --account
    # but works as usual when --account is not specified.
    print('{}: info: ignoring billing group "{}" in Taito'
          .format(args.prog, args.bill),
          file = sys.stderr)
elif args.bill in groupnames:
    print('{}: info: billing "{}" project'.format(args.prog, args.bill),
          file = sys.stderr)
elif args.bill:
    print('{}: error: billing group "{}" not in:'
          .format(args.prog, args.bill),
          *groupnames,
          file = sys.stderr)
    exit(1)
else:
    print('{}: error: no billing group'.format(args.prog),
          file = sys.stderr)
    exit(1)

def simplepath(path):
    '''Normalize path as relative to current directory if under the
    current directory and not starting with a component that starts
    with dots, else absolute.

    '''
    absolute = os.path.abspath(path)
    relative = os.path.relpath(path)
    return absolute if relative.startswith('..') else relative

def ensuredir(path):
    '''Attempt to establish that the directory path exists, by creating it
    if not. Do not bother to attempt to check whether anyone in
    particular can write in the directory.

    '''
    path = simplepath(path)
    try:
        os.makedirs(path, exist_ok = True)
    except OSError as exn:
        # in Python 3.4.0 (and earlier), a mode mismatch may lead here
        # even if path is all right
        if os.path.isdir(path): return
        raise exn

def ensuretempfile(outdir, outfile):
    handle, tmppath = mkstemp(prefix = outfile + '.', dir = outdir)
    os.close(handle)
    return tmppath

def setup(args):
    '''Make sure the relevant directories exist. Return suitable pathnames
    for the log dir, outfile (or None), and tempfile (or None). If
    there is to be a tempfile, create the tempfile in advance to claim
    the name. Except with the --cat option only output pathnames (a
    placeholder name for a tempfile) without creating any filesystem
    entries.

    '''

    args.cat or ensuredir(args.log)

    if args.out or args.accept:
        outfile = simplepath(args.out or args.accept)
        head, tail = os.path.split(outfile)
        if not tail:
            raise ArgumentError('game: not a filename: {}'
                                .format(args.out or
                                        args.accept))
    else:
        outfile = None

    if args.out and args.cat:
        # no good way out
        tempfile = os.path.join(head, tail + '.[random]')
    elif args.out:
        tempfile = ensuretempfile(head, tail)
    elif args.accept:
        tempfile = None
    else:
        tempfile = None

    return simplepath(args.log), outfile, tempfile

logdir, outfile, tempfile = setup(args)

# watch out! appending ' biojava' to the command when in Puhti,
# because otherwise Puhti batch jobs do not have java!
kieli = ({ 'yes' : 'module load kieli',
           'out' : '2>&1 module load kieli',
           'err' : '1>&2 module load kieli',
           'null' : '1> /dev/null 2>&1 module load kieli',
           'no' : '# module load kieli' }
         [args.kieli])

# appending ' biojava' to 'module load kieli' in Puhti
if os.path.exists('/appl/soft/ling'): kieli += ' biojava'

# Depending on the output options, one of:
# cmd arg ...
# cmd arg ... > outfile
# cmd arg ... > tempfile
# with appropriate quotation for the shell.
command = ' '.join(chain([quote(args.command)],
                         map(quote, args.argument),
                         ( []
                           if outfile is None else
                           ['>', quote(outfile)]
                           if tempfile is None else
                           ['>', quote(tempfile)] )))

logcommand = ' '.join(chain([quote(args.command)],
                            map(quote, args.argument)))

if tempfile is None:
    # also no outfile, or accept outfile as is
    finish = ''
else:
    # on success status, move output to outfile;
    # on error status, leave output in tempfile
    finish = '''\
test $status -eq 0 &&
mv {temp} {out}'''.format(temp = quote(tempfile),
                          out = quote(outfile))

# default a nour - and this could be done wholly in argparse? can an
# exclusive group have a common default?
time = ( '{:d}:{:02d}:00'
         .format(*(divmod(int(args.minutes), 60)
                   if args.minutes
                   else (int(args.hours), 0)
                   if args.hours else (1, 0))) )

# setting memory default here, not checking that argument is even a
# number let alone a natural number and not something outrageous,
# feeling bad about not checking TODO check and the whole processing
# could be moved to argparse machinery
memory = ( '{}M'.format(args.MiB) if args.MiB else
           '{}G'.format(args.GiB) if args.GiB else
           '8G' )

script = '''\
#! /bin/bash
#SBATCH --job-name={job}
#SBATCH --account={bill}
#SBATCH --partition={partition}
#SBATCH --nodes={nodes}
#SBATCH --ntasks={cores}
#SBATCH --time={time}
#SBATCH --mem={memory}
#SBATCH --out={out}
#SBATCH --error={err}
#SBATCH --chdir={workdir}

echo command: {logcommand}
echo outfile: {outfile}
echo workdir: {workdir}
echo partition: {partition}
echo nodes: {nodes}
echo cores: {cores}
echo time: {time}
echo memory: {memory}
echo load kieli: {whetherkieli}
echo billing group: {bill}
echo
date "+%F %T START"

{kieli}
{command}

status=$?
{finish}

T=$SECONDS
printf -v time %d:%02d:%02d $((T/3600)) $((T%3600/60)) $((T%60))
date "+%F %T FINISH IN $time WITH STATUS $status"
'''.format(job = args.job,
           bill = args.bill,
           partition = args.partition,
           nodes = '1',
           cores = args.cores,
           time = time,
           memory = memory,
           out = quote(os.path.join(args.log, '%j-{}.out'.format(args.job))),
           err = quote(os.path.join(args.log, '%j-{}.err'.format(args.job))),
           workdir = quote(os.getcwd()),
           logcommand = quote(logcommand),
           outfile = quote(outfile or '(stdout)'),
           kieli = kieli,
           whetherkieli = args.kieli,
           command = command,
           finish = finish)

if not os.path.exists('/appl/soft/ling'):
    # Taito does not seem to accept any value for --account,
    # but works as usual when --account is not specified.
    script = '\n'.join((line for line in script.split('\n')
                        if not line.startswith('#SBATCH --account=')))

with Popen(['cat' if args.cat else 'sbatch'], stdin = PIPE) as process:
    process.stdin.write(script.encode('UTF-8'))
