#!/usr/bin/env python3

from argparse import ArgumentParser
import signal
import subprocess
import sys
import time

SATNOGS_FLOWGRAPH_SCRIPTS = {
    'AFSK1K2': 'satnogs_afsk1200_ax25.py',
    'AMSAT_DUV': 'satnogs_amsat_fox_duv_decoder.py',
    'APT': 'satnogs_noaa_apt_decoder.py',
    'ARGOS_BPSK_PMT_A3': 'satnogs_argos_bpsk_ldr.py',
    'BPSK': 'satnogs_bpsk.py',
    'CW': 'satnogs_cw_decoder.py',
    'FM': 'satnogs_fm.py',
    'FSK': 'satnogs_fsk.py',
    'GFSK_RKTR': 'satnogs_reaktor_hello_world_fsk9600_decoder.py',
    'GFSK/BPSK': 'satnogs_qubik_telem.py',
    'SSTV': 'satnogs_sstv_pd120_demod.py'
}

SATNOGS_FLOWGRAPH_MODES = {
    'AFSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['AFSK1K2'],
        'has_baudrate': False,
        'has_framing': False
    },
    'APT': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['APT'],
        'has_baudrate': False,
        'has_framing': False,
    },
    'BPSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['BPSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax25'
    },
    'BPSK PMT-A3': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['ARGOS_BPSK_PMT_A3'],
        'has_baudrate': True,
        'has_framing': False
    },
    'CW': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['CW'],
        'has_baudrate': True,
        'has_framing': False
    },
    'DUV': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['AMSAT_DUV'],
        'has_baudrate': False,
        'has_framing': False
    },
    'FM': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FM'],
        'has_baudrate': False,
        'has_framing': False
    },
    'FSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax25'
    },
    'FSK AX.100 Mode 5': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax100_mode5'
    },
    'FSK AX.100 Mode 6': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax100_mode6'
    },
    'GFSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax25'
    },
    'GFSK Rktr': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['GFSK_RKTR'],
        'has_baudrate': True,
        'has_framing': False
    },
    'GFSK/BPSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['GFSK/BPSK'],
        'has_baudrate': True,
        'has_framing': False,
    },
    'GMSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax25'
    },
    'MSK': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax25'
    },
    'MSK AX.100 Mode 5': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax100_mode5'
    },
    'MSK AX.100 Mode 6': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['FSK'],
        'has_baudrate': True,
        'has_framing': True,
        'framing': 'ax100_mode6'
    },
    'SSTV': {
        'script_name': SATNOGS_FLOWGRAPH_SCRIPTS['SSTV'],
        'has_baudrate': False,
        'has_framing': False
    },
}

SATNOGS_FLOWGRAPH_MODE_DEFAULT = 'FM'

ZMQ_SUB_SCRIPT_NAME = 'zmq_frame_decoder_sub'
ZMQ_DEFAULT_PARAMETERS = {
    'zmq-dump-host': '127.0.0.1',
    'zmq-dump-port': 16887,
}


