#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
#
# satnogs_iq_to_ssb_audio.py
#
# Inspired by the idea shared as demod_usb2.py by hobisatelit / Bali:
# https://github.com/hobisatelit/
#
# Original demod_usb2.py notice:
# Copyright 2026 hobisatelit
# License: GPL-3.0-or-later
#
# This cleaned-up version keeps the same purpose:
# convert SatNOGS-style complex IQ raw recordings into SSB-demodulated audio
# for downstream tools such as SoundModem or post-processing scripts.
#
# Requirements:
#   pip install numpy scipy
#   ffmpeg is optional and only needed for OGG output.

from __future__ import annotations

import argparse
import json
import math
import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np
from scipy import signal
from scipy.io import wavfile


def _load_iq(path: Path, input_format: str, swap_iq: bool, conjugate: bool) -> np.ndarray:
    """Load interleaved complex IQ samples and normalize them to complex64."""

    if input_format == "cs16le":
        raw = np.fromfile(path, dtype="<i2")
        if raw.size < 2:
            raise ValueError("Input file does not contain enough int16 IQ samples.")
        if raw.size % 2:
            raw = raw[:-1]
        iq_i = raw[0::2].astype(np.float32)
        iq_q = raw[1::2].astype(np.float32)
        iq = (iq_i + 1j * iq_q) / 32768.0

    elif input_format == "cf32le":
        raw = np.fromfile(path, dtype="<f4")
        if raw.size < 2:
            raise ValueError("Input file does not contain enough float32 IQ samples.")
        if raw.size % 2:
            raw = raw[:-1]
        iq = raw[0::2].astype(np.float32) + 1j * raw[1::2].astype(np.float32)

    else:
        raise ValueError(f"Unsupported input format: {input_format}")

    if swap_iq:
        iq = iq.imag + 1j * iq.real

    if conjugate:
        iq = np.conj(iq)

    return iq.astype(np.complex64, copy=False)


def _frequency_shift(iq: np.ndarray, sample_rate: float, freq_offset_hz: float) -> np.ndarray:
    """Shift complex IQ by freq_offset_hz.

    Positive values move a signal at +freq_offset_hz down to baseband.
    This matches the convention used by many simple IQ correction scripts.
    """

    if freq_offset_hz == 0:
        return iq

    n = np.arange(iq.size, dtype=np.float64)
    osc = np.exp(-2j * np.pi * freq_offset_hz * n / sample_rate)
    return (iq * osc).astype(np.complex64, copy=False)


def _sos_filter(data: np.ndarray, sos: np.ndarray) -> np.ndarray:
    """Apply zero-phase filtering when possible, otherwise fall back to causal filtering."""

    if data.size > 3 * max(1, sos.shape[0]) * 2:
        try:
            return signal.sosfiltfilt(sos, data).astype(data.dtype, copy=False)
        except ValueError:
            pass

    return signal.sosfilt(sos, data).astype(data.dtype, copy=False)


def _complex_lowpass(iq: np.ndarray, sample_rate: float, cutoff_hz: float, order: int) -> np.ndarray:
    nyquist = sample_rate / 2.0
    cutoff_hz = min(max(float(cutoff_hz), 1.0), nyquist * 0.98)
    sos = signal.butter(order, cutoff_hz / nyquist, btype="lowpass", output="sos")
    real = _sos_filter(iq.real.astype(np.float32), sos)
    imag = _sos_filter(iq.imag.astype(np.float32), sos)
    return (real + 1j * imag).astype(np.complex64, copy=False)


def _audio_bandpass(audio: np.ndarray, sample_rate: float, low_hz: float, high_hz: float, order: int) -> np.ndarray:
    nyquist = sample_rate / 2.0
    low_hz = max(float(low_hz), 0.0)
    high_hz = min(float(high_hz), nyquist * 0.98)

    if high_hz <= 0 or high_hz <= low_hz:
        return audio.astype(np.float32, copy=False)

    if low_hz <= 0:
        sos = signal.butter(order, high_hz / nyquist, btype="lowpass", output="sos")
    else:
        sos = signal.butter(order, [low_hz / nyquist, high_hz / nyquist], btype="bandpass", output="sos")

    return _sos_filter(audio.astype(np.float32), sos).astype(np.float32, copy=False)


