from distutils.version import StrictVersion
import multiprocessing as mp
import queue
import signal
import sys
import datetime
import time
from typeguard import typechecked
from typing import (
    Type, TypeVar, Generic, Callable,
    Protocol, Tuple, Any, Iterable,
    List, Literal, runtime_checkable, TYPE_CHECKING,
    Union, Optional,
)



T = TypeVar("T")

Tgr = TypeVar("Tgr", bound="gr.top_block")
"""A gr top block or a subclass of such."""


if TYPE_CHECKING:
    from typing_extensions import ParamSpec, Concatenate, TypeAlias
    from gnuradio import gr  # type: ignore[import-untyped]
    from turtle import Turtle
    
    P = ParamSpec("P")

    TPFunc = Callable[Concatenate[T, P], None]
    """A function which takes at least one argument.
    In this particular project, the first argument is most often 
    an instance of `gr.top_block` or any subclass.

    Examples:

    ```python
    def f(tb) -> None: ...
    def f(tb, a) -> None: ...
    def f(tb, a, b) -> None: ...
    def f(tb, a, b, c) -> None: ...
    # etc...
    ```
    """

    _cmdqueue = mp.Queue[Tuple[TPFunc[T, P], Iterable[P.args]]]
    """A queue of tuples. Each tuple is (func, args), 
    where the function's first argument is of type `T` (generic),
    and the functions other args are of "the correct type". Example:
    ```python3
    def f(a: list, b: str, c: int):
        ...
    
    q = mp.Queue()
    args = ("foo", 33)
    q.put((f, args))
    ```
    """


@runtime_checkable
class _AboutToQuitAttr(Protocol):
    def connect(self, f: Callable[[], Any]) -> Any: ...

@runtime_checkable
class _QAppProt(Protocol):
    def exec_(self) -> Any: ...
    aboutToQuit: _AboutToQuitAttr


class ProcessTerminated(RuntimeError):
    ...
    

class UnableToLaunch(RuntimeError):
    ...



def _grc_main_prep(top_block_cls: "Type[Tgr]") -> "Tuple[Tgr, _QAppProt]":
    """This is a copy/paste of the main() function that is generated
    by GRC for Graphical (Qt) flowgraphs. It omits the last line
    of that main function, `qapp.exec_()`, because it is blocking.
    The developer using this function is responsible for running `.exec_()`.
    Example usage 1:
    ```python
    tb, qapp = _grc_main_prep(a_specific_gr_top_block)
    qapp.exec_()
    ```
    Example usage 2:
    ```python
    tb, qapp = _grc_main_prep(a_specific_gr_top_block)
    ## hypothetically, if your `tb` has a signal source block
    tb.signal_source.set_frequency(101.3e6)
    qapp.exec_()
    ```
    """
    from PyQt5 import Qt     # type: ignore[import-untyped]
    from gnuradio import gr

    if StrictVersion("4.5.0") <= StrictVersion(Qt.qVersion()) < StrictVersion("5.0.0"):
        style = gr.prefs().get_string('qtgui', 'style', 'raster')
        Qt.QApplication.setGraphicsSystem(style)
    qapp = Qt.QApplication(sys.argv)

    tb = top_block_cls()
    tb.start()
    tb.show()

    def sig_handler(sig=None, frame=None) -> None:  # type: ignore[no-untyped-def]
        Qt.QApplication.quit()

    signal.signal(signal.SIGINT, sig_handler)
    signal.signal(signal.SIGTERM, sig_handler)

    timer = Qt.QTimer()
    timer.start(500)
    timer.timeout.connect(lambda: None)

    def quitting() -> None:
        tb.stop()
        tb.wait()
    qapp.aboutToQuit.connect(quitting)
    assert isinstance(tb, gr.top_block)
    assert isinstance(qapp, _QAppProt)
    assert isinstance(qapp.aboutToQuit, _AboutToQuitAttr)
    return tb, qapp


