'''
The Proc sub is used to spin up worker processes that run hub referenced
coroutines.
'''
# Import python libs
import os
import sys
import atexit
import itertools
import asyncio
import subprocess
# Import third party libs
import msgpack


def __init__(hub):
    '''
    Create constants used by the client and server side of procs
    '''
    hub.proc.DELIM = b'd\xff\xcfCO)\xfe='
    hub.proc.D_FLAG = b'D'
    hub.proc.I_FLAG = b'I'
    hub.proc.Workers = {}
    hub.proc.WorkersIter = {}
    hub.proc.WorkersTrack = {}


def _get_cmd(hub, ind, ref, ret_ref, sock_dir):
    '''
    Return the shell command to execute that will start up the worker
    '''
    code = 'import sys; '
    code += 'import pop.hub; '
    code += 'hub = pop.hub.Hub(); '
    code += 'hub.pop.sub.add("pop.mods.proc"); '
    code += f'hub.proc.worker.start("{sock_dir}", "{ind}", "{ref}", "{ret_ref}")'
    cmd = f'{sys.executable} -c \'{code}\''
    return cmd


def mk_proc(hub, ind, workers, ret_ref, sock_dir):
    '''
    Create the process and add it to the passed in workers dict at the
    specified index
    '''
    ref = os.urandom(3).hex() + '.sock'
    workers[ind] = {'ref': ref}
    workers[ind]['path'] = os.path.join(sock_dir, ref)
    cmd = _get_cmd(hub, ind, ref, ret_ref, sock_dir)
    workers[ind]['proc'] = subprocess.Popen(cmd, shell=True)
    workers[ind]['pid'] = workers[ind]['proc'].pid


async def pool(hub, num, name='Workers', callback=None, sock_dir=None):
    '''
    Create a new local pool of process based workers

    :param num: The number of processes to add to this pool
    :param ref: The location on the hub to create the Workers dict used to
        store the worker pool, defaults to `hub.pop.proc.Workers`
    :param callback: The pop ref to call when the process communicates
        back
    '''
    ret_ref = os.urandom(3).hex() + '.sock'
    ret_sock_path = os.path.join(sock_dir, ret_ref)
    if not hub.proc.Tracker:
        hub.proc.init.mk_tracker()
    workers = {}
    if callback:
        await asyncio.start_unix_server(
                hub.proc.init.ret_work(callback),
                path=ret_sock_path)
    for ind in range(num):
        hub.proc.init.mk_proc(ind, workers, ret_ref, sock_dir)
    w_iter = itertools.cycle(workers)
    hub.proc.Workers[name] = workers
    hub.proc.WorkersIter[name] = w_iter
    hub.proc.WorkersTrack[name] = {
        'subs': [],
        'ret_ref': ret_ref,
        'sock_dir': sock_dir}
    up = set()
    while True:
        for ind in workers:
            if os.path.exists(workers[ind]['path']):
                up.add(ind)
        if len(up) == num:
            break
        await asyncio.sleep(0.01)
    # TODO: This seems to be spawning extra procs, this should be fixed
    #asyncio.ensure_future(hub.proc.init.maintain(name))


async def maintain(hub, name):
    '''
    Keep an eye on these processes
    '''
    workers = hub.proc.Workers[name]
    while True:
        for ind, data in workers.items():
            if not data['proc'].poll():
                hub.proc.init.mk_proc(ind, workers)
        await asyncio.sleep(2)


def mk_tracker(hub):
    '''
    Create the process tracker, this simply makes a data structure to hold
    process references and sets them to be terminated when the system is
    shutdown.
    '''
    hub.proc.Tracker = True
    atexit.register(hub.proc.init.clean)


def clean(hub):
    '''
    Clean up the processes registered in the tracker
    '''
    for name, workers in hub.proc.Workers.items():
        for ind in workers:
            workers[ind]['proc'].terminate()


def ret_work(hub, callback):
    async def work(reader, writer):
        '''
        Process the incoming work
        '''
        inbound = await reader.readuntil(hub.proc.DELIM)
        inbound = inbound[:-len(hub.proc.DELIM)]
        payload = msgpack.loads(inbound, raw=False)
        ret = await callback(payload)
        ret = msgpack.dumps(ret, use_bin_type=True)
        ret += hub.proc.DELIM
        writer.write(ret)
        await writer.drain()
        writer.close()
    return work
