"""
Synthetic Apollo Unified S-Band downlink signal generator.

Generates complex baseband representing a PM-modulated carrier with:
- 1.024 MHz BPSK subcarrier (PCM telemetry NRZ data)
- Optional 1.25 MHz FM voice subcarrier (test tone)
- Configurable SNR

Used for testing the entire demodulation chain against known data.
All parameters from IMPLEMENTATION_SPEC.md sections 2.3, 4.2, 5.1.
"""

import numpy as np

from apollo.constants import (
    PCM_HIGH_BIT_RATE,
    PCM_HIGH_WORDS_PER_FRAME,
    PCM_SUBCARRIER_HZ,
    PCM_SYNC_WORD_LENGTH,
    PCM_WORD_LENGTH,
    PM_PEAK_DEVIATION_RAD,
    SAMPLE_RATE_BASEBAND,
    VOICE_FM_DEVIATION_HZ,
    VOICE_SUBCARRIER_HZ,
)
from apollo.protocol import generate_sync_word, sync_word_to_bits


def generate_pcm_frame(
    frame_id: int = 1,
    odd: bool = False,
    data: bytes | None = None,
    words_per_frame: int = PCM_HIGH_WORDS_PER_FRAME,
) -> list[int]:
    """Generate a complete PCM frame as a list of bits (MSB first, NRZ).

    Args:
        frame_id: Frame number (1-50).
        odd: Whether this is an odd-numbered frame (complement sync core).
        data: Optional payload bytes (words 5-128/200). Random if None.
        words_per_frame: 128 (high rate) or 200 (low rate).

    Returns:
        List of bit values (0 or 1), length = words_per_frame * 8.
    """
    # Generate 32-bit sync word (words 1-4)
    sync = generate_sync_word(frame_id=frame_id, odd=odd)
    bits = sync_word_to_bits(sync)

    # Data words (words 5 through end)
    data_words = words_per_frame - (PCM_SYNC_WORD_LENGTH // PCM_WORD_LENGTH)
    if data is not None:
        payload = list(data[:data_words])
        # Pad if needed
        while len(payload) < data_words:
            payload.append(0x00)
    else:
        payload = [np.random.randint(0, 256) for _ in range(data_words)]

    for byte_val in payload:
        for bit_pos in range(7, -1, -1):  # MSB first
            bits.append((byte_val >> bit_pos) & 1)

    return bits


def generate_nrz_waveform(
    bits: list[int],
    bit_rate: float,
    sample_rate: float,
) -> np.ndarray:
    """Convert a bit sequence to an NRZ baseband waveform.

    NRZ: bit 1 → +1.0, bit 0 → -1.0.

    Args:
        bits: List of bit values (0 or 1).
        bit_rate: Bit rate in Hz.
        sample_rate: Output sample rate in Hz.

    Returns:
        Float array of NRZ samples.
    """
    samples_per_bit = sample_rate / bit_rate
    n_samples = int(len(bits) * samples_per_bit)
    waveform = np.empty(n_samples, dtype=np.float32)

    for i, bit in enumerate(bits):
        start = int(i * samples_per_bit)
        end = int((i + 1) * samples_per_bit)
        waveform[start:end] = 1.0 if bit == 1 else -1.0

    return waveform


def generate_bpsk_subcarrier(
    nrz_data: np.ndarray,
    subcarrier_freq: float,
    sample_rate: float,
) -> np.ndarray:
    """Generate a BPSK-modulated subcarrier.

    The 1.024 MHz subcarrier is bi-phase modulated by NRZ data:
    output(t) = data(t) * cos(2*pi*f_sc*t)

    Args:
        nrz_data: NRZ waveform (+1/-1 values).
        subcarrier_freq: Subcarrier frequency in Hz.
        sample_rate: Sample rate in Hz.

    Returns:
        Float array of BPSK subcarrier samples.
    """
    t = np.arange(len(nrz_data), dtype=np.float64) / sample_rate
    carrier = np.cos(2.0 * np.pi * subcarrier_freq * t)
    return (nrz_data * carrier).astype(np.float32)


def generate_fm_voice_subcarrier(
    n_samples: int,
    sample_rate: float,
    tone_freq: float = 1000.0,
    subcarrier_freq: float = VOICE_SUBCARRIER_HZ,
    fm_deviation: float = VOICE_FM_DEVIATION_HZ,
) -> np.ndarray:
    """Generate an FM voice subcarrier with a test tone.

    Voice path: audio → FM VCO → upconvert to 1.25 MHz.

    Args:
        n_samples: Number of output samples.
        sample_rate: Sample rate in Hz.
        tone_freq: Audio test tone frequency in Hz.
        subcarrier_freq: Voice subcarrier center frequency.
        fm_deviation: FM deviation in Hz.

    Returns:
        Float array of FM voice subcarrier samples.
    """
    t = np.arange(n_samples, dtype=np.float64) / sample_rate
    # FM modulation: instantaneous phase = 2*pi*fc*t + (dev/f_tone)*sin(2*pi*f_tone*t)
    modulation_index = fm_deviation / tone_freq
    phase = 2.0 * np.pi * subcarrier_freq * t + modulation_index * np.sin(
        2.0 * np.pi * tone_freq * t
    )
    return np.cos(phase).astype(np.float32)


def generate_usb_baseband(
    frames: int = 1,
    bit_rate: float = PCM_HIGH_BIT_RATE,
    sample_rate: float = SAMPLE_RATE_BASEBAND,
    pm_deviation: float = PM_PEAK_DEVIATION_RAD,
    voice_enabled: bool = False,
    voice_tone_hz: float = 1000.0,
    snr_db: float | None = None,
    frame_data: list[bytes] | None = None,
) -> tuple[np.ndarray, list[list[int]]]:
    """Generate a complete Apollo USB downlink baseband signal.

    Produces complex baseband representing a PM-modulated carrier with
    BPSK PCM subcarrier and optional FM voice subcarrier.

    Args:
        frames: Number of PCM frames to generate.
        bit_rate: PCM bit rate (51200 or 1600).
        sample_rate: Output sample rate in Hz.
        pm_deviation: Peak PM deviation in radians.
        voice_enabled: Include 1.25 MHz FM voice subcarrier.
        voice_tone_hz: Voice test tone frequency.
        snr_db: If not None, add AWGN at this SNR (dB).
        frame_data: Optional list of payload bytes per frame.

    Returns:
        Tuple of (complex baseband signal, list of bit sequences per frame).
    """
    words_per_frame = 128 if bit_rate == PCM_HIGH_BIT_RATE else 200

    all_bits = []
    all_frame_bits = []

    for i in range(frames):
        frame_id = (i % 50) + 1
        odd = (frame_id % 2) == 1
        data = frame_data[i] if frame_data and i < len(frame_data) else None
        frame_bits = generate_pcm_frame(
            frame_id=frame_id, odd=odd, data=data, words_per_frame=words_per_frame
        )
        all_frame_bits.append(frame_bits)
        all_bits.extend(frame_bits)

    # NRZ waveform at the output sample rate
    nrz = generate_nrz_waveform(all_bits, bit_rate, sample_rate)

    # BPSK subcarrier
    pcm_subcarrier = generate_bpsk_subcarrier(nrz, PCM_SUBCARRIER_HZ, sample_rate)

    # Composite modulating signal (scaled for PM deviation)
    # The PCM subcarrier level sets the PM deviation
    modulating = pcm_subcarrier * pm_deviation

    if voice_enabled:
        voice = generate_fm_voice_subcarrier(
            len(nrz), sample_rate, tone_freq=voice_tone_hz
        )
        # Voice subcarrier at reduced level relative to PCM
        # Per IMPL_SPEC: PCM=2.2Vpp, Voice=1.68Vpp → ratio 1.68/2.2 ≈ 0.76
        voice_level = pm_deviation * (1.68 / 2.2)
        modulating = modulating + voice * voice_level

    # PM modulation: s(t) = exp(j * modulating(t))
    # At baseband, the carrier is at DC, so this is just phase modulation
    signal = np.exp(1j * modulating).astype(np.complex64)

    # Add noise if requested
    if snr_db is not None:
        signal_power = np.mean(np.abs(signal) ** 2)
        noise_power = signal_power / (10.0 ** (snr_db / 10.0))
        noise = np.sqrt(noise_power / 2) * (
            np.random.randn(len(signal)) + 1j * np.random.randn(len(signal))
        )
        signal = (signal + noise).astype(np.complex64)

    return signal, all_frame_bits