#!/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

from radical.pilot.messages import HeartbeatMessage


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

      - uid : UID of bridge instance (unique to the hosting session)
      - name: name of the bridge
      - kind: type of bridge (`pubsub` or `queue`)

    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 bridge
    and are interpreted by the bridge implementation.

    After startup, the bridge's communication endpoint URLs are stored in the
    sessions registry under `bridges.<bridge_uid>`, in the form (shown for
    pubsub and queue type bridges):

        {
          'addr_pub': '$addr_pub',
          'addr_sub': '$addr_sub'
        }

        {
          'addr_put': '$addr_put',
          'addr_get': '$addr_get'
        }

    That config is formed so that any publishers, subscribers, putters or getters
    can obtain the respective bridge addresses automatically.  This also holds
    for command line tools like:

        > radical-pilot-bridge sid reg_addr test_pubsub.0000                 [1]
        > radical-pilot-sub    sid reg_addr test_pubsub.0000 foo &           [2]
        > radical-pilot-pub    sid reg_addr test_pubsub.0000 foo bar         [3]

    [1] establishes the pubsub channel 'command'
    [2] connect to the command channel, subscribe for topic `foo`
    [3] connect to the command channel, send messages for topic `foo`
    '''

    # 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('bridge_start', uid=uid)
        prof.disable()
        wrapped_main(sid, reg_addr, uid, log, prof)

    finally:
        prof.enable()
        prof.prof('bridge_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'])
    b_cfg  = ru.TypedDict(reg['bridges.%s.cfg' % uid])

    # create the instance and begin to work
    bridge = ru.zmq.Bridge.create(uid, cfg=b_cfg)

    reg['bridges.%s.addr_%s' % (uid, bridge.type_in )] = str(bridge.addr_in)
    reg['bridges.%s.addr_%s' % (uid, bridge.type_out)] = str(bridge.addr_out)

    reg.close()
    bridge.start()

    # re-enable the test below if timing issues crop up
  # if 'pubsub' in uid:
  #     d = ru.zmq.test_pubsub(bridge.channel, bridge.addr_pub, bridge.addr_sub)

    # bridge runs - send heartbeats so that cmgr 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):
        bridge.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)


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

