#!/usr/bin/env python3
"""
Apollo Full Downlink Demo -- PCM telemetry + crew voice on one carrier.

Reconstructs the complete Apollo USB downlink signal: PCM telemetry frames
on the 1.024 MHz BPSK subcarrier PLUS crew voice on the 1.25 MHz FM
subcarrier, both phase-modulated onto a single complex carrier.

Then receives the signal, splitting it into:
  - Decoded PCM telemetry frames (digital data)
  - Recovered crew voice audio (saved as WAV)

This is the full spacecraft-to-ground communications path:

    TX (spacecraft):
      pcm_frame_source -> nrz -> bpsk_mod (1.024 MHz) ─┐
      crew_audio -> fm_voice_mod (1.25 MHz, +/-29kHz) ──┤ scale 1.68/2.2
                                                         ├─> add -> pm_mod -> [RF]

    RX (ground station):
      [RF] -> pm_demod ──> subcarrier_extract -> bpsk_demod -> frame_sync ──> PCM frames
                       └─> voice_subcarrier_demod ──> crew audio (8 kHz)

Usage:
    uv run python examples/full_downlink_demo.py examples/audio/apollo11_crew.wav
    uv run python examples/full_downlink_demo.py input.wav --snr 25 --play
"""

import argparse
import time
from math import gcd

import numpy as np
import pmt
from gnuradio import blocks, gr
from scipy.io import wavfile
from scipy.signal import resample_poly