def _processcmds(instance: "T", q: "_cmdqueue[T, P]") -> None:
    """Pulls a command from the queue `q`, and calls it
    with `instance` as the first argument. `instance` is often a `gr.top_block`
    (which is conventionally named `tb` in the .py files generated by GRC.)"""
    try:
        cmd, args = q.get(block=False)
        cmd(instance, *args)
    except queue.Empty:
        pass


def _event_loop_gr(top_block_cls: "Type[Tgr]", q: "_cmdqueue[Tgr, P]") -> None:
    """Run the Qt/GR process.
    - Sets up the top_block similarly to the `main()` function in the .py file that GRC generates
    - Starts a repeating Qt timer to get and execute commands from `q`
    - Runs the Qt application (assumes that `top_block_cls` is a Qt-based GR flowgraph)"""
    from PyQt5 import Qt
    tb, qapp = _grc_main_prep(top_block_cls)
    timer = Qt.QTimer()
    timer.start(1)
    timer.timeout.connect(lambda: _processcmds(tb, q))
    qapp.aboutToQuit.connect(lambda: print("Gracefully exiting GRC flowgraph."))
    qapp.exec_()


class ParallelGR(Generic[Tgr]):
    def __init__(self, top_block_cls: "Type[Tgr]") -> None:
        self.__q: "_cmdqueue[Tgr, ...]" = mp.Queue()
        self.__proc = mp.Process(target=lambda: _event_loop_gr(top_block_cls, self.__q))

    def start(self) -> None:
        self.__proc.start()
    
    def terminate(self) -> None:
        self.__proc.terminate()

    def put_cmd(self, f: "TPFunc[Tgr, P]", *args: "P.args", **kwargs: "P.kwargs") -> None:
        """Put a command into the queue for the child
        process to execute. Any extra args provided will be
        passed to `f`."""
        if not self.__proc.is_alive():
            raise ProcessTerminated("The parallel process has terminated; cannot execute commands.")
        self.put_cmd_nolivecheck(f, *args, **kwargs)
    
    def put_cmd_nolivecheck(self, f: "TPFunc[Tgr, P]", *args: "P.args", **kwargs: "P.kwargs") -> None:
        """Same as `put_cmd`, but doesn't check that the process is alive."""
        self.__q.put((f, args))
    
    def is_alive(self) -> bool:
        return self.__proc.is_alive()


##### commands for student use

class hasPGR(Protocol):
    _pgr: ParallelGR

class hasCI(Protocol):
    _ci: "hasPGR | None"


if TYPE_CHECKING:
    from .specan.specan_gr3_8 import specan_gr3_8
    from .specan.specan_gr3_10 import specan_gr3_10
    from .wbfm_rx.wbfm_rx_gr3_8 import wbfm_rx_gr3_8
    from .wbfm_rx.wbfm_rx_gr3_10 import wbfm_rx_gr3_10
    from .noise_tx.noise_tx_gr3_8 import noise_tx_gr3_8
    from .noise_tx.noise_tx_gr3_10 import noise_tx_gr3_10
    from .psk_tx_loop.psk_tx_loop_gr3_8 import psk_tx_loop_gr3_8
    _SpecAn: TypeAlias = Union[specan_gr3_8, specan_gr3_10]
    _Noise_Tx: TypeAlias = Union[noise_tx_gr3_8, noise_tx_gr3_10]    
    _WBFM_Rx: TypeAlias = Union[wbfm_rx_gr3_8, wbfm_rx_gr3_10]
    _PSK_Tx: TypeAlias = psk_tx_loop_gr3_8



def _set_center_freq(tb: "_SpecAn", freq: float) -> None:
    tb.set_center_freq(freq)  # type: ignore[no-untyped-call]


def _set_if_gain(tb: "_SpecAn", if_gain: float) -> None:
    tb.set_if_gain(if_gain)


def _set_bb_gain(tb: "_SpecAn", bb_gain: float) -> None:
    tb.set_bb_gain(bb_gain)


def _set_samp_rate(tb: "_SpecAn", samp_rate: float) -> None:
    tb.set_samp_rate(samp_rate)  # type: ignore[no-untyped-call]


def _set_hw_bb_filt(tb: "_WBFM_Rx", val: float) -> None:
    tb.set_hw_bb_filt(val)


