#!/usr/bin/env python3

import os
import sys
import glob

import radical.utils as ru


# ------------------------------------------------------------------------------
#
def usage(msg):

    if msg:
        print('    Error: %s' % msg)

    print('''    Usage: %s <session_dir> [profile]

    This script will parse all logfiles in the session tree and search for PRTE
    profile entries.  It will parse those line, map the PRTE task IDs to RP IDs,
    and write RU formatted profile entries to the given profile (default:
    `radical.prte.prof`).  Note that the parsed profile entries will be
    appended if the given profile exists.

    The PRTE task ID to RP ID mapping is depends on the availability of the
    tasks STDERR: PRTE  will write a line line below to the task's stderr:

        [batch3:122527] JOB [3673,4] EXECUTING

    We parse that line and thus obtain the respective PRTE task ID.  The
    `agent.staging.output` default implementation will do that parsing and will
    write the ID mapping to `<x>.prte2pid.log`, where `<x>` is the runtime ID
    of the parsing output component.  This script will search the specified
    session directory for file that match that name and extract the mapping from
    those files.   This script will raise an error if a PRTE profile line cannot
    be mapped to an RP ID.

    ''' % sys.argv[0])

    if msg:
        sys.exit(1)


# ------------------------------------------------------------------------------
#
def handle_line(prof, line, pid, idmap):

    line = line.replace('[',      '|') \
               .replace(']',      '|') \
               .replace(' AT ',   '|') \
               .replace(' NULL ', '|')

    elems = list()
    for e in line.split('|'):
        e = e.strip()
        if e:
            elems.append(e)

    if elems[5] == 'ACTIVATE JOB' in line:
      # if '33357,24' in line:
      #     print '0%s' % elems[2:]
        elems.insert(7, pid)

        if len(elems) == 9:
            sys.stdout.write('#')
            elems.insert(7, elems[2])
        else:
            sys.stdout.write('+')

    elif elems[5] == 'ACTIVATE PROC' in line:
        if 'STATE COMMUNICATION FAILURE' in line:
            elems[3] = elems[7]
            elems[7] = ''
            sys.stdout.write('?')
        else:
            sys.stdout.write('-')

    while len(elems) < 10:
        elems.append('?')

    sys.stdout.flush()

 #  print elems[2:]
    #   ... DEBUG   : prte output: [batch3:80017] [[33357,0],0] [1565343424.463101] ACTIVATE JOB [33357,0] STATE PENDING ALLOCATION AT ../../../../../../../source/prrte-dev/orte/tools/prte/prte.c:497
    # ['... DEBUG   : prte output:', 'batch3:80017', '33357,0', ',0', '1565343424.463101', 'ACTIVATE JOB', '33357,0', 'STATE PENDING ALLOCATION', '../../../../../../../source/prrte-dev/orte/tools/prte/prte.c:497']
    # _     '2019-08-09 05:37:34,815: agent_0             : MainProcess : DVMWatcher     : DEBUG   : prte output:',
    # node  'batch3:80017',
    # dvm   '33357,0',
    # dvmd  ',0'
    # time  '1565343424.463101',
    # _     'ACTIVATE JOB',
    # tid   '33357,0',
    # proc  '',
    # event 'STATE PENDING ALLOCATION',
    #       '../../../../../../../source/prrte-dev/orte/tools/prte/prte.c:497']
    try:
        _, node, dvm, dvmd, tstamp, _, tid, proc, event, _ = elems

        dvm    = ('%s%s' % (dvm, dvmd)).replace(',', '.')
        event  = event.replace(' ', '_').lower()
        event  = event.replace('state_', '')
        tstamp = float(tstamp)

        uid    = idmap.get(tid, pid)
        uid    = uid.replace(',', '.')

        prof.prof('prte_%s' % event, uid=uid, state='AGENT_EXECUTING',
                  # timestamp=tstamp, comp='prte.dvm', tid=dvm, msg=node)
                  ts=tstamp, comp='prte.dvm', tid=dvm, msg=node)
    except Exception as e:
        print('skip line: %s' % str(e))
        print(line)


# ------------------------------------------------------------------------------
#
def handle_file(prof, fname, pid ,idmap):

  # print '    %-40s' % os.path.basename(fname),
    with open(fname, 'r') as fin:
        for line in fin.readlines():
            if ' ACTIVATE ' in line:
                handle_line(prof, line.strip(), pid, idmap)
  # print
  # print idmap


# ------------------------------------------------------------------------------
#
def map_ids(fname, idmap):

    with open(fname, 'r') as fin:
        for line in fin.readlines():
            if 'PRTE IDMAP:' in line:
                elems = line.split()
                tid, uid   = elems[-1].split(':',1)
                idmap[tid] = uid


# ------------------------------------------------------------------------------
#
def handle_dir(prof, dname):

    pid   = os.path.basename(dname)
    idmap = dict()

    for fname in glob.glob('%s/agent*staging*output*.log' % dname):
        print(fname)
        map_ids(fname, idmap)

    for fname in glob.glob('%s/*.log' % dname):
        handle_file(prof, fname, pid, idmap)

  # prof.prof('event', uid='uid', state='state',
  #           timestamp=0.0, comp='comp', tid='tid',
  #           msg='msg')

  # print idmap
    print()
    print()

# ------------------------------------------------------------------------------
#
def main():

    if len(sys.argv) not in [2, 3]:
        usage('wrong argument number')

    sid = sys.argv[1]

    if not os.path.isdir(sid):
        usage('no such dir: %s' % sid)

    if len(sys.argv) == 4:
        pname = sys.argv[2]
    else:
        pname = 'radical.prte'


    prof = ru.Profiler(pname, path=sid)

    # iterate through all `pilot.*` subdirs in the session tree
    for entry in glob.glob('%s/pilot.*' % sid):
        if os.path.isdir(entry):
            print()
            print('--> %s' %  entry)
            handle_dir(prof, entry)


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

    main()


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

