#!/usr/bin/env python3

__copyright__ = "Copyright 2014-2019, http://radical.rutgers.edu"
__license__   = "MIT"


import os
import sys
import time

import threading     as mt
import setproctitle  as spt

import radical.utils as ru
import radical.pilot as rp

from radical.pilot.messages import HeartbeatMessage

# ------------------------------------------------------------------------------
#
def main(sid, reg_addr, uid):
    '''
    This thin wrapper starts a RCT component It expects a single argument:
    a config to use for the component's configuration.  The config must contain:

      - uid : UID of component instance (unique to the hosting session)
      - name: name of the component
      - kind: type of component

    If the config contains a `heartbeat` section, that section must be formatted
    as follows:

        {
          'from'    : 'uid',
          'addr_pub': 'addr_pub',
          'addr_sub': 'addr_sub',
          'interval': <float>,
          'timeout' : <float>
        }

    If that section exists, the process will daemonize and heartbeats are used
    to manage the bridge lifetime: the lifetime of this bridge is then dependent
    on receiving heartbeats from the given `uid`: after `timeout` seconds of no
    heartbeats arriving, the bridge will terminate.  The bridge itself will
    publish heartbeats every `interval` seconds on the heartbeat channel under
    its own uid.

    If the heartbeat section is not present in the config file, the components
    lifetime is expected to be explicitly managed, i.e., that this wrapper
    process hosting the bridge is terminated externally.

    The config file may contain other entries which are passed to the component
    and are interpreted by the component implementation.
    '''

    # basic setup: logger and profiler
    log  = ru.Logger(name=uid, ns='radical.pilot', path=os.getcwd())
    prof = ru.Profiler(name=uid, ns='radical.pilot', path=os.getcwd())

    try:
        prof.prof('comp_start', uid=uid)
        prof.disable()
        wrapped_main(sid, reg_addr, uid, log, prof)

    finally:
        prof.enable()
        prof.prof('comp_stop', uid=uid)


# ------------------------------------------------------------------------------
#
def wrapped_main(sid, reg_addr, uid, log, prof):

    spt.setproctitle('rp.%s' % uid)

    term = mt.Event()
    reg  = ru.zmq.RegistryClient(url=reg_addr)

    hb_cfg = ru.TypedDict(reg['heartbeat'])
    c_cfg  = ru.TypedDict(reg['components.%s.cfg' % uid])

    reg.close()

    # start a non-primary session
    session = rp.Session(uid=sid, cfg=c_cfg,
                         _role=rp.Session._DEFAULT, _reg_addr=reg_addr)

    # create the instance and begin to work
    comp = rp.utils.BaseComponent.create(c_cfg, session)
    comp.start()

    # component runs - send heartbeats so that session knows about it
    hb_pub = ru.zmq.Publisher('heartbeat', hb_cfg.addr_pub, log=log, prof=prof)

    def hb_beat_cb():
        hb_pub.put('heartbeat', HeartbeatMessage(uid=uid))

    def hb_term_cb(hb_uid):
        comp.stop()
        term.set()
        return False

    hb = ru.Heartbeat(uid=uid,
                      timeout=hb_cfg.timeout,
                      interval=hb_cfg.interval,
                      beat_cb=hb_beat_cb,
                      term_cb=hb_term_cb,
                      log=log)
    hb.start()

    # always watch out for session heartbeat
    hb.watch(uid=sid)

    # react on session heartbeats
    def hb_sub_cb(topic, msg):
        hb_msg = HeartbeatMessage(from_dict=msg)
        if hb_msg.uid == sid:
            hb.beat(uid=sid)

    ru.zmq.Subscriber('heartbeat', hb_cfg.addr_sub,
                      topic='heartbeat', cb=hb_sub_cb,
                      log=log, prof=prof)

    # all is set up - we can sit idle 'til end of time.
    while not term.is_set():
        time.sleep(1)


# ------------------------------------------------------------------------------
#
if __name__ == "__main__":

    if len(sys.argv) != 4:
        sys.stderr.write('error: invalid arguments\n'
                         'usage: %s <sid> <reg_addr> <uid>\n'  % sys.argv[0])
        raise RuntimeError('invalid arguments: %s' % sys.argv)

    sid      = sys.argv[1]
    reg_addr = sys.argv[2]
    uid      = sys.argv[3]

    ru.daemonize(main=main, args=[sid, reg_addr, uid],
                 stdout='%s.out' % uid, stderr='%s.err' % uid)
    sys.exit(0)


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