def _set_freq_offset(tb: "_WBFM_Rx", freq_offset: float) -> None:
    tb.set_freq_offset(freq_offset)


def _set_channel_width(tb: "_WBFM_Rx", channel_width: float) -> None:
    tb.set_channel_width(channel_width)


def _set_noise_type(tb: "_Noise_Tx", noise_type: str) -> None:
    tb.set_noise_type(noise_type)  # type: ignore[no-untyped-call]


def _set_amplitude(tb: "_Noise_Tx", amplitude: float) -> None:
    tb.set_amplitude(amplitude)  # type: ignore[no-untyped-call]


def _set_filter_cutoff_freq(tb: "_Noise_Tx", filter_cutoff_freq: float) -> None:
    tb.set_filter_cutoff_freq(filter_cutoff_freq)  # type: ignore[no-untyped-call]


def _set_filter_transition_width(tb: "_Noise_Tx", filter_transition_width: float) -> None:
    tb.set_filter_transition_width(filter_transition_width)  # type: ignore[no-untyped-call]


def _set_data(tb: "_PSK_Tx", data: List[int]) -> None:
    tb.set_data(data)  # type: ignore[no-untyped-call]


def _set_modulation(tb: "_PSK_Tx", modulation: "_Modulations") -> None:
    tb.set_modulation(modulation)  # type: ignore[no-untyped-call]


def startnewinstance(cls: hasCI) -> None:
    cls._ci = cls()
    cls._ci._pgr.start()
    time.sleep(0.5)
    if not cls._ci._pgr.is_alive():
        raise UnableToLaunch("Possible fixes: plug in Hack RF; ensure there are no other programs using the Hack RF; press reset button on Hack RF")


# def decidemakenew(cls: hasCI) -> None:
#     if (cls._ci is None) or (not cls._ci._pgr.is_alive()):
#         startnewinstance(cls)
def alive_ish(cls: hasCI) -> bool:
    if cls._ci is None:
        return False
    return cls._ci._pgr.is_alive()


def decidelaunchkill(cls: hasCI, running: bool, setallfunc: Callable) -> None:
    if alive_ish(cls):
        if running:
            setallfunc()
            return { "status": "running", "timestamp": datetime.datetime.now() }
        else:
            assert cls._ci is not None
            cls._ci._pgr.terminate()
            return { "status": "terminated" }
    else:
        if running:
            startnewinstance(cls)
            setallfunc()
            return { "status": "launched", "timestamp": datetime.datetime.now() }
        else:
            return { "status": "not_running" }


