"""
Apollo USB Downlink Receiver — top-level hierarchical block.

Combines the full demod chain into a single convenient block:
    complex baseband → PM demod → subcarrier extract → BPSK demod → frame sync → demux

Input: complex baseband samples at 5.12 MHz
Output: telemetry PDUs on message ports (frames, telemetry, agc_data)

This is the "drop one block into GRC" convenience for the common case.
For finer control, use the individual blocks directly.

Reference: IMPLEMENTATION_SPEC.md — full downlink path
"""

from gnuradio import gr

from apollo.bpsk_demod import bpsk_demod
from apollo.constants import (
    PCM_HIGH_BIT_RATE,
    PCM_SUBCARRIER_HZ,
    SAMPLE_RATE_BASEBAND,
)
from apollo.pcm_demux import pcm_demux
from apollo.pcm_frame_sync import pcm_frame_sync
from apollo.pm_demod import pm_demod
from apollo.subcarrier_extract import subcarrier_extract


class usb_downlink_receiver(gr.hier_block2):
    """Apollo USB downlink receiver — complex baseband to telemetry PDUs.

    Inputs:
        complex — baseband IQ samples at sample_rate (default 5.12 MHz)

    Message outputs (no streaming output):
        frames     — complete PCM frame PDUs (from frame sync)
        telemetry  — individual word PDUs with channel metadata
        agc_data   — AGC channel data (ch 34/35/57)
        raw_frame  — full frame passthrough

    The block chains: PM demod → subcarrier extract → BPSK demod → frame sync → demux.
    The BPSK demodulator recovers NRZ bits, which the frame sync correlates against the
    32-bit sync pattern. Locked frames are demultiplexed and emitted on message ports.
    """

    def __init__(
        self,
        sample_rate: float = SAMPLE_RATE_BASEBAND,
        bit_rate: int = PCM_HIGH_BIT_RATE,
        carrier_pll_bw: float = 0.02,
        subcarrier_bw: float = 150_000,
        bpsk_loop_bw: float = 0.045,
        max_bit_errors: int = 3,
        output_format: str = "raw",
    ):
        gr.hier_block2.__init__(
            self,
            "apollo_usb_downlink_receiver",
            gr.io_signature(1, 1, gr.sizeof_gr_complex),
            gr.io_signature(0, 0, 0),  # message-only output
        )

        # Register message output ports (pass raw strings — the method interns them)
        self.message_port_register_hier_out("frames")
        self.message_port_register_hier_out("telemetry")
        self.message_port_register_hier_out("agc_data")
        self.message_port_register_hier_out("raw_frame")

        # Stage 1: PM demodulator — carrier PLL + phase extraction
        self.pm = pm_demod(
            carrier_pll_bw=carrier_pll_bw,
            sample_rate=sample_rate,
        )

        # Stage 2: Subcarrier extractor — bandpass + downconvert 1.024 MHz
        self.sc_extract = subcarrier_extract(
            center_freq=PCM_SUBCARRIER_HZ,
            bandwidth=subcarrier_bw,
            sample_rate=sample_rate,
        )

        # Stage 3: BPSK demodulator — Costas loop + symbol sync + slicer
        self.bpsk = bpsk_demod(
            symbol_rate=bit_rate,
            sample_rate=sample_rate,
            loop_bw=bpsk_loop_bw,
        )

        # Stage 4: PCM frame synchronizer — 32-bit correlator
        self.frame_sync = pcm_frame_sync(
            bit_rate=bit_rate,
            max_bit_errors=max_bit_errors,
        )

        # Stage 5: PCM demultiplexer — word extraction + AGC channel ID
        self.demux = pcm_demux(
            output_format=output_format,
        )

        # Connect streaming chain: complex in → PM → subcarrier → BPSK → frame sync
        self.connect(self, self.pm, self.sc_extract, self.bpsk, self.frame_sync)

        # Connect message ports: frame_sync → demux → hier output ports
        self.msg_connect(self.frame_sync, "frames", self.demux, "frames")
        self.msg_connect(self.demux, "telemetry", self, "telemetry")
        self.msg_connect(self.demux, "agc_data", self, "agc_data")
        self.msg_connect(self.demux, "raw_frame", self, "raw_frame")

        # Also forward raw frames from frame_sync directly
        self.msg_connect(self.frame_sync, "frames", self, "frames")