#!python

"""Parse and Launch AER Communication Programs.

Author: Yuhuang Hu
Email : yuhuang.hu@ini.uzh.ch
"""

from __future__ import print_function, absolute_import

import os
import sys
import argparse
import time

from pyaer.utils import expandpath
from pyaer.utils import ordered_yml_load
from pyaer.comm import AERProcess
from pyaer import log

HUB = "Hub"
PUB = "Publisher"
SUB = "Subscriber"
PUBSUB = "PubSuber"
SAVER = "Saver"

USE_DEFAULT = "use_default"
USE_DEFAULT_PUB = "use_default_pub"
USE_DEFAULT_SUB = "use_default_sub"
DEFAULT_URL = "tcp://127.0.0.1"
DEFAULT_PUBLISHER_PORT = "5100"
DEFAULT_SUBSCRIBER_PORT = "5099"
DEFAULT_AER_HUB_NAME = "PyAER Message Hub"
DEFAULT_AER_PUBLISHER_PROGRAM = "aer_publisher"
DEFAULT_AER_SUBSCRIBER_PROGRAM = "aer_subscriber"
DEFAULT_AER_PUBSUBER_PROGRAM = "aer_pubsuber"

DEFAULT_HDF5_MODE = "w"
DEFAULT_LIBVER_VERSION = "earliest"

PROGRAM_KEY = "program"

launch_logger = log.get_logger(
    "AER Launcher",
    log.INFO, stream=sys.stdout)


def parse_hub(hub_desc):
    parsed_cmd = ["aer_hub"]

    if hub_desc[USE_DEFAULT] is True:
        parsed_cmd += ["--url", DEFAULT_URL]
        parsed_cmd += ["--publisher_port", DEFAULT_PUBLISHER_PORT]
        parsed_cmd += ["--subscriber_port", DEFAULT_SUBSCRIBER_PORT]
        parsed_cmd += ["--aer_hub_name", DEFAULT_AER_HUB_NAME]

        return parsed_cmd

    parsed_cmd += ["--url", hub_desc["url"]]
    parsed_cmd += ["--publisher_port", str(hub_desc["publisher_port"])]
    parsed_cmd += ["--subscriber_port", str(hub_desc["subscriber_port"])]
    parsed_cmd += ["--aer_hub_name", hub_desc["aer_hub_name"]]

    return parsed_cmd


def parse_publisher(pub_desc):
    parsed_cmd = []

    # select program
    if PROGRAM_KEY in pub_desc:
        parsed_cmd += [expandpath(pub_desc[PROGRAM_KEY])]
    else:
        parsed_cmd += [DEFAULT_AER_PUBLISHER_PROGRAM]

    parsed_cmd += [
        "--url",
        DEFAULT_URL if "url" not in pub_desc else pub_desc.pop("url")]
    parsed_cmd += [
        "--port",
        DEFAULT_PUBLISHER_PORT if "port" not in pub_desc else
        str(pub_desc.pop("port"))]

    parsed_cmd += ["--name", pub_desc.pop("name")]
    parsed_cmd += ["--master_topic", pub_desc.pop("master_topic")]

    if "custom_pub" in pub_desc:
        # parse custom command
        # follows the form /file/path/class_name
        custom_pub, custom_class = os.path.split(pub_desc.pop("custom_pub"))

        parsed_cmd += ["--custom_pub", expandpath(custom_pub)]
        parsed_cmd += ["--custom_class", custom_class]
    else:
        parsed_cmd += ["--use_default_pub"]

    pub_desc.pop(USE_DEFAULT_PUB)

    try:
        parsed_cmd += ["--device", pub_desc.pop("device")]
        if pub_desc.pop("noise_filter", default=False):
            parsed_cmd += ["--noise_filter"]
        if "bias_file" in pub_desc:
            parsed_cmd += [
                "--bias_file", expandpath(pub_desc.pop("bias_file"))]
    except Exception:
        pass

    try:
        for (option, value) in pub_desc.items():
            parsed_cmd += ["--{}".format(option), str(value)]
    except Exception:
        pass

    return parsed_cmd


def parse_subscriber(sub_desc):
    parsed_cmd = []

    # select program
    if PROGRAM_KEY in sub_desc:
        parsed_cmd += [expandpath(sub_desc[PROGRAM_KEY])]
    else:
        parsed_cmd += [DEFAULT_AER_SUBSCRIBER_PROGRAM]

    parsed_cmd += [
        "--url",
        DEFAULT_URL if "url" not in sub_desc else sub_desc.pop("url")]
    parsed_cmd += [
        "--port",
        DEFAULT_SUBSCRIBER_PORT if "port" not in sub_desc else
        str(sub_desc.pop("port"))]
    parsed_cmd += ["--name", sub_desc.pop("name")]
    parsed_cmd += ["--topic", sub_desc.pop("topic")]

    if "custom_sub" in sub_desc:
        # parse custom command
        # follows the form /file/path/class_name
        custom_sub, custom_class = os.path.split(sub_desc.pop("custom_sub"))

        parsed_cmd += ["--custom_sub", expandpath(custom_sub)]
        parsed_cmd += ["--custom_class", custom_class]
    else:
        parsed_cmd += ["--use_default_sub"]

    sub_desc.pop(USE_DEFAULT_SUB)

    try:
        for (option, value) in sub_desc.items():
            parsed_cmd += ["--{}".format(option), str(value)]
    except Exception:
        pass

    return parsed_cmd


