#!python

"""Parse and Launch AER Communication Programs.

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

from __future__ import print_function, absolute_import

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"

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 pub_desc[USE_DEFAULT] is True:
        parsed_cmd += [DEFAULT_AER_PUBLISHER_PROGRAM]
    else:
        parsed_cmd += [expandpath(pub_desc[PROGRAM_KEY])]
    pub_desc.pop(USE_DEFAULT)

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

    if pub_desc[USE_DEFAULT_PUB] is True:
        parsed_cmd += ["--use_default_pub"]
    else:
        # parse custom command
        parsed_cmd += ["--custom_pub", expandpath(pub_desc.pop("custom_pub"))]
        parsed_cmd += ["--custom_class", pub_desc.pop("custom_class")]

    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"]
        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 sub_desc[USE_DEFAULT] is True:
        parsed_cmd += [DEFAULT_AER_SUBSCRIBER_PROGRAM]
    else:
        parsed_cmd += [expandpath(sub_desc[PROGRAM_KEY])]
    sub_desc.pop(USE_DEFAULT)

    parsed_cmd += ["--url", sub_desc.pop("url")]
    parsed_cmd += ["--port", str(sub_desc.pop("port"))]
    parsed_cmd += ["--topic", sub_desc.pop("topic")]
    parsed_cmd += ["--name", sub_desc.pop("name")]

    if sub_desc[USE_DEFAULT_SUB] is True:
        parsed_cmd += ["--use_default_sub"]
    else:
        # parse custom command
        parsed_cmd += ["--custom_sub", expandpath(sub_desc.pop("custom_sub"))]
        parsed_cmd += ["--custom_class", sub_desc.pop("custom_class")]

    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 pubsuber_desc[USE_DEFAULT] is True:
        parsed_cmd += [DEFAULT_AER_PUBSUBER_PROGRAM]
    else:
        parsed_cmd += [expandpath(pubsuber_desc[PROGRAM_KEY])]
    pubsuber_desc.pop(USE_DEFAULT)

    parsed_cmd += ["--url", pubsuber_desc.pop("url")]

    parsed_cmd += ["--pub_port", str(pubsuber_desc.pop("pub_port"))]
    parsed_cmd += ["--pub_topic", pubsuber_desc.pop("pub_topic")]
    parsed_cmd += ["--pub_name", pubsuber_desc.pop("pub_name")]

    parsed_cmd += ["--sub_port", str(pubsuber_desc.pop("sub_port"))]
    parsed_cmd += ["--sub_topic", pubsuber_desc.pop("sub_topic")]
    parsed_cmd += ["--sub_name", pubsuber_desc.pop("sub_name")]

    parsed_cmd += ["--custom_pubsuber", expandpath(
        pubsuber_desc.pop("custom_pubsuber"))]
    parsed_cmd += ["--custom_class", pubsuber_desc.pop("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"]
        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", saver_desc["url"]]
    parsed_cmd += ["--port", str(saver_desc["port"])]
    parsed_cmd += ["--topic", saver_desc["topic"]]
    parsed_cmd += ["--name", saver_desc["name"]]

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

    try:
        if saver_desc["hdf5"] is True:
            parsed_cmd += ["--hdf5"]
    except Exception:
        pass

    try:
        if saver_desc["zarr"] is True:
            parsed_cmd += ["--zarr"]
    except Exception:
        pass

    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):
        process.stop()


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