class FlowgraphDispatcher:
    """
    Generate & execute SatNOGS Flowgraphs

    This class performs two tasks:
        1. Run a GNU Radio flowgraph with the appropriate set of parameters.
        2. Run ZMQ subscriber process, that receives decoded frames and writes them to files.

    :param parameters: The parameters needed for blocks instantiation
    :type parameters: argparse.Namespace
    """

    def __init__(self, parameters):
        self.parameters = {}
        self.flowgraph_script = ''
        self.flowgraph_process = None
        self.zmq_sub_process = None

        for arg in vars(parameters):
            self.parameters[arg] = getattr(parameters, arg)

        self.zmq_parameters = ZMQ_DEFAULT_PARAMETERS
        self.zmq_parameters['decoded-data-file-path'] = self.parameters['decoded-data-file-path']

        mode = self.parameters['mode']
        baud = self.parameters['baud']
        self.parameters['mode'] = None
        self.parameters['baud'] = None
        self.parameters['wpm'] = None

        if mode in SATNOGS_FLOWGRAPH_MODES:
            self.flowgraph_script = SATNOGS_FLOWGRAPH_MODES[mode]['script_name']
            if baud and SATNOGS_FLOWGRAPH_MODES[mode]['has_baudrate']:
                # If this is a CW observation pass the WPM parameter
                if mode == 'CW':
                    self.parameters['wpm'] = str(baud)
                else:
                    self.parameters['baudrate'] = str(baud)
            # Apply framing mode
            if SATNOGS_FLOWGRAPH_MODES[mode]['has_framing']:
                self.parameters['framing'] = SATNOGS_FLOWGRAPH_MODES[mode]['framing']
        else:
            self.flowgraph_script = SATNOGS_FLOWGRAPH_MODES[SATNOGS_FLOWGRAPH_MODE_DEFAULT]['script_name']

    def __del__(self):
        """
        Stop all subprocesses
        """
        self.__terminate_subprocesses()

    def __run_flowgraph(self):
        """
        Starts the flowgraph subprocess

        Starts the flowgraph script already assigned to the FlowGraph
        class from the constructor function. It uses a non-blocking call
        to Popen.
        Exits on OSError with exit message.
        """
        args = [self.flowgraph_script]
        for parameter, value in self.parameters.items():
            if value is not None:
                args.append('--{}={}'.format(parameter, value))

        try:
            self.flowgraph_process = subprocess.Popen(args)
        except OSError as error:
            self.__terminate_subprocesses()
            sys.exit("\033[31mError starting flowgraph: {}\nAborting Observation!\033[0m"
                     .format(error))

    def __run_zmq_sub(self):
        """
        Starts the ZMQ subcriber subprocess

        Starts the flowgraph script already assigned to the FlowGraph
        class from the constructor function. It uses a non-blocking call
        to Popen.
        Exits on OSError with exit message.
        """
        args = [ZMQ_SUB_SCRIPT_NAME]
        for parameter, value in self.zmq_parameters.items():
            if value is not None:
                args.append('--{}={}'.format(parameter, value))
        try:
            self.zmq_sub_process = subprocess.Popen(args)
        except OSError as error:
            self.__terminate_subprocesses()
            sys.exit("\033[31mError starting ZMQ subscriber: {}\nAborting Observation!\033[0m"
                     .format(error))

    def __terminate_flowgraph(self):
        """
        Terminates the flowgraph subprocess

        Send a SIGINT if the process exists and wait up to two seconds
        to terminate. If the process has not terminated in the two
        seconds send SIGKILL singal to kill the process, and print some
        warning
        """
        if ((self.flowgraph_process is not None)
                and (self.flowgraph_process.poll() is None)):
            self.flowgraph_process.send_signal(signal.SIGINT)
            try:
                self.flowgraph_process.wait(2)
            except subprocess.TimeoutExpired:
                self.flowgraph_process.send_signal(signal.SIGKILL)
                print("\033[31mFlowgraph script did not close properly!")
                print("Some frames might be missing\033[0m")

    def __terminate_zmq_subscriber(self):
        """
        Terminates the ZMQ subscriber subprocess

        Send a SIGINT if the process exists and wait up to two seconds
        to terminate. If the process has not terminated in the two
        seconds send SIGKILL singal to kill the process, and print some
        warning
        """
        if ((self.zmq_sub_process is not None)
                and (self.zmq_sub_process.poll() is None)):
            self.zmq_sub_process.send_signal(signal.SIGINT)
            try:
                self.zmq_sub_process.wait(2)
            except subprocess.TimeoutExpired:
                self.zmq_sub_process.send_signal(signal.SIGKILL)
                print("\033[31mZMQ subscriber script did not close properly!")
                print("Some frames might be missing\033[0m")

    def __terminate_subprocesses(self):
        """
        Terminates the subprocesses

        Calls the respective functions to terminate the subprocesses.
        See :func:`__terminate_flowgraph`
        See :func:`__terminate_zmq_subscriber`
        """
        self.__terminate_flowgraph()
        self.__terminate_zmq_subscriber()

    def run(self):
        """
        Starts all necessary subprocesses

        Calls the specific run function for each needed subprocess. ZMQ
        subscriber must be started before the flowgraph to not miss any
        messages.
        See :func:`__run_zmq_sub`
        See :func:`__run_flowgraph`
        """
        self.__run_zmq_sub()
        self.__run_flowgraph()

    def signal_handler(self, sig, frame):
        """
        Singal handler for SIGINT

        If SIGINT is received try to stop the flowgraph and then the
        ZMQ subscriber script
        """
        self.__terminate_subprocesses()