def _resample_if_needed(audio: np.ndarray, sample_rate: int, audio_rate: int) -> np.ndarray:
    if sample_rate == audio_rate:
        return audio.astype(np.float32, copy=False)

    gcd = math.gcd(sample_rate, audio_rate)
    up = audio_rate // gcd
    down = sample_rate // gcd
    return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)


def _normalize(audio: np.ndarray, mode: str, peak: float, rms: float) -> np.ndarray:
    audio = audio.astype(np.float32, copy=False)

    if mode == "none":
        return np.clip(audio, -1.0, 1.0)

    if mode == "peak":
        max_abs = float(np.max(np.abs(audio))) if audio.size else 0.0
        if max_abs > 0:
            audio = audio / max_abs * float(peak)
        return np.clip(audio, -1.0, 1.0)

    if mode == "rms":
        current_rms = float(np.sqrt(np.mean(audio * audio))) if audio.size else 0.0
        if current_rms > 0:
            audio = audio / current_rms * float(rms)
        max_abs = float(np.max(np.abs(audio))) if audio.size else 0.0
        if max_abs > 1.0:
            audio = audio / max_abs * float(peak)
        return np.clip(audio, -1.0, 1.0)

    raise ValueError(f"Unsupported normalization mode: {mode}")


def _write_wav(path: Path, sample_rate: int, audio: np.ndarray) -> None:
    pcm = np.int16(np.clip(audio, -1.0, 1.0) * 32767.0)
    wavfile.write(path, sample_rate, pcm)


def _write_ogg_with_ffmpeg(wav_path: Path, ogg_path: Path, quality: int) -> None:
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg not found. WAV was written, but OGG conversion cannot be performed.")

    cmd = [
        ffmpeg,
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(wav_path),
        "-codec:a",
        "libvorbis",
        "-qscale:a",
        str(quality),
        str(ogg_path),
    ]
    subprocess.run(cmd, check=True)


def _write_sidecar(path: Path, metadata: dict) -> None:
    path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def build_argparser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Convert complex IQ raw samples to SSB-demodulated audio. "
            "Designed for SatNOGS post-processing experiments."
        )
    )

    parser.add_argument("input_file", type=Path, help="Input raw IQ file.")
    parser.add_argument("output_wav", type=Path, help="Output WAV file.")

    parser.add_argument(
        "--input-format",
        choices=["cs16le", "cf32le"],
        default="cs16le",
        help="Input IQ format. SatNOGS IQ dumps are commonly cs16le. Default: cs16le.",
    )
    parser.add_argument("--sample-rate", type=int, default=48000, help="IQ sample rate in Hz. Default: 48000.")
    parser.add_argument(
        "--freq-offset",
        type=float,
        default=-1230.0,
        help=(
            "Frequency correction in Hz. Positive values shift a signal at +offset down to baseband. "
            "Default: -1230.0, matching the original Bali script."
        ),
    )
    parser.add_argument("--sideband", choices=["usb", "lsb"], default="usb", help="SSB sideband. Default: usb.")

    parser.add_argument("--swap-iq", action="store_true", help="Swap I and Q before processing.")
    parser.add_argument("--conjugate", action="store_true", help="Conjugate IQ before processing.")

    parser.add_argument(
        "--ssb-bandwidth",
        type=float,
        default=5000.0,
        help="Complex low-pass bandwidth after frequency shift, in Hz. Default: 5000.",
    )
    parser.add_argument("--audio-low", type=float, default=300.0, help="Audio high-pass cutoff in Hz. Default: 300.")
    parser.add_argument("--audio-high", type=float, default=4200.0, help="Audio low-pass cutoff in Hz. Default: 4200.")
    parser.add_argument("--filter-order", type=int, default=5, help="Butterworth filter order. Default: 5.")

    parser.add_argument(
        "--audio-rate",
        type=int,
        default=48000,
        help="Output audio sample rate in Hz. Default: 48000.",
    )
    parser.add_argument(
        "--normalize",
        choices=["peak", "rms", "none"],
        default="peak",
        help="Audio normalization mode. Default: peak.",
    )
    parser.add_argument("--peak", type=float, default=0.80, help="Peak normalization target. Default: 0.80.")
    parser.add_argument("--rms", type=float, default=0.18, help="RMS normalization target. Default: 0.18.")

    parser.add_argument(
        "--write-json",
        action="store_true",
        help="Write a JSON sidecar next to the WAV with processing parameters.",
    )
    parser.add_argument(
        "--ogg",
        type=Path,
        default=None,
        help="Optional OGG/Vorbis output path. Requires ffmpeg.",
    )
    parser.add_argument(
        "--ogg-quality",
        type=int,
        default=5,
        help="OGG/Vorbis quality used when --ogg is provided. Default: 5.",
    )

    return parser


