"""
Apollo PCM Frame Source -- generates a continuous NRZ bit stream of PCM frames.

The transmit-side counterpart to pcm_frame_sync.  Produces a steady stream of
128-word (high rate, 51.2 kbps) or 200-word (low rate, 1.6 kbps) PCM frames,
each beginning with the standard 32-bit sync word:

    [5-bit A][15-bit core][6-bit B][6-bit frame ID]

Frame IDs cycle 1 through 50 (one subframe), with the 15-bit core complemented
on odd-numbered frames.  An optional message input allows dynamic payload
injection; otherwise frames carry zero-fill data.

The core logic lives in FrameSourceEngine (pure Python, testable without GNU
Radio).  The GR sync_block wrapper bridges frame-granularity generation with
GR's sample-granularity scheduler via an internal bit buffer.

Reference: IMPLEMENTATION_SPEC.md sections 5.1, 5.2
"""

from collections import deque

import numpy as np

from apollo.constants import (
    PCM_HIGH_BIT_RATE,
    PCM_HIGH_WORDS_PER_FRAME,
    PCM_LOW_WORDS_PER_FRAME,
)
from apollo.usb_signal_gen import generate_pcm_frame


class FrameSourceEngine:
    """PCM frame generation engine (pure Python, no GR dependency).

    Maintains a rolling frame counter (1-50) and generates complete frames
    on demand via next_frame().  Odd-numbered frames get a complemented
    sync core automatically.

    Args:
        bit_rate: PCM bit rate in bps (51200 or 1600).
    """

    def __init__(self, bit_rate: int = PCM_HIGH_BIT_RATE):
        self.bit_rate = bit_rate
        if bit_rate == PCM_HIGH_BIT_RATE:
            self.words_per_frame = PCM_HIGH_WORDS_PER_FRAME
        else:
            self.words_per_frame = PCM_LOW_WORDS_PER_FRAME

        self.frame_counter = 1

    def next_frame(self, data: bytes | None = None) -> list[int]:
        """Generate the next PCM frame as a list of bits (0/1 values, MSB first).

        Args:
            data: Optional payload bytes for data words.  If None, the frame
                  carries zero-fill (deterministic, unlike the random fill in
                  generate_pcm_frame when data=None for signal-gen use).

        Returns:
            List of bit values, length = words_per_frame * 8.
        """
        frame_id = self.frame_counter
        odd = (frame_id % 2) == 1

        # Default to zero-fill rather than random for a transmit source --
        # downstream blocks and tests need deterministic output.
        if data is None:
            data = bytes(self.words_per_frame)

        bits = generate_pcm_frame(
            frame_id=frame_id,
            odd=odd,
            data=data,
            words_per_frame=self.words_per_frame,
        )

        # Advance counter: 1 -> 2 -> ... -> 50 -> 1
        self.frame_counter = (self.frame_counter % 50) + 1

        return bits


# ---------------------------------------------------------------------------
# GNU Radio block wrapper (optional -- only if gnuradio is available)
# ---------------------------------------------------------------------------

try:
    import pmt
    from gnuradio import gr

    class pcm_frame_source(gr.sync_block):
        """GNU Radio source block: continuous PCM frame bit stream.

        Outputs a stream of bytes (values 0 or 1) representing NRZ-encoded
        PCM telemetry frames.  Frame IDs cycle 1-50 automatically.

        An optional ``frame_data`` message input accepts PMT u8vector payloads
        that will be used as the data words for the next generated frame.

        Parameters:
            bit_rate: 51200 (128 words/frame) or 1600 (200 words/frame).
        """

        def __init__(self, bit_rate: int = PCM_HIGH_BIT_RATE):
            gr.sync_block.__init__(
                self,
                name="apollo_pcm_frame_source",
                in_sig=None,
                out_sig=[np.byte],
            )

            self._engine = FrameSourceEngine(bit_rate=bit_rate)
            self._bit_buffer: deque[int] = deque()
            self._pending_data: bytes | None = None

            # Message input for dynamic payload injection
            self.message_port_register_in(pmt.intern("frame_data"))
            self.set_msg_handler(
                pmt.intern("frame_data"), self._handle_frame_data
            )

        def _handle_frame_data(self, msg):
            """Store incoming PMT payload bytes for the next frame."""
            if pmt.is_u8vector(msg):
                self._pending_data = bytes(pmt.u8vector_elements(msg))
            elif pmt.is_pair(msg):
                # Accept PDU (car=meta, cdr=payload)
                payload = pmt.cdr(msg)
                if pmt.is_u8vector(payload):
                    self._pending_data = bytes(pmt.u8vector_elements(payload))

        def work(self, input_items, output_items):
            out = output_items[0]
            n_out = len(out)
            produced = 0

            while produced < n_out:
                if not self._bit_buffer:
                    frame_bits = self._engine.next_frame(data=self._pending_data)
                    self._pending_data = None
                    self._bit_buffer.extend(frame_bits)

                chunk = min(n_out - produced, len(self._bit_buffer))
                for i in range(chunk):
                    out[produced + i] = self._bit_buffer.popleft()
                produced += chunk

            return produced

except ImportError:
    pass