from apollo.bpsk_subcarrier_mod import bpsk_subcarrier_mod
from apollo.constants import (
    PCM_HIGH_BIT_RATE,
    PCM_HIGH_WORDS_PER_FRAME,
    PCM_SUBCARRIER_HZ,
    PCM_WORD_LENGTH,
    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_demux import DemuxEngine
from apollo.pcm_frame_source import pcm_frame_source
from apollo.pm_demod import pm_demod
from apollo.pm_mod import pm_mod
from apollo.usb_downlink_receiver import usb_downlink_receiver
from apollo.voice_subcarrier_demod import voice_subcarrier_demod


def load_and_upsample_audio(audio_path, sample_rate):
    """Load audio file and upsample to baseband rate."""
    input_rate, audio_data = wavfile.read(audio_path)
    if audio_data.ndim > 1:
        audio_data = audio_data[:, 0]

    # Normalize to [-1, 1]
    if audio_data.dtype == np.int16:
        audio_float = audio_data.astype(np.float32) / 32768.0
    else:
        audio_float = audio_data.astype(np.float32)

    duration = len(audio_float) / input_rate

    # Resample to 8 kHz first
    audio_rate = 8000
    if input_rate != audio_rate:
        g = gcd(audio_rate, input_rate)
        audio_float = resample_poly(audio_float, audio_rate // g, input_rate // g).astype(
            np.float32
        )

    # Upsample to baseband
    g = gcd(sample_rate, audio_rate)
    upsampled = resample_poly(audio_float, sample_rate // g, audio_rate // g).astype(np.float32)

    return upsampled, duration, audio_rate


def build_tx_signal(audio_samples, n_samples, sample_rate, snr_db):
    """Build the combined TX signal: PCM + voice -> PM modulation.

    Assembles the individual blocks manually (not using usb_signal_source)
    so we can inject external audio into the voice channel.
    """
    tb = gr.top_block()

    # --- PCM telemetry path ---
    frame_src = pcm_frame_source(bit_rate=PCM_HIGH_BIT_RATE)
    nrz = nrz_encoder(bit_rate=PCM_HIGH_BIT_RATE, sample_rate=sample_rate)
    bpsk = bpsk_subcarrier_mod(
        subcarrier_freq=PCM_SUBCARRIER_HZ,
        sample_rate=sample_rate,
    )
    tb.connect(frame_src, nrz, bpsk)

    # --- Voice subcarrier path (external audio) ---
    voice_src = blocks.vector_source_f(audio_samples[:n_samples].tolist())
    voice_mod = fm_voice_subcarrier_mod(
        sample_rate=sample_rate,
        subcarrier_freq=VOICE_SUBCARRIER_HZ,
        fm_deviation=VOICE_FM_DEVIATION_HZ,
        audio_input=True,
    )
    # Scale voice relative to PCM: 1.68/2.2 per IMPL_SPEC
    voice_gain = blocks.multiply_const_ff(1.68 / 2.2)
    tb.connect(voice_src, voice_mod, voice_gain)

    # --- Sum subcarriers ---
    adder = blocks.add_ff(1)
    tb.connect(bpsk, (adder, 0))
    tb.connect(voice_gain, (adder, 1))

    # --- PM modulation ---
    pm = pm_mod(pm_deviation=PM_PEAK_DEVIATION_RAD, sample_rate=sample_rate)
    head = blocks.head(gr.sizeof_gr_complex, n_samples)
    tb.connect(adder, pm, head)

    # --- Optional AWGN ---
    if snr_db is not None:
        import math

        noise_power = 1.0 / (10.0 ** (snr_db / 10.0))
        noise_amp = math.sqrt(noise_power / 2.0)
        noise = blocks.vector_source_c(
            (np.random.randn(n_samples) + 1j * np.random.randn(n_samples)).astype(np.complex64)
            * noise_amp
        )
        summer = blocks.add_cc(1)
        snk = blocks.vector_sink_c()
        tb.connect(head, (summer, 0))
        tb.connect(noise, (summer, 1))
        tb.connect(summer, snk)
    else:
        snk = blocks.vector_sink_c()
        tb.connect(head, snk)

    tb.run()
    return np.array(snk.data())


def receive_pcm(signal_data, sample_rate):
    """Run the PCM receive chain and return decoded frames."""
    tb = gr.top_block()
    src = blocks.vector_source_c(signal_data.tolist())
    rx = usb_downlink_receiver(
        sample_rate=sample_rate,
        bit_rate=PCM_HIGH_BIT_RATE,
        output_format="raw",
    )
    snk = blocks.message_debug()
    tb.connect(src, rx)
    tb.msg_connect(rx, "frames", snk, "store")
    tb.run()
    return snk


def receive_voice(signal_data, sample_rate, audio_rate=8000):
    """Run the voice receive chain and return recovered audio samples."""
    tb = gr.top_block()
    src = blocks.vector_source_c(signal_data.tolist())
    pm = pm_demod(sample_rate=sample_rate)
    voice = voice_subcarrier_demod(sample_rate=sample_rate, audio_rate=audio_rate)
    snk = blocks.vector_sink_f()
    tb.connect(src, pm, voice, snk)
    tb.run()
    return np.array(snk.data(), dtype=np.float32)


def main():
    parser = argparse.ArgumentParser(description="Full Apollo downlink: PCM telemetry + crew voice")
    parser.add_argument("audio", help="Input audio WAV file (crew voice)")
    parser.add_argument(
        "--output", "-o", default=None, help="Output WAV path (default: <input>_fullchain.wav)"
    )
    parser.add_argument("--snr", type=float, default=None, help="Add AWGN noise at this SNR in dB")
    parser.add_argument("--play", action="store_true", help="Play recovered voice with aplay")
    args = parser.parse_args()

    if args.output is None:
        stem = args.audio.rsplit(".", 1)[0]
        args.output = f"{stem}_fullchain.wav"

    sample_rate = int(SAMPLE_RATE_BASEBAND)
    audio_rate = 8000

    print("=" * 60)
    print("Apollo Full Downlink Demo")
    print("  PCM telemetry (1.024 MHz BPSK) + crew voice (1.25 MHz FM)")
    print("=" * 60)
    print()

    # Load and prepare audio
    print("Loading crew voice audio...")
    audio_upsampled, duration, _ = load_and_upsample_audio(args.audio, sample_rate)
    print(f"  Source:    {args.audio} ({duration:.2f}s)")
    print(f"  Upsampled: {len(audio_upsampled):,} samples at {sample_rate / 1e6:.2f} MHz")
    print()

    # Calculate frame timing
    bits_per_frame = PCM_HIGH_WORDS_PER_FRAME * PCM_WORD_LENGTH
    samples_per_frame = int(bits_per_frame * sample_rate / PCM_HIGH_BIT_RATE)
    n_frames = int(duration * 50) + 2  # 50 fps + margin
    n_samples = min(len(audio_upsampled), n_frames * samples_per_frame)

    print(f"  PCM frames: ~{n_frames} at 50 fps")
    print(f"  Signal:     {n_samples:,} samples ({n_samples / sample_rate:.2f}s)")
    snr_desc = f"{args.snr} dB" if args.snr is not None else "clean"
    print(f"  SNR:        {snr_desc}")
    print()

    # === TRANSMIT ===
    print("TX: Building combined PCM + voice signal...")
    t0 = time.time()
    signal = build_tx_signal(audio_upsampled, n_samples, sample_rate, args.snr)
    t_tx = time.time() - t0
    print(f"  Generated {len(signal):,} complex samples ({t_tx:.1f}s)")

    # Verify constant envelope (PM property)
    envelope = np.abs(signal[:10000])
    if args.snr is None:
        env_std = np.std(envelope)
        print(f"  PM envelope std: {env_std:.6f} (should be ~0 for clean)")
    print()

    # === RECEIVE: PCM telemetry ===
    print("RX: Decoding PCM telemetry frames...")
    t0 = time.time()
    frame_sink = receive_pcm(signal, sample_rate)
    t_pcm = time.time() - t0
    n_recovered = frame_sink.num_messages()
    print(f"  Recovered {n_recovered} PCM frames ({t_pcm:.1f}s)")

    if n_recovered > 0:
        demux = DemuxEngine(output_format="raw")
        print()
        for i in range(min(n_recovered, 5)):
            msg = frame_sink.get_message(i)
            payload = pmt.cdr(msg) if pmt.is_pair(msg) else msg
            frame_bytes = bytes(pmt.u8vector_elements(payload))
            result = demux.process_frame(frame_bytes)
            fid = result.get("sync", {}).get("frame_id", 0)
            n_words = len(result.get("words", []))
            parity = "odd" if fid % 2 == 1 else "even"
            print(f"    Frame {i + 1}: ID={fid:>2} ({parity}), {n_words} data words")
        if n_recovered > 5:
            print(f"    ... ({n_recovered - 5} more frames)")
    print()

    # === RECEIVE: crew voice ===
    print("RX: Demodulating crew voice (1.25 MHz FM)...")
    t0 = time.time()
    recovered_audio = receive_voice(signal, sample_rate, audio_rate)
    t_voice = time.time() - t0
    print(f"  Recovered {len(recovered_audio):,} audio samples ({t_voice:.1f}s)")
    print(f"  Duration:  {len(recovered_audio) / audio_rate:.2f}s at {audio_rate} Hz")

    # Normalize and save
    peak = np.max(np.abs(recovered_audio))
    if peak > 0:
        recovered_audio = recovered_audio / peak * 0.9
    recovered_int16 = (recovered_audio * 32767).astype(np.int16)
    wavfile.write(args.output, audio_rate, recovered_int16)
    print(f"  Saved:     {args.output}")
    print()

    # === SUMMARY ===
    print("=" * 60)
    print(f"  TX: {n_samples / sample_rate:.2f}s of combined PCM + voice")
    print(f"  RX: {n_recovered} PCM frames + {len(recovered_audio) / audio_rate:.2f}s crew voice")
    print(f"  SNR: {snr_desc}")
    print("=" * 60)

    if args.play:
        import subprocess

        print()
        print("Playing recovered crew voice...")
        subprocess.run(["aplay", args.output], check=False)
    else:
        print()
        print(f"Play voice: aplay {args.output}")


if __name__ == "__main__":
    main()