#!/usr/bin/env python3

import os
import sys
import glob
import logging
import radical.utils       as ru
import radical.pilot       as rp


# ------------------------------------------------------------------------------
#
def run_record(rec):

    session = None
    rep = ru.Reporter('radical')

    try:

        rep.title('Session Replay: %s' % rec)

        rep.header('create session')
        s_dict  = ru.read_json_str("%s/session.json" % rec)
        dburl   = s_dict.get('dburl')
        rep.info('session dburl: %s' % dburl)

        session = rp.Session()
        rep.ok('session uid  : %s' % session.uid)

        pmgr    = rp.PilotManager(session=session)
        rep.ok('pilot manager: %s' % pmgr.uid)

        tmgr    = rp.TaskManager(session=session)
        rep.ok('task manager : %s' % tmgr.uid)

        rep.header('create pilots')
        pds = list()
        for pd_json in glob.glob("%s/pilot.*.json" % rec):
            pd_dict = ru.read_json(pd_json)
            pd      = rp.PilotDescription()
            for key, val in pd_dict.items():
                pd.set_attribute(key, val)
            pds.append(pd)
            rep.info('%-15s [%3d cores]' % (pd.resource, pd.cores))
        pilots = pmgr.submit_pilots(pds)
        rep.ok('pilots submitted')

        rep.header('using pilots')
        tmgr.add_pilots(pilots)

        batch = 0
        while True:
            ud_json_list = glob.glob("%s/task.*.batch.%03d.json" % (rec, batch))

            if not ud_json_list:
                rec.header('no more task batches found')
                break

            rep.header('submit tasks [batch %d]' % batch)
            uds = list()
            for ud_json in ud_json_list:
                ud_dict = ru.read_json(ud_json)
                ud      = rp.TaskDescription()
                for key, val in ud_dict.items():
                    ud.set_attribute(key, val)
                uds.append(ud)
                args = ud_dict.get('arguments', [])
                rep.info('%s  %s [%3d cores]' % (ud.executable, ' '.join(args), ud.cores))
            tasks = tmgr.submit_tasks(uds)
            rep.ok('tasks submitted  [batch %d]' % batch)

            rep.info('wait for tasks [batch %d]' % batch)
            tmgr.wait_tasks()
            rep.ok('tasks all done   [batch %d]' % batch)

            for u in tasks:
                rep.info("%s (@ %s) state %s, exit %s"
                    % (u.uid, u.execution_locations, u.state, u.exit_code))

            batch += 1


    except Exception as e:
        logging.exception('error')
        rep.error("Exception caught: %s" % e)

    finally:

        if session:
            rep.info('closing session %s' % session.uid)
            session.close()
            rep.ok('session closed')


# ------------------------------------------------------------------------------
#
def usage (msg=None, noexit=False) :

    if  msg :
        print("\n      Error: %s" % msg)

    print("""
      usage      : %s [-r rec]
      example    : %s -r /tmp/recorded_session
      options :
        -r <rec> : run the session recorded in the directory 'rec'
                   if not specified, we use the value of the env variable
                   RADICAL_PILOT_SESSION_RECORD (if available)

""" % (sys.argv[0], sys.argv[0]))

    if  msg :
        sys.exit (1)

    if  not noexit :
        sys.exit (0)


# ------------------------------------------------------------------------------
#
if __name__ == '__main__' :

    import optparse
    parser = optparse.OptionParser (add_help_option=False)

    parser.add_option('-r', '--record', dest='rec')
    parser.add_option('-h', '--help',   dest='help', action="store_true")

    options, args = parser.parse_args ()

    if  options.help :
        usage ()

    if  options.rec :
        rec = options.rec
    else:
        print('looking for RADICAL_PILOT_SESSION_RECORD')
        rec = os.environ.get('RADICAL_PILOT_RECORD_SESSION')
        if not rec:
            print('not found')

    if not rec:
        usage ("No record specified")

    # don't record a recorded session
    if 'RADICAL_PILOT_RECORD_SESSION' in os.environ:
        del os.environ['RADICAL_PILOT_RECORD_SESSION']

    run_record(rec)


# ------------------------------------------------------------------------------

