from gnuradio import analog
from gnuradio import blocks
from gnuradio import digital
from gnuradio import filter
from gnuradio import gr
from gnuradio import soapy
from gnuradio import uhd
from gnuradio import zeromq
from gnuradio import fec
import gnuradio.satnogs as satnogs

import math
import Hamlib

from gnuradio import network

class FlowgraphBuilder(gr.top_block):
    """
    Generate SatNOGS Flowgraphs.
    This class instantiates necessary blocks according to arguments passed
    to it and then connects them accordingly to create the final top block
    flowgraph.
    It provides a function to run the created flowgraph.
    """

    def __init__(self, radio_config, transmitter, flowgraph_config, rx):
        flowgraph_name = transmitter.name
        if (rx):
            flowgraph_name += " receiver"
        else:
            flowgraph_name += " transmitter"
        gr.top_block.__init__(self, flowgraph_name, catch_exceptions=True)
        self._running = False

        # Instantiate correct SDR device
        if (rx):
            self.sdr_device = self.__init_sdr_source(radio_config)
            uri = flowgraph_config.zmq_command_pull_msg_uri
            timeout_us = flowgraph_config.zmq_command_pull_msg_timeout_us
            self.zmq_command_pull_msg = self.__init_zmq_pull_msg(uri,
                                                                 timeout_us)
            # Connect command pull with SDR
            self.msg_connect((self.zmq_command_pull_msg, 'out'),
                             (self.sdr_device, 'cmd'))
            # Init ZMQ msg push for frames
            uri = flowgraph_config.zmq_frame_pub_msg_uri
            timeout_us = flowgraph_config.zmq_frame_pub_msg_timeout_us

            self.zmq_frame_pub_msg = zeromq.pub_msg_sink(uri, timeout_us)

            self.udp_sink = network.udp_sink(gr.sizeof_gr_complex, 1, '127.0.0.1', 1234, 0, 1472, False)
            self.connect((self.sdr_device, 0), (self.udp_sink, 0))
        else:
            self.sdr_device = self.__init_sdr_sink(radio_config)

        # Switch based on transmitter name
        # Only supports phasma
        transmitter_name = transmitter.name
        match transmitter_name.lower():
            case "phasma":
                if (rx):
                    # Initialize Decoder and Scrambler objects
                    self.decoder_definition = self.__init_decoder_definition(transmitter,
                                                                             flowgraph_config)
                    self.scrambler_obj = self.__init_scrambler_object(transmitter,
                                                                      flowgraph_config)
                    # Init Doppler compensation
                    self.doppler_compensation = ( 
                        self.__init_doppler_compensation(radio_config,
                                                         flowgraph_config,
                                                         transmitter)
                    )
                    # Connect SDR with Doppler Compensation
                    self.connect((self.sdr_device, 0), (self.doppler_compensation, 0))

                    # Init quadrature Demod
                    self.quad_demod = self.__init_quad_demod(radio_config,
                                                             flowgraph_config,
                                                             transmitter)
                    # Connect Doppler Compensation with Quadrature Demod
                    self.connect((self.doppler_compensation, 0), (self.quad_demod, 0))
                    
                    # Init Symbol Sync
                    self.symbol_sync = self.__init_symbol_sync(radio_config,
                                                               flowgraph_config,
                                                               transmitter)
                    # Connect Quadrature Demod  with Symbol Sync
                    self.connect((self.quad_demod, 0), (self.symbol_sync, 0))
                    
                    # Init binary slicer
                    self.binary_slicer = digital.binary_slicer_fb()
                    # Connect Symbol Sync with Binary Slicer
                    self.connect((self.symbol_sync, 0), (self.binary_slicer, 0))
                    
                    # Init Frame Decoder
                    self.frame_decoder = satnogs.frame_decoder(self.decoder_definition,
                                                               1*1)
                    # Connect Binary Slicer to Frame Decoder
                    self.connect((self.binary_slicer, 0), (self.frame_decoder, 0))

                    last_block = self.frame_decoder
                    
                    # Init FEC Async Encoder
                    self.fec_encoder = self.__init_fec_encoder(radio_config,
                                                               flowgraph_config,
                                                               transmitter)
                    if self.fec_encoder is not None:
                        self.msg_connect((self.frame_decoder, 'out'), (self.fec_encoder, 'in'))
                        last_block = self.fec_encoder
                    
                    # Init CRC Async Check
                    self.crc_check = self.__init_crc_check(radio_config,
                                                           flowgraph_config,
                                                           transmitter)
                    if self.crc_check is not None:
                        self.msg_connect((last_block, 'out'), (self.crc_check, 'in'))
                        self.msg_connect((self.crc_check, 'ok'),
                                         (self.zmq_frame_pub_msg, 'in'))
                    else:
                        pass
 

    def __init_sdr_source(self, radio_config):
        """
        Instantiates correct SDR device based on the device driver passed to the
        radio_config object

        @return sdr_source: the sdr_source_object
        """
        # UHD
        if ("uhd" in radio_config.dev_driver):
            spec = "A:A" # Hardcoded, check if needs to be a parameter

            sdr_source = uhd.sdr_source(radio_config.dev_args,
                                         uhd.stream_args(
                                            cpu_format=radio_config.output_stream_format,
                                            args=radio_config.stream_args,
                                            channels=list(range(radio_config.channels))
                                        ))
            sdr_source.set_subdev_spec(spec, 0)
            sdr_source.set_samp_rate(radio_config.samp_rate)
            sdr_source.set_time_uknown_pps(uhd.time_spec(0))
            sdr_source.set_center_freq(radio_config(freq))
            sdr_source.set_antenna(radio_config.antenna)
            sdr_source.set_gain(radio_config.gain)
        elif ("rtlsdr" in radio_config.dev_driver):
            sdr_source = soapy.source(device="driver=rtlsdr",
                                      type=radio_config.output_stream_format,
                                      nchan=radio_config.channels,
                                      dev_args=radio_config.dev_args,
                                      stream_args=radio_config.stream_args,
                                      tune_args=radio_config.tune_args.split(','),
                                      other_settings=radio_config.other_settings.split(','))
            sdr_source.set_sample_rate(0, radio_config.samp_rate)
            sdr_source.set_frequency(0, radio_config.freq)
            sdr_source.set_frequency_correction(0, radio_config.ppm)
            sdr_source.set_gain(0, radio_config.gain)

        return sdr_source

    def __init_zmq_pull_msg(self, uri, timeout_us):
        return zeromq.pull_msg_source(uri, timeout_us, False)

    def __init_doppler_compensation(self, radio_config, flowgraph_config, transmitter):
        # Init rigctl variable
        modulation_str = transmitter.transmission_mode[0].modulation.WhichOneof('modulation')
        modulation = getattr(transmitter.transmission_mode[0].modulation,
                             modulation_str)
        if modulation.HasField('baudrate'):
            baudrate = modulation.baudrate
        else:
            baudate = flowgraph_config.audio_samp_rate
        self.rig = satnogs.rigctl_doppler_correction.make(radio_config.freq,
                                                      Hamlib.RIG_MODEL_NETRIGCTL)
        self.rig.set_param("rig_pathname", flowgraph_config.rigctl_ip +
                      ":" + str(flowgraph_config.rigctl_port))
        samp_rate = radio_config.samp_rate
        rx_freq = radio_config.freq
        lo_offset = flowgraph_config.lo_offset
        # target_samp_rate = flowgraph_config.sps * baudrate# samp_rate
        target_samp_rate = 200e3
        compensate = True

        print('samp_rate: ', samp_rate)
        print('rx_freq: ', rx_freq)
        print('lo_offset:', lo_offset)
        print('target_samp: ', target_samp_rate)
        print('compensate: ', compensate)

        return satnogs.doppler_compensation(self.rig, samp_rate, rx_freq,
                                            lo_offset, target_samp_rate,
                                            compensate)

    def __init_quad_demod(self, radio_config, flowgraph_config, transmitter):
        modulation_str = (transmitter.transmission_mode[0].modulation.
                          WhichOneof('modulation'))
        modulation = getattr(transmitter.transmission_mode[0].modulation,
                             modulation_str)
        baudrate = modulation.baudrate 
        sps = flowgraph_config.sps
        gain = (2 * math.pi)/sps

        return analog.quadrature_demod_cf(gain)

    def __init_symbol_sync(self, radio_config, flowgraph_config, transmitter):
        sps = flowgraph_config.sps
        error_detector = digital.TED_GARDNER
        ted_gain = 1.0
        loop_bandwidth = 2 * math.pi / 200
        damping_factor = 1 / math.sqrt(2)
        max_deviation = 0.1
        out_sps_per_symbol = 1
        interpolating_resampler = digital.IR_PFB_NO_MF
        filterbank_arms = 128

        return digital.symbol_sync_ff(detector_type=error_detector,
                                      sps=sps,
                                      loop_bw=loop_bandwidth,
                                      damping_factor=damping_factor,
                                      max_deviation=max_deviation,
                                      interp_type=interpolating_resampler)

    def __init_decoder_definition(self, transmitter, flowgraph_config):
        decoder_str = (transmitter.transmission_mode[0].framing.
                       WhichOneof('framing'))
        decoder = getattr(transmitter.transmission_mode[0].framing,
                          decoder_str)
        preamble = list(decoder.preamble) * decoder.preamble_length
        preamble_threshold = flowgraph_config.preamble_threshold
        sync_word = list(decoder.sync_word)
        sync_word_threshold = flowgraph_config.sync_word_threshold
        max_frame_len = decoder.max_frame_length
        
        return satnogs.ieee802_15_4_variant_decoder.make(preamble=preamble,
                                                         preamble_threshold=preamble_threshold,
                                                         sync=sync_word,
                                                         sync_threshold=sync_word_threshold,
                                                         max_len=max_frame_len)

    def __init_scrambler_object(self, transmitter, flowgraph_config):
        coder_str = (transmitter.transmission_mode[0].framing.WhichOneof('framing'))
        coder = getattr(transmitter.transmission_mode[0].framing,
                        coder_str)
        scrambler = coder.generic_options.scrambler
        mask = scrambler.mask
        seed = scrambler.seed
        order = scrambler.order
        msb = scrambler.msb
        self_synchronizing = scrambler.self_synchronizing

        return satnogs.scrambler.make(mask=mask,
                                      seed=seed,
                                      order=order,
                                      msb=msb,
                                      self_sync=self_synchronizing)

    def __init_fec_encoder(self, radio_config, flowgraph_config, transmitter):
        coder_str = (transmitter.transmission_mode[0].framing.WhichOneof('framing'))
        coder = getattr(transmitter.transmission_mode[0].framing,
                        coder_str)
        frame_len = coder.max_frame_length
        if coder.generic_options.HasField('scrambler'):
            return fec.async_encoder(self.scrambler_obj, True, False, False, frame_len + 4)
        
        return None

    def __init_crc_check(self, radio_config, flowgraph_config, transmitter):
        coder_str = (transmitter.transmission_mode[0].framing.WhichOneof('framing'))
        coder = getattr(transmitter.transmission_mode[0].framing,
                       coder_str)
        crc = coder.generic_options.crc
        if coder.generic_options.HasField('crc'):
            return digital.crc_check(num_bits=crc.size_in_bits,
                                     poly=int.from_bytes(crc.polynomial),
                                     initial_value=int.from_bytes(crc.init_value),
                                     final_xor=int.from_bytes(crc.final_xor),
                                     input_reflected=crc.reflect_input,
                                     result_reflected=crc.reflect_output,
                                     swap_endianness=False,
                                     discard_crc=False)

        else:
            return None


    @property
    def running(self):
        """Property indicating whether the flowgraph is running"""
        return self._running

    @running.setter
    def running(self, running):
        """Sets the running property"""
        self._running = running

    def run(self):
        """
        Starts the flowgraph and sets the running property to True

        @throws RuntimeError if flowgraph already running or cannot start
        """
        if (self._running):
            raise RuntimeError("Flowgraph already running")
        try:
            gr.top_block.start(self)
        except RuntimeError as error:
            raise error
        self._running = True
        # gr.top_block.wait(self)

    def stop(self):
        """
        Stops the flowgraph if already running. If not do nothing.
        """
        if (self._running):
            gr.top_block.stop(self)
            gr.top_block.wait(self)
            self._running = False
