"""
The Proc sub is used to spin up worker processes that run hub referenced
coroutines.
"""
import asyncio
import atexit
import itertools
import os
import subprocess
import sys

import msgpack

import pop.hub


def __init__(hub: "pop.hub.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: "pop.hub.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: "pop.hub.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: "pop.hub.Hub", num, name: str = "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 hasattr(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: "pop.hub.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: "pop.hub.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: "pop.hub.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: "pop.hub.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