class _EXPLANATIONS:
    _introductory = """\n\nIf there is not an instance of this Paragradio app running, launch a new one, and set the settings. If one is already running, update the settings of the existing one. Returns a dictionary containing the timestamp of the update.\n"""
    amplitude = "\n**amplitude**: Update the amplitude of the generated signal. Units are difficult to explain; try setting it to `1` and then adjust higher or lower based on what fits your situation.\n"
    bb_gain = "\n**bb_gain**: Update the Baseband gain of the SDR peripheral. Units are dB.\n"
    center_freq = "\n**center_freq**: Update the center frequency of the SDR peripheral and any associated GUI elements. Units are Hz.\n"
    channel_width = """\n**channel_width**: Update the width of the software bandpass filter. Units are Hz.
        This reduces the interference from nearby stations by filtering to a single radio station. In the United States, Broadcast FM radio stations are 200 kHz wide, so a channel width of 200 kHz is a good option. A slightly narrower width sometimes helps improve the Signal to Noise Ratio.\n"""
    data = "\n**data**: Update what binary data is being repeatedly transmitted. Example: [1, 0, 1, 1] \n"
    filter_cutoff_freq = "\n**filter_cutoff_freq**: Update the cutoff frequency of the filter that shapes the generated noise before transmitting it. Units are Hz.\n"
    filter_transition_width = "\n**filter_transition_width**: Update the transition width of the filter that shapes the generated noise before transmitting it. Units are Hz.\n"
    freq_offset = """\n**freq_offset**: Update the frequency offset. Units are Hz.
        When tuning the FM Radio, you'll often get a clearer sound if
        you tune offset to avoid the DC Spike.\n"""
    hw_bb_filt = """\n**hw_bb_filt**: Update the Hardware Baseband Filter. Units are Hz. 
    Details: The HackRF One and many other SDR peripherals have a built-in filter that precedes the Analog to Digital conversion. It is able to reduce or prevent aliasing, which software filters cannot do. Typically, you should set the baseband filter as low as possible for the signals that you wish to receive -- a tighter filter will more effectively reduce aliasing. For more info, see the Hack RF One documentation about Sampling Rate and Baseband Filters <https://hackrf.readthedocs.io/en/latest/sampling_rate.html>.\n"""
    if_gain = "\n**if_gain**: Update the Intermediate Frequency gain of the SDR peripheral. Units are dB.\n"
    modulation = "\n**modulation**: Update the modulation. Options: 'BPSK', 'QPSK', 'DQPSK', '8PSK', '16QAM'.\n"
    noise_type = """\n**noise_type**: Update the noise type to 'uniform' or 'gaussian'.
    In the underlying GNU Radio noise generator, there are two further options, 'impulse' and 'laplacian'. Those two options do not function (as of 2025 January), therefore paragradio does not provide them.\n"""
    running = "\n**running**: Runs the GNU Radio app if this is set to True. Terminates the app if set to False. \n"
    samp_rate = "\n**samp_rate**: Update the sample rate of the SDR peripheral and the bandwidth (the amount of viewable spectrum) of the GUI spectrum view. Units are samples per second (which, in this context, is roughly the same as Hz).\n"


class SpecAn():
    _ci: "Optional[SpecAn]" = None
    "Current instance"

    def __init__(self) -> None:
        """Create a Paragradio Spectrum Analyzer."""
        from .specan import specan_fg
        self._pgr = ParallelGR(specan_fg)
    
    @classmethod
    def __set_all(cls, center_freq, if_gain, bb_gain, samp_rate, hw_bb_filt):
        cls._ci._pgr.put_cmd(_set_center_freq, center_freq)
        cls._ci._pgr.put_cmd(_set_if_gain, if_gain)
        cls._ci._pgr.put_cmd(_set_bb_gain, bb_gain)
        cls._ci._pgr.put_cmd(_set_samp_rate, samp_rate)
        cls._ci._pgr.put_cmd(_set_hw_bb_filt, hw_bb_filt)

    @classmethod
    def __config(cls, running, center_freq, if_gain, bb_gain, samp_rate, hw_bb_filt):
        def setallfunc():
            cls.__set_all(center_freq, if_gain, bb_gain, samp_rate, hw_bb_filt)
        return decidelaunchkill(cls, running, setallfunc)

    @typechecked
    @staticmethod
    def config(
            *,
            running: bool,
            center_freq: float = 93e6,
            if_gain: int = 24,
            bb_gain: int = 32,
            samp_rate: float = 2e6,
            hw_bb_filt: float = 2.75e6,
        ) -> dict:
        """To view the docs for this method, run `from paragradio.v2025_03 import SpecAn; help(SpecAn)` in your Python editor. If you are using a marimo notebook, you can view the docstring with rich formatting by running this in a cell: `mo.md(SpecAn.config.__doc__)`"""
        return SpecAn.__config(running, center_freq, if_gain, bb_gain, samp_rate, hw_bb_filt)



SpecAn.config.__doc__ += _EXPLANATIONS._introductory + _EXPLANATIONS.center_freq + _EXPLANATIONS.if_gain +_EXPLANATIONS.bb_gain + _EXPLANATIONS.samp_rate + _EXPLANATIONS.hw_bb_filt

