#!/usr/bin/env python

"""
Cosmos Script
"""

import argparse
import json
import os
from functools import partial

from cosmos.api import Cosmos, default_get_submit_args


def get_db_url_or_path_to_sqlite(db_url_or_path_to_sqlite):
    if '://' not in db_url_or_path_to_sqlite:
        assert os.path.exists(
            db_url_or_path_to_sqlite), '`%s` does not exist on the filesystem' % db_url_or_path_to_sqlite
        db_url_or_path_to_sqlite = os.path.abspath(os.path.expanduser(db_url_or_path_to_sqlite))

    database_url = db_url_or_path_to_sqlite if ':' in db_url_or_path_to_sqlite else 'sqlite:///%s' % db_url_or_path_to_sqlite

    return database_url


def runweb(db_url_or_path_to_sqlite, host, port):
    database_url = get_db_url_or_path_to_sqlite(db_url_or_path_to_sqlite)

    cosmos_app = Cosmos(database_url=database_url)

    cosmos_app.runweb(host, port)


def shell(db_url_or_path_to_sqlite):
    database_url = get_db_url_or_path_to_sqlite(db_url_or_path_to_sqlite)

    cosmos_app = Cosmos(database_url=database_url)
    cosmos_app.shell()


def run(in_jobs, default_drm, default_queue, restart, max_cores,
        default_job_class=None, default_drm_options=None):
    """
    Create an embarassingly parallel workflow from all jobs in `in_jobs`.  `in_jobs` should be a json of a dict
    keyed by uid => command.
    """
    with open(in_jobs) as fp:
        jobs_dict = json.load(fp)

    cosmos_app = Cosmos(database_url='sqlite:///cosmos.sqlite',
                        default_drm=default_drm,
                        default_drm_options=default_drm_options,
                        default_job_class=default_job_class,
                        default_queue=default_queue,
                        get_submit_args=partial(default_get_submit_args, parallel_env='smp'))
    cosmos_app.initdb()
    wf = cosmos_app.start(name='workflow', restart=restart, skip_confirm=True)

    for uid, cmd in jobs_dict.iteritems():
        wf.add_task(func=lambda cmd: cmd,
                    params=dict(cmd=cmd),
                    stage_name='lambda',
                    uid=uid)
    wf.run(max_cores=max_cores)


if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__)
    sps = p.add_subparsers(title="Commands", dest='cmd')
    sps.required = True

    sp = sps.add_parser('runweb')
    sp.add_argument('db_url_or_path_to_sqlite')
    sp.add_argument('--host', '-H', default='0.0.0.0')
    sp.add_argument('--port', '-p', default=3232, type=int)

    sp = sps.add_parser('shell')
    sp.add_argument('db_url_or_path_to_sqlite')

    sp = sps.add_parser('run')
    sp.add_argument('in_jobs')
    sp.add_argument('--default-drm', '-drm', default='ge')
    sp.add_argument('--job-class', '-j')
    sp.add_argument('--queue', '-q')

    sp.add_argument('--max_cores', '--max-cores', '-c', type=int,
                    help="Maximum number (based on the sum of Task.core_req) of cores to use at once.  0 means unlimited.",
                    default=None)
    sp.add_argument('--restart', '-r', action='store_true',
                    help="Completely restart the Workflow.  Note this will delete all record of the Workflow in the database.")

    args = p.parse_args()
    print(vars(args))
    kwargs = dict(args._get_kwargs())
    cmd = kwargs.pop('cmd')
    globals()[cmd](**kwargs)
