# Copyright (c) 2021 Satlab A/S <satlab@satlab.com>
# SPDX-License-Identifier: GPL-3.0

import json
import threading
import xmlrpc.server

from functools import partial

import numpy as np

from gnuradio import gr, blocks, analog, fft, uhd, zeromq, satlab
from gnuradio.fft import logpwrfft


class SRSRadio(gr.top_block):
    def __init__(self,
                 rx_freq,
                 rx_bw,
                 rx_size,
                 rx_gain,
                 tx_freq,
                 tx_bw,
                 tx_size,
                 tx_gain,
                 tun,
                 udp_bind,
                 udp_port,
                 clock_source,
                 gpio_control,
                 xmlrpc_enable,
                 xmlrpc_bind,
                 xmlrpc_port,
                 zmq_enable,
                 zmq_port,
                 pdu_logger,
                 fft_websocket_enable,
                 fft_collect_enable,
                 init_gated,
                 tsb_tag_name):

        # Call parent constructor
        gr.top_block.__init__(self, 'radio')

        # Variables
        self.tx_gain = tx_gain
        self.tx_freq = tx_freq
        self.tx_doppler = 0
        self.rx_gain = rx_gain
        self.rx_freq = rx_freq
        self.rx_doppler = 0

        # Variables
        uhd_device = 'num_send_frames=32,num_recv_frames=256'
        fft_size = 2048

        # USRP blocks
        self.usrp_source = uhd.usrp_source(
            ','.join((uhd_device, '')),
            uhd.stream_args(
                cpu_format='fc32',
                args='',
                channels=[0],
            ),
        )
        self.usrp_source.set_clock_source(clock_source, 0)
        self.usrp_source.set_center_freq(rx_freq, 0)
        self.usrp_source.set_rx_agc(False, 0)
        self.usrp_source.set_gain(rx_gain, 0)
        self.usrp_source.set_antenna('RX2', 0)
        self.usrp_source.set_bandwidth(rx_bw, 0)
        self.usrp_source.set_samp_rate(self.samp_rate)
        self.usrp_source.set_clock_rate(self.clock_rate, uhd.ALL_MBOARDS)

        self.usrp_sink = uhd.usrp_sink(
            ','.join((uhd_device, '')),
            uhd.stream_args(
                cpu_format='fc32',
                args='',
                channels=[0],
            ),
            tsb_tag_name
        )
        self.usrp_sink.set_clock_source(clock_source, 0)
        self.usrp_sink.set_center_freq(tx_freq, 0)
        self.usrp_sink.set_gain(tx_gain, 0)
        self.usrp_sink.set_antenna('TX/RX', 0)
        self.usrp_sink.set_bandwidth(tx_bw, 0)
        self.usrp_sink.set_samp_rate(self.samp_rate)
        self.usrp_sink.set_clock_rate(self.clock_rate, uhd.ALL_MBOARDS)

        # Configure RX/TX GPIOs. TX on FPGA_TXD0 (GPIO 8), RX on FPGA_RXD0 (GPIO 9)
        if gpio_control:
            gpio_tx = 1 << 8
            gpio_rx = 1 << 9
            gpio_mask = gpio_tx | gpio_rx
            self.usrp_sink.set_gpio_attr(  # Both ATR
                'FP0', 'CTRL', gpio_tx | gpio_rx, gpio_mask)
            self.usrp_sink.set_gpio_attr(  # Both output
                'FP0', 'DDR', gpio_tx | gpio_rx, gpio_mask)
            self.usrp_sink.set_gpio_attr(  # Both off normal
                'FP0', 'ATR_0X', 0, gpio_mask)
            self.usrp_sink.set_gpio_attr(  # TX on in TX
                'FP0', 'ATR_TX', gpio_tx, gpio_mask)
            self.usrp_sink.set_gpio_attr(  # RX on in RX
                'FP0', 'ATR_RX', gpio_rx, gpio_mask)
            self.usrp_sink.set_gpio_attr(  # Both on in RXTX
                'FP0', 'ATR_XX', gpio_tx | gpio_rx, gpio_mask)

        # Doppler compensation mixers and LOs
        self.tx_doppler_mixer = blocks.multiply_vcc(1)
        self.rx_doppler_mixer = blocks.multiply_vcc(1)
        self.rx_gate = blocks.copy(gr.sizeof_gr_complex)
        self.tx_gate = blocks.copy(gr.sizeof_gr_complex)
        self.tx_doppler_lo = analog.sig_source_c(self.samp_rate, analog.GR_COS_WAVE, self.tx_doppler, 1, 0, 0)
        self.rx_doppler_lo = analog.sig_source_c(self.samp_rate, analog.GR_COS_WAVE, self.rx_doppler, 1, 0, 0)

        # Set initial gated state
        self.rx_gate.set_enabled(not init_gated)
        self.tx_gate.set_enabled(not init_gated)

        # UDP socket for CSP frames
        self.socket = blocks.socket_pdu('UDP_SERVER', udp_bind, str(udp_port), 1500, False)

        # ZMQ socket for raw samples
        if zmq_enable:
            self.zmq_samples = zeromq.pub_sink(gr.sizeof_gr_complex, 1, f'tcp://*:{zmq_port}', 100, False, -1)
            self.connect((self.rx_doppler_mixer, 0), (self.zmq_samples, 0))

        # Create FFT block
        if fft_websocket_enable or fft_collect_enable:
            self.logpwrfft = logpwrfft.logpwrfft_c(
                sample_rate=self.samp_rate,
                fft_size=fft_size,
                ref_scale=1,
                frame_rate=25,
                avg_alpha=1.0,
                average=False,
                # Build in spectrum rotator by fs/2 in window to avoid shift
                win=lambda l: fft.window.blackmanharris(l) * np.tile([1, -1], l // 2)
            )
            self.connect((self.rx_doppler_mixer, 0), self.logpwrfft)

            # Create websocket spectrum
            if fft_websocket_enable:
                self.ws_sink = satlab.websocket_sink(gr.sizeof_float * fft_size, 8081)
                self.ws_sink.set_greeting(
                    json.dumps({
                        'center': self.rx_freq,
                        'span': self.samp_rate
                    })
                )
                self.connect(self.logpwrfft, self.ws_sink)

            # Create FFT collector
            if fft_collect_enable:
                self.fft_collector = satlab.RowCollector(row_size=fft_size)
                self.connect(self.logpwrfft, self.fft_collector)

        # Connections
        self.connect((self.usrp_source, 0), (self.rx_doppler_mixer, 1))
        self.connect((self.rx_doppler_lo, 0), (self.rx_doppler_mixer, 0))

        self.connect((self.tx_doppler_lo, 0), (self.tx_doppler_mixer, 0))
        self.connect((self.tx_doppler_mixer, 0), (self.tx_gate, 0))
        self.connect((self.tx_gate, 0), (self.usrp_sink, 0))

        # Create packet loggers
        self.pdulog = satlab.PDULogger()

        # Connect transmitter chain
        self.connect((self.transmitter, 0), (self.tx_doppler_mixer, 1))
        self.msg_connect((self.socket, 'pdus'), (self.transmitter, 'csp_in'))

        # Connect receiver chain
        self.connect((self.rx_doppler_mixer, 0), (self.rx_gate, 0))
        self.connect((self.rx_gate, 0), (self.receiver, 0))
        self.msg_connect((self.receiver, 'csp_out'), (self.socket, 'pdus'))

        # Connect packet logger
        self.msg_connect((self.receiver, 'csp_out'), (self.pdulog, 'rx'))
        self.msg_connect((self.socket, 'pdus'), (self.pdulog, 'tx'))

        # Connect TUN interface
        if tun is not None:
            self.tuntap = blocks.tuntap_pdu(tun, tx_size, True)
            self.msg_connect((self.tuntap, 'pdus'), (self.transmitter, 'ip_in'))
            self.msg_connect((self.receiver, 'ip_out'), (self.tuntap, 'pdus'))

        # Direct function mappings
        self.get_rx_frequency = self.usrp_source.get_center_freq
        self.get_rx_gain = partial(self.usrp_source.get_gain, 0)
        self.get_tx_frequency = self.usrp_sink.get_center_freq
        self.get_tx_gain = partial(self.usrp_sink.get_gain, 0)

        # Get setters and getters in self.rx and self.transmitter and map
        # them to self.[get|set]_[rx|tx]_[attribute]
        for rxtx in [(self.receiver, 'rx'), (self.transmitter, 'tx')]:
            for getset in ['get', 'set']:
                for attr in [x for x in dir(rxtx[0]) if callable(getattr(rxtx[0], x)) and x.startswith(f'{getset}_')]:
                    f = f'{getset}_{rxtx[1]}_{attr[4:]}'
                    if not hasattr(self, f):
                        setattr(self, f, getattr(rxtx[0], attr))

        # Start XMLRPC server
        if xmlrpc_enable:
            self.xmlrpc_server = xmlrpc.server.SimpleXMLRPCServer((xmlrpc_bind, xmlrpc_port),
                                                                  allow_none=True, logRequests=False)
            # Register setters and getters
            for func in [x for x in dir(self) if callable(getattr(self, x)) and (x.startswith('get_') or x.startswith('set_'))]:
                if func not in ('get_receiver', 'get_transmitter'):  # not intended for remote
                    self.xmlrpc_server.register_function(getattr(self, func), func)
            self.xmlrpc_server.register_introspection_functions()
            self.xmlrpc_server_thread = threading.Thread(target=self.xmlrpc_server.serve_forever)
            self.xmlrpc_server_thread.daemon = True
            self.xmlrpc_server_thread.start()
            print(f'Serving XML-RPC on {xmlrpc_bind}:{xmlrpc_port}')

    def get_receiver(self):
        return self.receiver

    def get_transmitter(self):
        return self.transmitter

    def set_tx_gated(self, gated):
        self.tx_gate.set_enabled(not gated)

    def set_tx_gain(self, gain):
        self.usrp_sink.set_gain(gain, 0)

    def set_tx_freq(self, freq):
        self.usrp_sink.set_center_freq(freq, 0)

    def set_tx_doppler(self, tx_doppler):
        self.tx_doppler_lo.set_frequency(tx_doppler)

    def set_rx_gated(self, gated):
        self.rx_gate.set_enabled(not gated)

    def set_rx_gain(self, gain):
        self.usrp_source.set_gain(gain, 0)

    def set_rx_freq(self, rx_freq):
        self.usrp_source.set_center_freq(rx_freq, 0)
        if hasattr(self, 'ws_sink'):
            self.ws_sink.set_greeting(
                json.dumps({
                    'center': rx_freq,
                    'span': self.samp_rate
                })
            )

    def set_rx_doppler(self, rx_doppler):
        self.rx_doppler_lo.set_frequency(rx_doppler)