class SpecAnSim():

    _ci: "Optional[SpecAnSim]" = None
    "Current instance"

    def __init__(self) -> None:
        """Create a Paragradio Simulated Spectrum Analyzer with simulated activity on 93.5 MHz.
        """
        from .specansim import specansim_fg
        self._pgr = ParallelGR(specansim_fg)

    @classmethod
    def __set_all(cls, center_freq):
        cls._ci._pgr.put_cmd(_set_center_freq, center_freq)
    
    @classmethod
    def __config(cls, running, center_freq):
        def setallfunc():
            cls.__set_all(center_freq)
        return decidelaunchkill(cls, running, setallfunc)

    @typechecked
    @staticmethod
    def config(
            *,
            running: bool,
            center_freq: float = 93e6,
        ) -> dict:
        """To view the docs for this method, run `from paragradio.v2025_03 import SpecAnSim; help(SpecAnSim)` in your Python editor. If you are using a marimo notebook, you can view the docstring with rich formatting by running this in a cell: `mo.md(SpecAnSim.config.__doc__)`"""
        return SpecAnSim.__config(running, center_freq)

    
SpecAnSim.config.__doc__ +=_EXPLANATIONS._introductory + """Set the center frequency of simulated spectrum view."""

class WBFM_Rx():
    _ci: "Optional[WBFM_Rx]" = None
    "Current instance"

    def __init__(self) -> None:
        """Create a Paragradio Wideband FM Receiver."""
        from .wbfm_rx import wbfm_rx_fg
        self._pgr = ParallelGR(wbfm_rx_fg)
    
    @classmethod
    def __set_all(cls, center_freq, if_gain, bb_gain, hw_bb_filt, freq_offset, channel_width):
        cls._ci._pgr.put_cmd(_set_center_freq, center_freq)
        cls._ci._pgr.put_cmd(_set_if_gain, if_gain)
        cls._ci._pgr.put_cmd(_set_bb_gain, bb_gain)
        cls._ci._pgr.put_cmd(_set_hw_bb_filt, hw_bb_filt)
        cls._ci._pgr.put_cmd(_set_freq_offset, freq_offset)
        cls._ci._pgr.put_cmd(_set_channel_width,channel_width)

    @typechecked
    @staticmethod
    def config(
            *,
            running: bool,
            center_freq: float = 93e6,
            if_gain: int = 24,
            bb_gain: int = 32,
            hw_bb_filt: float = 2.75e6,
            freq_offset: float = 0.0,
            channel_width: float = 200e3,
            # Note: Can't add set_samp_rate because the rational resampler doesn't update at runtime
        ) -> dict:
        """To view the docs for this method, run `from paragradio.v2025_03 import WBFM_Rx; help(WBFM_Rx)` in your Python editor. If you are using a marimo notebook, you can view the docstring with rich formatting by running this in a cell: `mo.md(WBFM_Rx.config.__doc__)`"""
        if running == False:
            WBFM_Rx._ci._pgr.terminate()
            return {"terminated": "terminated"}
        decidemakenew(WBFM_Rx)
        WBFM_Rx.__set_all(center_freq, if_gain, bb_gain, hw_bb_filt, freq_offset, channel_width)
        return {
            "timestamp": datetime.datetime.now(),
        }
    

WBFM_Rx.config.__doc__ += _EXPLANATIONS._introductory + _EXPLANATIONS.center_freq + _EXPLANATIONS.if_gain + _EXPLANATIONS.bb_gain + _EXPLANATIONS.hw_bb_filt + _EXPLANATIONS.freq_offset + _EXPLANATIONS.channel_width