def client_argument_parser():
    """Initialize argument parser for satnogs-client arguments."""
    parser = ArgumentParser()
    parser.add_argument("--antenna",
                        help="Set SoapySDR device antenna name for RX")
    parser.add_argument("--baud",
                        help="Baudrate or wpm for CW flowgraph")
    parser.add_argument("--bb-freq",
                        help="Set SoapySDR device baseband CORDIC frequency",
                        type=float,
                        dest="bb-freq")
    parser.add_argument("--bw",
                        help="Set SoapySDR device RF filter bandwidth",
                        type=float)
    parser.add_argument("--dc-removal",
                        help="Set SoapySDR automatic DC removal",
                        type=int,
                        dest="dc-removal")
    parser.add_argument("--decoded-data-file-path",
                        help="Set decoded frames filepath",
                        dest="decoded-data-file-path")
    parser.add_argument("--dev-args",
                        help="Set SoapySDR device arguments",
                        dest="dev-args")
    parser.add_argument("--doppler-correction-per-sec",
                        help="Set number of Doppler corrections per sec",
                        type=int,
                        dest="doppler-correction-per-sec")
    parser.add_argument("--enable-iq-dump",
                        help="Enable/Disable iq data dumps",
                        type=int,
                        dest="enable-iq-dump")
    parser.add_argument("--file-path",
                        help="Set the audio recordings file path",
                        dest="file-path")
    parser.add_argument("--framing",
                        help="For compatibility with client args. NOT USED!")
    parser.add_argument("--gain",
                        help="Set SoapySDR device overall gain",
                        type=float)
    parser.add_argument("--gain-mode",
                        help="Set SoapySDR device gain mode",
                        dest="gain-mode")
    parser.add_argument("--iq-file-path",
                        help="Set the iq dumps file path",
                        dest="iq-file-path")
    parser.add_argument("--lo-offset",
                        help="Set SoapySDR local oscillator offset",
                        dest="lo-offset")
    parser.add_argument("--mode", help="The modulation scheme")
    parser.add_argument("--other-settings",
                        help="Set SoapySDR channel other settings",
                        dest="other-settings")
    parser.add_argument(
        "--ppm",
        help="Set SoapySDR device oscillator frequency correction in ppm",
        type=float)
    parser.add_argument("--rigctl-host",
                        help="Set Hamlib rotctld host IP",
                        dest="rigctl-host")
    parser.add_argument("--rigctl-port",
                        help="Set Hamlib rigctld port",
                        type=int,
                        dest="rigctl-port")
    parser.add_argument("--rx-freq",
                        help="Set SoapySDR device center frequency",
                        type=float,
                        dest="rx-freq")
    parser.add_argument("--samp-rate-rx",
                        help="Set SoapySDR device sampling rate",
                        type=float,
                        dest="samp-rate-rx")
    parser.add_argument("--soapy-rx-device",
                        help="Set SoapySDR device options",
                        dest="soapy-rx-device")
    parser.add_argument("--stream-args",
                        help="Set SoapySDR device stream args",
                        dest="stream-args")
    parser.add_argument("--tune-args",
                        help="Set SoapySDR device tune args",
                        dest="tune-args")
    parser.add_argument("--udp-dump-host",
                        help="Set UDP Doppler corrected destination IP",
                        dest="udp-dump-host")
    parser.add_argument("--udp-dump-port",
                        help="Set UDP Doppler corrected destination port",
                        type=int,
                        dest="udp-dump-port")
    parser.add_argument("--waterfall-file-path",
                        help="Set waterfall data file path",
                        dest="waterfall-file-path")

    return parser


def main():
    #####################################
    # Argument parser
    ####################################
    args = client_argument_parser().parse_args()

    flowgraph_d = FlowgraphDispatcher(args)
    signal.signal(signal.SIGINT, flowgraph_d.signal_handler)
    flowgraph_d.run()

    # Spinlock on flowgraph process
    # Might be changed with more elaborate multithreading
    while (flowgraph_d.flowgraph_process.poll() is None):
        time.sleep(0.25)


if __name__ == '__main__':
    main()