def parse_pubsuber(pubsuber_desc):
    parsed_cmd = []

    # select program
    if PROGRAM_KEY in pubsuber_desc:
        parsed_cmd += [expandpath(pubsuber_desc[PROGRAM_KEY])]
    else:
        parsed_cmd += [DEFAULT_AER_PUBSUBER_PROGRAM]

    parsed_cmd += [
        "--url",
        DEFAULT_URL if "url" not in pubsuber_desc else
        pubsuber_desc.pop("url")]
    parsed_cmd += [
        "--pub_port",
        DEFAULT_PUBLISHER_PORT if "pub_port" not in pubsuber_desc else
        str(pubsuber_desc.pop("pub_port"))]
    parsed_cmd += ["--pub_name", pubsuber_desc.pop("pub_name")]
    parsed_cmd += ["--pub_topic", pubsuber_desc.pop("pub_topic")]

    parsed_cmd += [
        "--sub_port",
        DEFAULT_SUBSCRIBER_PORT if "sub_port" not in pubsuber_desc else
        str(pubsuber_desc.pop("sub_port"))]
    parsed_cmd += ["--sub_name", pubsuber_desc.pop("sub_name")]
    parsed_cmd += ["--sub_topic", pubsuber_desc.pop("sub_topic")]

    custom_pubsuber, custom_class = os.path.split(
        pubsuber_desc.pop("custom_pubsuber"))
    parsed_cmd += ["--custom_pubsuber", expandpath(custom_pubsuber)]
    parsed_cmd += ["--custom_class", custom_class]

    # Leave this here, maybe useful
    try:
        parsed_cmd += ["--device", pubsuber_desc.pop("device")]
        if pubsuber_desc.pop("noise_filter", default=False):
            parsed_cmd += ["--noise_filter"]
        if "bias_file" in pubsuber_desc:
            parsed_cmd += ["--bias_file", expandpath(
                pubsuber_desc.pop("bias_file"))]
    except Exception:
        pass

    try:
        for (option, value) in pubsuber_desc.items():
            parsed_cmd += ["--{}".format(option), str(value)]
    except Exception:
        pass

    return parsed_cmd


def parse_saver(saver_desc):
    parsed_cmd = ["aer_saver"]

    parsed_cmd += [
        "--url",
        DEFAULT_URL if "url" not in saver_desc else saver_desc["url"]]
    parsed_cmd += [
        "--port",
        DEFAULT_SUBSCRIBER_PORT if "port" not in saver_desc else
        str(saver_desc["port"])]
    parsed_cmd += ["--name", saver_desc["name"]]
    parsed_cmd += ["--topic", saver_desc["topic"]]

    parsed_cmd += ["--filename", expandpath(saver_desc["filename"])]

    parsed_cmd += [
        "--mode",
        DEFAULT_HDF5_MODE if "mode" not in saver_desc else
        saver_desc["mode"]]
    parsed_cmd += [
        "--libver",
        DEFAULT_LIBVER_VERSION if "libver" not in saver_desc else
        saver_desc["libver"]]

    # Use HDF5 as the default saver
    try:
        if saver_desc["zarr"] is True:
            parsed_cmd += ["--zarr"]
    except Exception:
        parsed_cmd += ["--hdf5"]

    return parsed_cmd


parser = argparse.ArgumentParser("AER Launch")
parser.add_argument("--launch_file", type=expandpath,
                    help="AER Launch File")

args = parser.parse_args()

# load launch file
with open(args.launch_file, "r") as launch_file:
    launch_desc = ordered_yml_load(launch_file)

# main parsing loop
process_collector = []  # collect all active processes

try:
    for pg_i, (pg_type, pg_desc) in enumerate(launch_desc.items()):
        # if the first one is not a hub type, then start the default hub
        if pg_i == 0 and pg_type != HUB:
            parsed_hub_cmd = parse_hub({"use_default": True})
            process_collector.append(
                AERProcess(parsed_hub_cmd, daemon=True))

        if pg_type == HUB:
            parsed_hub_cmd = parse_hub(pg_desc)
            process_collector.append(
                AERProcess(parsed_hub_cmd))
        elif PUB in pg_type:
            parsed_pub_cmd = parse_publisher(pg_desc)
            process_collector.append(
                AERProcess(parsed_pub_cmd))
        elif SUB in pg_type:
            parsed_sub_cmd = parse_subscriber(pg_desc)
            process_collector.append(
                AERProcess(parsed_sub_cmd))
        elif PUBSUB in pg_type:
            parsed_pubsuber_cmd = parse_pubsuber(pg_desc)
            process_collector.append(
                AERProcess(parsed_pubsuber_cmd))
        elif SAVER in pg_type:
            parsed_saver_cmd = parse_saver(pg_desc)
            process_collector.append(
                AERProcess(parsed_saver_cmd))
        else:
            launch_logger.error("Unsupported Type {}".format(pg_type))

    # launching
    for process in process_collector:
        process.run()
except Exception:
    launch_logger.info("Error launching, terminating all processes.")
    for process in reversed(process_collector):
        try:
            process.stop()
        except Exception:
            pass
        del process
    sys.exit(1)

while True:
    try:
        time.sleep(1)
    except Exception:
        launch_logger.info("Exiting launcher, terminating all processes.")
        for process in reversed(process_collector):
            try:
                process.stop()
            except Exception:
                pass
            del process
        break