class Noise_Tx():
    _ci: "Optional[Noise_Tx]" = None
    "Current instance"

    def __init__(self) -> None:
        from .noise_tx import noise_tx_fg
        self._pgr = ParallelGR(noise_tx_fg)

    @classmethod
    def __set_all(cls, center_freq, amplitude, if_gain, noise_type, filter_cutoff_freq, filter_transition_width, samp_rate):
        cls._ci._pgr.put_cmd(_set_center_freq, center_freq)
        cls._ci._pgr.put_cmd(_set_amplitude, amplitude)
        cls._ci._pgr.put_cmd(_set_if_gain, if_gain)
        cls._ci._pgr.put_cmd(_set_noise_type, noise_type)
        cls._ci._pgr.put_cmd(_set_filter_cutoff_freq, filter_cutoff_freq)
        cls._ci._pgr.put_cmd(_set_filter_transition_width, filter_transition_width)
        cls._ci._pgr.put_cmd(_set_samp_rate, samp_rate)


    @typechecked
    @staticmethod
    def config(
            *,
            running: bool,
            center_freq: float = 2.4e9,
            amplitude: float = 0.01,
            if_gain: int = 0,
            noise_type: Literal["uniform", "gaussian"] = "uniform",
            filter_cutoff_freq: float = 50e3,
            filter_transition_width: float = 200e3,
            samp_rate: float = 2e6,
        ) -> dict:
        """To view the docs for this method, run `from paragradio.v2025_03 import Noise_Tx; help(Noise_Tx)` in your Python editor. If you are using a marimo notebook, you can view the docstring with rich formatting by running this in a cell: `mo.md(Noise_Tx.config.__doc__)`"""
        if running == False:
            Noise_Tx._ci._pgr.terminate()
            return {"terminated": "terminated"}
        decidemakenew(Noise_Tx)
        Noise_Tx.__set_all(center_freq, amplitude, if_gain, noise_type, filter_cutoff_freq, filter_transition_width, samp_rate)
        return {
            "timestamp": datetime.datetime.now(),
        }

Noise_Tx.config.__doc__ += _EXPLANATIONS._introductory + _EXPLANATIONS.center_freq + _EXPLANATIONS.amplitude + _EXPLANATIONS.if_gain + _EXPLANATIONS.noise_type + _EXPLANATIONS.filter_cutoff_freq + _EXPLANATIONS.filter_transition_width + _EXPLANATIONS.samp_rate


if TYPE_CHECKING:
    _Modulations: TypeAlias = Literal["BPSK", "QPSK", "DQPSK", "8PSK", "16QAM"]
modulationslist = ["BPSK", "QPSK", "DQPSK", "8PSK", "16QAM"]


class PSK_Tx_loop():
    _ci: "Optional[PSK_Tx_loop]" = None
    "Current instance"

    def __init__(self) -> None:
        from .psk_tx_loop import psk_tx_loop_fg
        self._pgr = ParallelGR(psk_tx_loop_fg)

    @classmethod
    def __set_all(cls, center_freq, if_gain, amplitude, data, samp_rate, modulation):
        cls._ci._pgr.put_cmd(_set_center_freq, center_freq)
        cls._ci._pgr.put_cmd(_set_if_gain, if_gain)
        cls._ci._pgr.put_cmd(_set_amplitude, amplitude)
        cls._ci._pgr.put_cmd(_set_data, data)
        cls._ci._pgr.put_cmd(_set_samp_rate, samp_rate)
        cls._ci._pgr.put_cmd(_set_modulation, modulation)

    @typechecked
    @staticmethod
    def config(
            *,
            running: bool,
            center_freq: float = 2.4e9,
            if_gain: int = 0,
            amplitude: float = 0.01,
            data: List[int],
            samp_rate: float = 2e6,
            modulation: Literal["BPSK", "QPSK", "DQPSK", "8PSK", "16QAM"] = "BPSK",
        ) -> dict:
        """To view the docs for this method, run `from paragradio.v2025_03 import PSK_Tx_loop; help(PSK_Tx_loop)` in your Python editor. If you are using a marimo notebook, you can view the docstring with rich formatting by running this in a cell: `mo.md(PSK_Tx_loop.config.__doc__)`"""
        if running == False:
            PSK_Tx_loop._ci._pgr.terminate()
            return {"terminated": "terminated"}
        decidemakenew(PSK_Tx_loop)
        PSK_Tx_loop.__set_all(center_freq, if_gain, amplitude, data, samp_rate, modulation)
        return {
            "timestamp": datetime.datetime.now(),
        }
    

PSK_Tx_loop.config.__doc__ += _EXPLANATIONS._introductory + _EXPLANATIONS.center_freq + _EXPLANATIONS.if_gain + _EXPLANATIONS.amplitude + _EXPLANATIONS.data + _EXPLANATIONS.samp_rate + _EXPLANATIONS.modulation
