
import io
import os
import sys
import time
import shlex

import threading         as mt

import radical.utils     as ru

from .. import constants as rpc

from ..pytask           import PythonTask
from ..task_description import TASK_FUNCTION, TASK_EXEC
from ..task_description import TASK_PROC, TASK_SHELL, TASK_EVAL


# ------------------------------------------------------------------------------
#
class Worker(object):
    '''
    Implement the Raptor protocol for dispatching multiple Tasks on persistent
    resources.
    '''

    # --------------------------------------------------------------------------
    #
    def __init__(self, cfg, manager, rank):

        if cfg is None:
            cfg = dict()

        if isinstance(cfg, str): self._cfg = ru.Config(cfg=ru.read_json(cfg))
        else                   : self._cfg = ru.Config(cfg=cfg)

        self._rank = rank
        self._sbox = os.environ['RP_TASK_SANDBOX']
        self._uid  = os.environ['RP_TASK_ID']

        self._log  = ru.Logger(name=self._uid,   ns='radical.pilot.worker',
                               level='DEBUG', targets=['.'], path=self._sbox)
        self._prof = ru.Profiler(name=self._uid, ns='radical.pilot.worker',
                                 path=self._sbox)

        # register for lifetime management messages on the control pubsub
        psbox = os.environ['RP_PILOT_SANDBOX']
        ctrl_cfg = ru.read_json('%s/%s.cfg' % (psbox, rpc.CONTROL_PUBSUB))

        ru.zmq.Subscriber(rpc.CONTROL_PUBSUB, url=ctrl_cfg['sub'],
                          log=self._log, prof=self._prof, cb=self._control_cb)

        # we push hertbeat and registration messages on that pubsub also
        self._ctrl_pub = ru.zmq.Publisher(rpc.CONTROL_PUBSUB,
                                          url=ctrl_cfg['pub'],
                                          log=self._log,
                                          prof=self._prof)
        # let ZMQ settle
        time.sleep(1)

        # the manager (rank 0) registers the worker with the master
        if manager:

            self._ctrl_pub.put(rpc.CONTROL_PUBSUB, {'cmd': 'worker_register',
                                                    'arg': {'uid' : self._uid}})


          # # FIXME: we never unregister on termination
          # self._ctrl_pub.put(rpc.CONTROL_PUBSUB, {'cmd': 'worker_unregister',
          #                                         'arg': {'uid' : self._uid}})

        # run heartbeat thread in all ranks (one hb msg every `n` seconds)
        self._hb_delay  = 5
        self._hb_thread = mt.Thread(target=self._hb_worker)
        self._hb_thread.daemon = True
        self._hb_thread.start()

        self._log.debug('heartbeat thread started %s:%s', self._uid, self._rank)

        self._modes = dict()

        # run worker initialization *before* starting to work on requests.
        # the worker provides these builtin methods:
        #     eval:  evaluate a piece of python code with `eval`
        #     exec:  evaluate a piece of python code with `exec`
        #     call:  execute  a method or function call
        #     proc:  execute  a command line (fork/exec)
        #     shell: execute  a shell command
        self.register_mode(TASK_FUNCTION, self._dispatch_func)
        self.register_mode(TASK_EVAL,     self._dispatch_eval)
        self.register_mode(TASK_EXEC,     self._dispatch_exec)
        self.register_mode(TASK_PROC,     self._dispatch_proc)
        self.register_mode(TASK_SHELL,    self._dispatch_shell)

        # prepare base env dict used for all tasks
        self._task_env = dict()
        for k,v in os.environ.items():
            if k.startswith('RP_'):
                self._task_env[k] = v


    # --------------------------------------------------------------------------
    #
    def _hb_worker(self):

        while True:

            self._ctrl_pub.put(rpc.CONTROL_PUBSUB,
                    {'cmd': 'worker_rank_heartbeat',
                     'arg': {'uid' : self._uid,
                             'rank': self._rank}})

            time.sleep(self._hb_delay)


    # --------------------------------------------------------------------------
    #
    def _control_cb(self, topic, msg):

        if msg['cmd'] == 'terminate':
            self.stop()
            self.join()
            sys.exit()

        elif msg['cmd'] == 'worker_terminate':
            self._log.debug('worker_terminate signal')

            if msg['arg']['uid'] == self._uid:
                self.stop()
                self.join()
                sys.exit()


    # --------------------------------------------------------------------------
    #
    def get_master(self):
        '''
        The worker can submit tasks back to the master - this method will
        return a small shim class to provide that capability.  That class has
        a single method `run_task` which accepts a single `rp.TaskDescription`
        from which a `rp.Task` is created and executed.  The call then waits for
        the task's completion before returning it in a dict representation, the
        same as when passed to the master's `result_cb`.

        Note: the `run_task` call is running in a separate thread and will thus
              not block the master's progress.
        '''

        # ----------------------------------------------------------------------
        class Master(object):

            def __init__(self, addr):
                self._task_service_ep = ru.zmq.Client(url=addr)

            def run_task(self, td):
                return self._task_service_ep.request('run_task', td)
        # ----------------------------------------------------------------------

        return Master(self._cfg.ts_addr)


    # --------------------------------------------------------------------------
    #
    def start(self):

        raise NotImplementedError('`start()` must be implemented by child class')


    # --------------------------------------------------------------------------
    #
    def stop(self):

        raise NotImplementedError('`run()` must be implemented by child class')


    # --------------------------------------------------------------------------
    #
    def join(self):

        raise NotImplementedError('`join()` must be implemented by child class')


    # --------------------------------------------------------------------------
    #
    def register_mode(self, name, dispatcher):

        if name in self._modes:
            raise ValueError('mode %s already registered' % name)

        self._modes[name] = dispatcher


    # --------------------------------------------------------------------------
    #
    def get_dispatcher(self, name):

        if name not in self._modes:
            raise ValueError('mode %s unknown' % name)

        return self._modes[name]


    # --------------------------------------------------------------------------
    #
    def _dispatch_func(self, task):
        '''
        We expect three attributes: 'function', containing the name of the
        member method or free function to call, `args`, an optional list of
        unnamed parameters, and `kwargs`, and optional dictionary of named
        parameters.

        *function* is resolved first against `locals()`, then `globals()`, then
        attributes of the implementation class (member functions of *base*, as
        provided to `MPIWorkerRank()`). Finally, an attempt is made to
        deserialize a PythonTask from *function*. The first non-null resolution
        of *function* is used as the callable.

        Raises
        ------
        ValueError
            if *function* cannot be resolved.

        NOTE: MPI function tasks will get a private communicator passed as first
              unnamed argument.
        '''

        uid  = task['uid']
        func = task['description']['function']

        to_call = ''
        names   = ''
        args    = task['description'].get('args',   [])
        kwargs  = task['description'].get('kwargs', {})
        py_func = False

        self._log.debug('orig args: %s : %s', args, kwargs)

        # check if we have a serialized object
        self._log.debug('func serialized: %d: %s', len(func), func)
        try:
            # FIXME: can we have a better test than try/except?  This hides
            #        potential errors...
            # FIXME: ensure we did not get args and kwargs from above
            to_call, args, kwargs = PythonTask.get_func_attr(func)
            py_func = True
        except:
            pass

        # check if `func_name` is a global name
        if not to_call:
            assert func
            names   = dict(list(globals().items()) + list(locals().items()))
            to_call = names.get(func)

        # if not, check if this is a class method of this worker implementation
        if not to_call:
            to_call = getattr(self, func, None)

        # check if we have a serialized object
        if not to_call:
            self._log.debug('func serialized: %d: %s', len(func), func)

            try:
                to_call, _args, _kwargs = PythonTask.get_func_attr(func)

            except Exception:
                self._log.exception('function is not a PythonTask [%s] ', uid)

            else:
                py_func = True
                if args or kwargs:
                    raise ValueError('`args` and `kwargs` must be empty for'
                                     'PythonTask function [%s]' % uid)
                else:
                    args   = _args
                    kwargs = _kwargs

        if not to_call:
            self._log.error('no %s in \n%s\n\n%s', func, names, dir(self))
            raise ValueError('%s callable %s not found: %s' % (uid, func, task))

        comm = task.get('mpi_comm')
        if comm:
            # we have an MPI communicator we need to inject into the function's
            # arguments.
            if py_func:
                # For a `py_func` we add the communicator as `comm` kwarg if
                # that is set to None, and otherwise as first `arg` if that is
                # None.  If neither is true we'll error out.
                # NOTE that we don't change the number of arguments either way.
                if 'comm' in kwargs and kwargs['comm'] is None:
                    kwargs['comm'] = comm
                elif args and args[0] is None:
                    args[0] = comm
                else:
                    raise RuntimeError("can't inject communicator for %s: %s: %s",
                                       task['uid'], args, kwargs)
            else:
                args.insert(0, comm)

        # make sure we capture stdout / stderr
        bak_stdout = sys.stdout
        bak_stderr = sys.stderr

        strout = None
        strerr = None

        # set the task environment
        old_env = os.environ.copy()

        for k, v in task['description'].get('environment', {}).items():
            os.environ[k] = str(v)

        try:
            # redirect stdio to capture them during execution
            sys.stdout = strout = io.StringIO()
            sys.stderr = strerr = io.StringIO()

            self._prof.prof('rank_start', uid=uid)
            self._log.debug('to call %s: %s : %s', to_call, args, kwargs)
            val = to_call(*args, **kwargs)
            self._prof.prof('rank_stop', uid=uid)
            out = strout.getvalue()
            err = strerr.getvalue()
            exc = (None, None)
            ret = 0

        except Exception as e:
            self._log.exception('_call failed: %s', task['uid'])
            val = None
            out = strout.getvalue()
            err = strerr.getvalue() + ('\ncall failed: %s' % e)
            exc = (repr(e), '\n'.join(ru.get_exception_trace()))
            ret = 1

        finally:
            # restore stdio
            sys.stdout = bak_stdout
            sys.stderr = bak_stderr

            # remove communicator from args again
            if comm:
                if py_func:
                    if 'comm' in kwargs:
                        del kwargs['comm']
                    elif args:
                        args[0] = None
                else:
                    args.pop(0)

            os.environ = old_env

        self._log.debug('%s: got %s', uid, out)

        return out, err, ret, val, exc


    # --------------------------------------------------------------------------
    #
    def _dispatch_eval(self, task):
        '''
        We expect a single attribute: 'code', containing the Python
        code to be eval'ed
        '''

        uid  = task['uid']
        code = task['description']['code']
        assert code

        bak_stdout = sys.stdout
        bak_stderr = sys.stderr

        strout = None
        strerr = None

        old_env = os.environ.copy()

        for k, v in task['description'].get('environment', {}).items():
            os.environ[k] = str(v)

        try:
            # redirect stdio to capture them during execution
            sys.stdout = strout = io.StringIO()
            sys.stderr = strerr = io.StringIO()

            self._log.debug('eval [%s] [%s]', code, task['uid'])

            self._prof.prof('rank_start', uid=uid)
            val = eval(code)
            self._prof.prof('rank_stop', uid=uid)
            out = strout.getvalue()
            err = strerr.getvalue()
            exc = (None, None)
            ret = 0

        except Exception as e:
            self._log.exception('_eval failed: %s', task['uid'])
            val = None
            out = strout.getvalue()
            err = strerr.getvalue() + ('\neval failed: %s' % e)
            exc = (repr(e), '\n'.join(ru.get_exception_trace()))
            ret = 1

        finally:
            # restore stdio
            sys.stdout = bak_stdout
            sys.stderr = bak_stderr

            os.environ = old_env

        return out, err, ret, val, exc



    # --------------------------------------------------------------------------
    #
    def _dispatch_exec(self, task):
        '''
        We expect a single attribute: 'code', containing the Python code to be
        exec'ed.  The optional attribute `pre_exec` can be used for any import
        statements and the like which need to run before the executed code.
        '''

        bak_stdout = sys.stdout
        bak_stderr = sys.stderr

        strout = None
        strerr = None

        old_env = os.environ.copy()

        for k, v in task['description'].get('environment', {}).items():
            os.environ[k] = str(v)

        try:
            # redirect stdio to capture them during execution
            sys.stdout = strout = io.StringIO()
            sys.stderr = strerr = io.StringIO()

            uid  = task['uid']
            pre  = task['description'].get('pre_exec', [])
            code = task['description']['code']

            # create a wrapper function around the given code
            lines = code.split('\n')
            outer = 'def _my_exec():\n'
            for line in lines:
                outer += '    ' + line + '\n'

            # call that wrapper function via exec, and keep the return value
            src = '%s\n\n%s\n\nresult=_my_exec()' % ('\n'.join(pre), outer)

            # assign a local variable to capture the code's return value.
            loc = dict()
            self._prof.prof('rank_start', uid=uid)
            exec(src, {}, loc)                # pylint: disable=exec-used # noqa
            self._prof.prof('rank_stop', uid=uid)
            val = loc['result']
            out = strout.getvalue()
            err = strerr.getvalue()
            exc = (None, None)
            ret = 0

        except Exception as e:
            self._log.exception('_exec failed: %s', task['uid'])
            val = None
            out = strout.getvalue()
            err = strerr.getvalue() + ('\nexec failed: %s' % e)
            exc = (repr(e), '\n'.join(ru.get_exception_trace()))
            ret = 1

        finally:
            # restore stdio
            sys.stdout = bak_stdout
            sys.stderr = bak_stderr

            os.environ = old_env

        return out, err, ret, val, exc


    # --------------------------------------------------------------------------
    #
    def _dispatch_proc(self, task):
        '''
        We expect two attributes: 'executable', containing the executabele to
        run, and `arguments` containing a list of arguments (strings) to pass as
        command line arguments.  We use `sp.Popen` to run the fork/exec, and to
        collect stdout, stderr and return code
        '''

        try:
            import subprocess as sp

            uid  = task['uid']
            exe  = task['description']['executable']
            args = task['description'].get('arguments', list())
            tenv = task['description'].get('environment', dict())

            cmd  = '%s %s' % (exe, ' '.join([shlex.quote(arg) for arg in args]))
          # self._log.debug('proc: --%s--', args)
            self._prof.prof('rank_start', uid=uid)
            proc = sp.Popen(cmd, env=tenv,  stdin=None,
                            stdout=sp.PIPE, stderr=sp.PIPE,
                            close_fds=True, shell=True)
            out, err = proc.communicate()
            ret      = proc.returncode
            exc      = (None, None)
            self._prof.prof('rank_stop', uid=uid)

        except Exception as e:
            self._log.exception('proc failed: %s', task['uid'])
            out = None
            err = 'exec failed: %s' % e
            exc = (repr(e), '\n'.join(ru.get_exception_trace()))
            ret = 1

        return out, err, ret, None, exc


    # --------------------------------------------------------------------------
    #
    def _dispatch_shell(self, task):
        '''
        We expect a single attribute: 'command', containing the command
        line to be called as string.
        '''

        try:
            uid = task['uid']
            cmd = task['description']['command']
            env = task['description']['environment']
          # self._log.debug('shell: --%s--', cmd)
            self._prof.prof('rank_start', uid=uid)
            out, err, ret = ru.sh_callout(cmd, shell=True, env=env)
            exc = (None, None)
            self._prof.prof('rank_stop', uid=uid)

        except Exception as e:
            self._log.exception('_shell failed: %s', task['uid'])
            out = None
            err = 'shell failed: %s' % e
            exc = (repr(e), '\n'.join(ru.get_exception_trace()))
            ret = 1

      # os.environ = old_env

        return out, err, ret, None, exc


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

