#!/usr/bin/env python


from mycluster import mycluster
from mycluster.persist import JobDB
import argparse


def create_db():
    job_db = None
    try:
        job_db = JobDB()
    except Exception as e:
        print(('Database failed to initialise. Error Message: ' + str(e)))
        job_db = None

    return job_db


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-q', '--queues',
                        help='List all available queues', action="store_true")
    create_group = parser.add_argument_group(title='Create job file')
    create_group.add_argument('-c', '--create', help='Name of job file')
    create_group.add_argument('--jobqueue', help='Job queue (required)')
    create_group.add_argument('--jobname', default='myclusterjob',
                              help='Job name')
    create_group.add_argument('--ntasks', type=int, default=1,
                              help='Number of tasks')
    create_group.add_argument('--nthreadspertask', type=int, default=-1,
                              help='Number of threads per task')
    create_group.add_argument('--taskpernode', type=int, help='Task per node')
    create_group.add_argument('--script', help='script file (required)')
    create_group.add_argument('--project', default='default',
                              help='Project/Account used of accounting usage(required)')
    create_group.add_argument('--maxtime', default='12:00:00',
                              help='Maximum runtime (hh:mm::ss)')
    create_group.add_argument('--shared', default=False, action='store_true',
                              help='Non exclusive node use')
    create_group.add_argument('--ompiargs', default='-bysocket -bind-to-socket', help='Arguments for OpenMPI')
    create_group.add_argument('--qos', help='QOS level to submit to')
    parser.add_argument('-s', '--submit', help='Submit job file')
    parser.add_argument('-i', '--immediate', action='store_true',
                        help='Submit job immediately')
    parser.add_argument('--depends', default=None,
                        help='List of job dependencies jobA_id:jobB_id:jobC_id (slurm only)')
    parser.add_argument('-d', '--delete', help='Delete job')
    parser.add_argument('--email', help='Email update')
    parser.add_argument('--firstname', help='Name update')
    parser.add_argument('--lastname', help='Name update')
    parser.add_argument('-p', '--printjobs', type=int, nargs='?',
                        const=0,
                        help='Print jobs 0 = all, 10 = first ten, -10 = last ten')
    parser.add_argument('-e', '--export', type=int, help='Export job data')
    parser.add_argument('--remote', help='Add remote site - user@site')
    parser.add_argument('--nodb', default=False, action="store_true", help="Don't use a database to store job status")
    parser.add_argument('--sysscribe', type=int, help=argparse.SUPPRESS)
    parser.add_argument('--jobid', type=int, help=argparse.SUPPRESS)
    parser.add_argument('--appname', help=argparse.SUPPRESS)
    parser.add_argument('--appdata', type=int, help=argparse.SUPPRESS)
    parser.add_argument('--silent', default=False, action='store_true',
                        help='Non exclusive node use')
    args = parser.parse_args()

    job_db = None
    if not args.nodb:
        job_db = create_db()
    if job_db is not None:
        mycluster.update_db(job_db=job_db)

    """
    Module initialiser function
    """
    mycluster.init(args.silent, job_db=job_db)

    if args.queues:
        mycluster.print_queue_info()
    if args.submit:
        mycluster.submit(args.submit, args.immediate, args.depends, job_db=job_db)
    if args.sysscribe:
        mycluster.sysscribe_update(args.sysscribe, job_db=job_db)
    if args.jobid:
        if args.appname:
            mycluster.appname_update(args.jobid, args.appname, job_db)
        if args.appdata:
            mycluster.appdata_update(args.jobid, args.appdata, job_db)
        if not args.appname and not args.appdata:
            parser.print_help()
    if args.export:
        mycluster.export(args.export)
    if args.email:
        mycluster.email_update(args.email, job_db)
    if args.firstname:
        mycluster.firstname_update(args.firstname, job_db)
    if args.lastname:
        mycluster.lastname_update(args.lastname, job_db)
    if args.printjobs is not None:
        mycluster.printjobs(args.printjobs, job_db)
    if args.delete:
        mycluster.delete(args.delete, job_db)
    if args.remote:
        mycluster.add_remote(args.remote, job_db)
    if args.create:
        if args.create and args.jobqueue and args.script and not args.taskpernode:
            mycluster.create_submit(args.jobqueue,
                                    script_name=args.create,
                                    job_db=job_db,
                                    my_script=args.script,
                                    my_name=args.jobname,
                                    my_output=args.jobname+'.out',
                                    num_tasks=args.ntasks,
                                    num_threads_per_task=args.nthreadspertask,
                                    project_name=args.project,
                                    wall_clock=args.maxtime,
                                    shared=args.shared,
                                    no_syscribe=True,
                                    openmpi_args = args.ompiargs,
                                    qos = args.qos)
        elif args.create and args.jobqueue and args.script and args.taskpernode:
            mycluster.create_submit(args.jobqueue,
                                    script_name=args.create,
                                    job_db=job_db,
                                    my_script=args.script,
                                    my_name=args.jobname,
                                    my_output=args.jobname+'.out',
                                    num_tasks=args.ntasks,
                                    project_name=args.project,
                                    wall_clock=args.maxtime,
                                    tasks_per_node=args.taskpernode,
                                    num_threads_per_task=args.nthreadspertask,
                                    shared=args.shared,
                                    no_syscribe=True,
                                    openmpi_args = args.ompiargs,
                                    qos = args.qos)
        else:
            parser.print_help()

if __name__ == "__main__":
    main()
