"""
Apollo USB Signal Source -- complete transmit chain in one block.

The transmit-side counterpart to usb_downlink_receiver. Wires together the
full modulation chain:

    pcm_frame_source -> nrz_encoder -> bpsk_subcarrier_mod -+-> add_ff -> pm_mod -> [complex out]
                                                             |
                             fm_voice_subcarrier_mod --------+
                             (optional, scaled by 1.68/2.2)

This mirrors CuriousMarc's physical bench topology: the individual composable
blocks map 1:1 to Keysight instruments (EXG signal generator for PM, two
33522B AWGs for subcarrier modulation).

For finer control, use the individual blocks directly.

Reference: IMPLEMENTATION_SPEC.md -- full downlink transmit path
"""

import math

from gnuradio import analog, blocks, gr

from apollo.bpsk_subcarrier_mod import bpsk_subcarrier_mod
from apollo.constants import (
    PCM_HIGH_BIT_RATE,
    PCM_SUBCARRIER_HZ,
    PM_PEAK_DEVIATION_RAD,
    SAMPLE_RATE_BASEBAND,
    VOICE_FM_DEVIATION_HZ,
    VOICE_SUBCARRIER_HZ,
)
from apollo.fm_voice_subcarrier_mod import fm_voice_subcarrier_mod
from apollo.nrz_encoder import nrz_encoder
from apollo.pcm_frame_source import pcm_frame_source
from apollo.pm_mod import pm_mod


class usb_signal_source(gr.hier_block2):
    """Apollo USB downlink signal source -- complex baseband output.

    Outputs:
        complex -- PM-modulated baseband at sample_rate (default 5.12 MHz)

    Message inputs:
        frame_data -- forwarded to pcm_frame_source for dynamic payload injection

    The block generates PCM telemetry frames, NRZ-encodes them, BPSK-modulates
    onto a 1.024 MHz subcarrier, optionally adds a 1.25 MHz FM voice subcarrier,
    and applies PM modulation to produce complex baseband.

    Optional AWGN noise can be added by setting snr_db to a finite value.
    """

    def __init__(
        self,
        sample_rate: float = SAMPLE_RATE_BASEBAND,
        bit_rate: int = PCM_HIGH_BIT_RATE,
        pm_deviation: float = PM_PEAK_DEVIATION_RAD,
        voice_enabled: bool = False,
        voice_tone_hz: float = 1000.0,
        snr_db: float | None = None,
    ):
        gr.hier_block2.__init__(
            self,
            "apollo_usb_signal_source",
            gr.io_signature(0, 0, 0),              # source -- no input
            gr.io_signature(1, 1, gr.sizeof_gr_complex),
        )

        self._sample_rate = sample_rate
        self._voice_enabled = voice_enabled

        # Forward the frame_data message port from pcm_frame_source
        self.message_port_register_hier_in("frame_data")

        # --- PCM telemetry path ---

        # Stage 1: Generate PCM frame bits (0/1 byte stream)
        self.frame_src = pcm_frame_source(bit_rate=bit_rate)

        # Forward message port: hier input -> pcm_frame_source
        self.msg_connect(self, "frame_data", self.frame_src, "frame_data")

        # Stage 2: NRZ encode (byte 0/1 -> float +1/-1 at sample_rate)
        self.nrz = nrz_encoder(bit_rate=bit_rate, sample_rate=sample_rate)

        # Stage 3: BPSK modulate onto 1.024 MHz subcarrier
        self.bpsk = bpsk_subcarrier_mod(
            subcarrier_freq=PCM_SUBCARRIER_HZ,
            sample_rate=sample_rate,
        )

        # Connect PCM chain: frame_src -> nrz -> bpsk
        self.connect(self.frame_src, self.nrz, self.bpsk)

        # --- Subcarrier summing ---

        if voice_enabled:
            # Voice subcarrier level relative to PCM:
            # Per IMPL_SPEC: PCM = 2.2 Vpp, Voice = 1.68 Vpp
            # The BPSK subcarrier has unity amplitude, so voice is scaled
            # by 1.68/2.2 to maintain the correct power ratio.
            voice_scale = 1.68 / 2.2

            self.voice = fm_voice_subcarrier_mod(
                sample_rate=sample_rate,
                subcarrier_freq=VOICE_SUBCARRIER_HZ,
                fm_deviation=VOICE_FM_DEVIATION_HZ,
                tone_freq=voice_tone_hz,
            )
            self.voice_gain = blocks.multiply_const_ff(voice_scale)
            self.adder = blocks.add_ff(1)

            # PCM subcarrier -> adder port 0
            self.connect(self.bpsk, (self.adder, 0))
            # Voice subcarrier (scaled) -> adder port 1
            self.connect(self.voice, self.voice_gain, (self.adder, 1))

            composite = self.adder
        else:
            composite = self.bpsk

        # --- PM modulation ---

        self.pm = pm_mod(pm_deviation=pm_deviation, sample_rate=sample_rate)
        self.connect(composite, self.pm)

        # --- Optional AWGN ---

        if snr_db is not None:
            # Signal power is 1.0 (PM constant envelope)
            noise_power = 1.0 / (10.0 ** (snr_db / 10.0))
            noise_amplitude = math.sqrt(noise_power / 2.0)

            self.noise = analog.noise_source_c(
                analog.GR_GAUSSIAN, noise_amplitude, 0,
            )
            self.sum_noise = blocks.add_cc(1)

            self.connect(self.pm, (self.sum_noise, 0))
            self.connect(self.noise, (self.sum_noise, 1))
            self.connect(self.sum_noise, self)
        else:
            self.connect(self.pm, self)