def main() -> int:
    parser = build_argparser()
    args = parser.parse_args()

    try:
        if args.sample_rate <= 0:
            raise ValueError("--sample-rate must be positive.")
        if args.audio_rate <= 0:
            raise ValueError("--audio-rate must be positive.")
        if args.filter_order < 1:
            raise ValueError("--filter-order must be >= 1.")

        iq = _load_iq(args.input_file, args.input_format, args.swap_iq, args.conjugate)

        # Remove static DC before tuning.
        iq = iq - np.mean(iq)

        iq = _frequency_shift(iq, args.sample_rate, args.freq_offset)

        # For LSB, flip the analytic spectrum before extracting real audio.
        if args.sideband == "lsb":
            iq = np.conj(iq)

        iq = _complex_lowpass(iq, args.sample_rate, args.ssb_bandwidth, args.filter_order)

        audio = iq.real.astype(np.float32, copy=False)
        audio = audio - float(np.mean(audio))

        audio = _audio_bandpass(audio, args.sample_rate, args.audio_low, args.audio_high, args.filter_order)
        audio = _resample_if_needed(audio, args.sample_rate, args.audio_rate)
        audio = _normalize(audio, args.normalize, args.peak, args.rms)

        args.output_wav.parent.mkdir(parents=True, exist_ok=True)
        _write_wav(args.output_wav, args.audio_rate, audio)

        metadata = {
            "input_file": str(args.input_file),
            "output_wav": str(args.output_wav),
            "input_format": args.input_format,
            "sample_rate_hz": args.sample_rate,
            "audio_rate_hz": args.audio_rate,
            "freq_offset_hz": args.freq_offset,
            "sideband": args.sideband,
            "swap_iq": args.swap_iq,
            "conjugate": args.conjugate,
            "ssb_bandwidth_hz": args.ssb_bandwidth,
            "audio_low_hz": args.audio_low,
            "audio_high_hz": args.audio_high,
            "filter_order": args.filter_order,
            "normalize": args.normalize,
            "peak": args.peak,
            "rms": args.rms,
            "samples_in": int(iq.size),
            "samples_out": int(audio.size),
        }

        if args.write_json:
            _write_sidecar(args.output_wav.with_suffix(args.output_wav.suffix + ".json"), metadata)

        if args.ogg is not None:
            args.ogg.parent.mkdir(parents=True, exist_ok=True)
            _write_ogg_with_ffmpeg(args.output_wav, args.ogg, args.ogg_quality)
            metadata["output_ogg"] = str(args.ogg)
            if args.write_json:
                _write_sidecar(args.ogg.with_suffix(args.ogg.suffix + ".json"), metadata)

        return 0

    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())