voice-cat/core/src/audio/audio_engine.h

/*
 * audio/audio_engine.h — capture/playback + DSP + jitter buffer + mixer.
 *
 * Design: docs/voice.md §8–11. Real-time path:
 *   capture(miniaudio) → APM(AEC/NS/AGC/VAD, send-side) → Opus encode → ...
 *   ... → Opus decode → per-user recv NS (listener-chosen) → gain/mute → mix → playback
 *
 * REAL-TIME RULE: audio-callback threads never allocate, lock, or block (architecture.md §3).
 * The JitterBuffer and per-stream maps are accessed only under a try_lock; a failed lock
 * causes PLC for that period (acceptable for M2; lock-free ring buffer is the M3 upgrade).
 */
#ifndef VOICECAT_AUDIO_AUDIO_ENGINE_H
#define VOICECAT_AUDIO_AUDIO_ENGINE_H

#include <algorithm>
#include <atomic>
#include <cstring>
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <unordered_map>
#include <vector>

#ifdef VOICECAT_HAS_AUDIO
// miniaudio single-header — MINIAUDIO_IMPLEMENTATION defined in audio_engine.cpp
#include <miniaudio.h>
#endif

#ifdef VOICECAT_HAS_OPUS
#include "codec/opus_codec.h"
#endif

#include "audio/apm_processor.h"

namespace voicecat::audio {

// ── JitterBuffer ─────────────────────────────────────────────────────────────
// Per-ssrc adaptive jitter buffer. Thread-safe via internal mutex.
class JitterBuffer {
   public:
    struct Frame {
        uint16_t             seq;
        uint32_t             timestamp;
        bool                 fec_present;
        std::vector<uint8_t> payload;
    };

    // Insert an incoming frame. Thread-safe.
    void push(Frame f);

    // Return the next frame whose timestamp <= playout_ts, or nullopt (caller should PLC).
    // Drops frames that are too old (more than kLateDropSamples late).
    std::optional<Frame> pop(uint32_t playout_ts);

    // Timestamp of the earliest buffered frame, or nullopt if empty/contended. Lets the playout
    // clock seed/re-sync itself to the arriving stream rather than free-running (see
    // AudioEngine::on_playback). Uses try_lock — never blocks the real-time callback.
    std::optional<uint32_t> peek_front_ts() const;

    uint32_t target_depth_ms()const { return target_depth_ms_.load(); }
    uint32_t packets_lost()   const { return lost_.load(); }
    void     reset();

   private:
    static constexpr uint32_t kLateDropSamples = 48000 / 2;  // 500 ms @48 kHz

    mutable std::mutex       mu_;
    std::map<uint32_t, Frame> buf_;   // keyed by timestamp (u32 wraps are handled below)

    std::atomic<uint32_t> target_depth_ms_{40};
    std::atomic<uint32_t> lost_{0};

    // Jitter estimation (EWMA).
    uint32_t last_push_ts_  = 0;      // local clock estimate on last push
    uint32_t jitter_est_    = 0;      // EWMA jitter in samples
    bool     first_push_    = true;
};

// ── AudioParams ──────────────────────────────────────────────────────────────
struct AudioParams {
    uint32_t    sample_rate        = 48000;
    uint32_t    capture_channels   = 1;  // no stereo capture device (mic) in this pass
    uint32_t    playback_channels  = 2;  // true stereo output (see audio_engine.cpp on_playback)
    uint32_t    frame_ms           = 20;
    std::string capture_device_id;   // "" = default; opaque id from AudioEngine::enumerate_devices
    std::string playback_device_id;  // "" = default; opaque id from AudioEngine::enumerate_devices
};

// One enumerated device, returned by AudioEngine::enumerate_devices(). `id` is an internal,
// opaque hex-encoded ma_device_id — callers must always round-trip an id that came from
// enumerate_devices(); never construct one by hand (names aren't guaranteed unique, so the id
// is the only stable handle miniaudio accepts back for device selection).
struct DeviceInfo {
    std::string id;
    std::string name;
    bool        is_default = false;
};

// ── AudioEngine ──────────────────────────────────────────────────────────────
// Owns miniaudio capture/playback, per-ssrc jitter buffers + Opus decoders, and the mixer.
class AudioEngine {
   public:
    // Callback type for encoded capture frames ready to be sent. `kind` identifies which
    // local stream this PCM belongs to (a vc_stream_kind value; 0 = MIC for the real capture
    // device, which is always the "primary" tap). `channels` is the channel count of the PCM
    // buffer (1 = mono, 2 = stereo interleaved) — the mic capture device is mono in v1, but
    // the WASAPI loopback path (SCREEN_AUDIO) captures in the channel's mode when stereo, so
    // the encoder sees real interleaved L/R PCM rather than a mono upmix. M3: multiple
    // concurrent local streams are possible (e.g. MIC + SCREEN_AUDIO), each fed via its own
    // injection tap (see inject_capture) since there is only one real hardware capture device.
    using CaptureCallback = std::function<void(int kind, const int16_t* pcm, int samples,
                                              int channels)>;

    AudioEngine();
    ~AudioEngine();

    AudioEngine(const AudioEngine&)            = delete;
    AudioEngine& operator=(const AudioEngine&) = delete;

    // Start capture+playback devices. capture_cb is called on the encode thread
    // for each capture frame (not on the audio callback thread).
    bool start(const AudioParams& p, CaptureCallback capture_cb = nullptr);
    void stop();

    bool running() const { return running_.load(std::memory_order_acquire); }

    // Enumerate input or output devices (for UI pickers / vc_list_devices). Static: works
    // before any AudioEngine instance is running (device pickers need to populate pre-connect).
    // Inits a throwaway ma_context if VOICECAT_HAS_AUDIO; returns {} otherwise. See DeviceInfo
    // above for the `id` encoding contract.
    static std::vector<DeviceInfo> enumerate_devices(bool capture);

    // Real desktop-audio loopback capture (Windows/WASAPI only, VOICECAT_HAS_LOOPBACK). Feeds
    // `kind`'s capture_cb_ directly, same pattern as the real mic capture device — NOT routed
    // through inject_capture()'s test-only ring. `channels` is the channel count to open the
    // loopback device with (1 = mono downmix of the system mix, 2 = stereo capture when the
    // channel is configured stereo); the accumulator and capture_cb_ invocation are shaped to
    // match. No-op (returns false) when unsupported.
    bool start_loopback_capture(int kind, int channels);
    void stop_loopback_capture();

    // Inject synthetic PCM directly into the capture pipeline (bypasses real device).
    // Thread-safe; can be called from any thread including tests. `kind` selects which local
    // stream's injection tap to feed (each gets its own ring buffer); the 2-arg overload
    // targets kind 0 (MIC) for source compatibility with existing callers.
    void inject_capture(int kind, const int16_t* pcm, size_t n);
    void inject_capture(const int16_t* pcm, size_t n) { inject_capture(0, pcm, n); }

    // Called by the net thread when a decoded voice frame arrives for a remote stream.
    void push_recv_frame(uint32_t ssrc, JitterBuffer::Frame f);

    // Per-stream receive-side controls (safe from any thread).
    void set_stream_gain(uint32_t ssrc, float gain);      // 0.0–2.0, default 1.0
    void set_stream_mute(uint32_t ssrc, bool mute);
    // Listener-chosen, local-only noise reduction on a specific remote stream (docs/voice.md
    // §10) — lazily instantiates an ApmProcessor on first enable, frees it on disable.
    void set_stream_noise_reduction(uint32_t ssrc, bool enable);
    void remove_stream(uint32_t ssrc);

    // Edge-triggered talk-state transitions since the last call (docs/voice.md §7: talk state
    // is derived from recent frame arrival, no protocol message). Call from a lightweight
    // poller, not the audio callback thread. Returns {ssrc, now_talking} for each stream whose
    // state flipped.
    std::vector<std::pair<uint32_t, bool>> poll_talk_transitions();

    // Get stats for a remote stream's jitter buffer.
    uint32_t stream_packets_lost(uint32_t ssrc) const;
    uint32_t stream_target_depth_ms(uint32_t ssrc) const;

#ifdef VOICECAT_HAS_OPUS
    // Configure the Opus decoder for an incoming ssrc (must be called before
    // push_recv_frame for that ssrc). Thread-safe.
    void init_recv_stream(uint32_t ssrc, const codec::OpusParams& p);
#endif

#ifdef VOICECAT_HAS_AUDIO
    // TEST-ONLY — exposes the playback mixer without a real ma_device, so tests can verify
    // stereo mixing end-to-end (no audio hardware needed). Same logic the real playback
    // callback uses; safe to call any time after start() (no ma_device touched).
    void mix_for_test(int16_t* out, uint32_t frames) { on_playback(out, frames); }

    // TEST-ONLY — drives the capture-side frame accumulator directly with an explicit
    // callback, bypassing capture_cb_. Call after engine.start(p) WITHOUT a capture
    // callback so the real mic (if any) never touches capture_accum_ (on_capture returns
    // early when capture_cb_ is null). The explicit `cb` is invoked only when a full
    // frame_samples_ chunk is ready — that is the invariant under test.
    void feed_capture_for_test(const int16_t* pcm, int frames, const CaptureCallback& cb) {
        if (frame_samples_ <= 0 || !cb || capture_accum_.buf.empty()) return;
        const int16_t* src = pcm;
        auto remaining = frames;
        while (remaining > 0) {
            int space = frame_samples_ - capture_accum_.count;
            int copy  = std::min(remaining, space);
            std::memcpy(capture_accum_.buf.data() + capture_accum_.count, src,
                        static_cast<size_t>(copy) * sizeof(int16_t));
            capture_accum_.count += copy;
            src       += copy;
            remaining -= copy;
            if (capture_accum_.count == frame_samples_) {
                cb(0, capture_accum_.buf.data(), frame_samples_, 1);
                capture_accum_.count = 0;
            }
        }
    }
#endif

#ifdef VOICECAT_HAS_LOOPBACK
    // TEST-ONLY — drives the loopback accumulator directly with an explicit callback, the
    // loopback analogue of feed_capture_for_test. Self-contained: sizes the accumulator and
    // sets loopback_channels_ itself, so it works on headless CI where start_loopback_capture
    // can't init a real WASAPI device. `channels` selects mono (1) or interleaved stereo (2).
    // PCM is interleaved L/R when channels==2. Invokes `cb` once per full frame_samples_
    // per-channel chunk, with the channel count passed through so the encoder branch in
    // on_capture_frame sees real stereo (channels==2) rather than a mono upmix.
    void feed_loopback_for_test(const int16_t* pcm, int frames_per_channel, int channels,
                                const CaptureCallback& cb) {
        if (frame_samples_ <= 0 || !cb) return;
        const int ch = std::max(1, channels);
        const int full = frame_samples_ * ch;
        // Size the accumulator for the requested channel count (off the RT thread; this is a
        // test-only path). start_loopback_capture() does the same sizing when it opens a real
        // device, but on headless CI that init fails — so do it here too.
        if (static_cast<int>(loopback_accum_.buf.size()) != full) {
            loopback_accum_.buf.assign(static_cast<size_t>(full), 0);
            loopback_accum_.count = 0;
        }
        loopback_channels_ = ch;
        const int16_t* src = pcm;
        auto remaining = frames_per_channel * ch;
        while (remaining > 0) {
            int space = full - loopback_accum_.count;
            int copy  = std::min(remaining, space);
            std::memcpy(loopback_accum_.buf.data() + loopback_accum_.count, src,
                        static_cast<size_t>(copy) * sizeof(int16_t));
            loopback_accum_.count += copy;
            src       += copy;
            remaining -= copy;
            if (loopback_accum_.count == full) {
                cb(loopback_kind_, loopback_accum_.buf.data(), frame_samples_, ch);
                loopback_accum_.count = 0;
            }
        }
    }
#endif

   private:
#ifdef VOICECAT_HAS_AUDIO
    static void capture_data_cb(ma_device*, void*, const void*, ma_uint32);
    static void playback_data_cb(ma_device*, void*, const void*, ma_uint32);
    void on_capture(const int16_t* pcm, ma_uint32 frames);
    void on_playback(int16_t* out, ma_uint32 frames);

    ma_device capture_device_{};
    ma_device playback_device_{};
    bool      capture_started_  = false;
    bool      playback_started_ = false;

#ifdef VOICECAT_HAS_LOOPBACK
    // Desktop-audio loopback capture (SCREEN_AUDIO) — own lifecycle, decoupled from
    // capture_device_/playback_device_ start/stop (a screen-share can start/stop independently
    // of the mic and of whether anything is currently playing back).
    static void loopback_data_cb(ma_device*, void*, const void*, ma_uint32);
    void        on_loopback(const int16_t* pcm, ma_uint32 frames);

    ma_device loopback_device_{};
    bool      loopback_started_ = false;
    int       loopback_kind_    = 0;
    int       loopback_channels_ = 1;  // channel count the loopback device was opened with
#endif
#endif

    AudioParams      params_{};
    CaptureCallback  capture_cb_;
    std::atomic<bool> running_{false};

    // Capture-side frame accumulators: miniaudio fires the capture (and loopback) callback at
    // whatever period the hardware/driver chooses — commonly 480 samples (10 ms) on WASAPI
    // shared mode, while the Opus encoder requires exactly frame_samples_ per call (960 for
    // 20 ms @ 48 kHz). Accumulate incoming PCM until a full frame is ready, then call
    // capture_cb_. This mirrors the RemoteStream::ring fix on the playback side. Both
    // accumulators are pre-allocated once in start(); never resized from the RT callback
    // thread (satisfies architecture.md §3 — no allocation on RT threads).
    struct CaptureAccum {
        std::vector<int16_t> buf;    // pre-sized to frame_samples_ in start()
        int                  count = 0;
    };
    CaptureAccum capture_accum_;   // mic / real capture device (on_capture)
    CaptureAccum loopback_accum_;  // screen-audio WASAPI loopback (on_loopback)

    // Inject ring(s): stores raw int16 PCM written by inject_capture(), one ring per local
    // stream kind so e.g. MIC and SCREEN_AUDIO can each be fed independently in tests.
    // The encode thread reads from these (no real capture device needed in tests).
    struct InjectTap {
        std::vector<int16_t> ring;            // circular, size = kInjectCapSamples
        std::atomic<size_t>  write{0};
        std::atomic<size_t>  read{0};
    };
    std::mutex                         inject_mu_;
    std::unordered_map<int, std::unique_ptr<InjectTap>> inject_taps_;
    static constexpr size_t kInjectCapSamples = 48000 * 2;  // 2 s @48 kHz mono

    // Per remote stream (protected by streams_mu_).
    struct RemoteStream {
        JitterBuffer     jitter;
#ifdef VOICECAT_HAS_OPUS
        codec::OpusDecoder decoder;
#endif
        float            gain  = 1.0f;
        bool             mute  = false;
        uint32_t         playout_ts = 0;
        // playout_ts free-runs (advances every callback via PLC), so it must be seeded from, and
        // periodically re-synced to, the actual stream timeline — otherwise it drifts past the
        // jitter buffer's drop window across VAD/PTT gaps and late joins and every frame is
        // dropped/never-due (silent playback). false until the first frame seeds it (on_playback).
        bool             playout_started = false;

        // M3: listener-chosen, local-only noise reduction (docs/voice.md §10). Lazily
        // created only when enabled — bounded by how many remote streams this listener
        // subscribes to, so no separate instance cap is needed.
        bool                           noise_reduction_enabled = false;
        std::unique_ptr<ApmProcessor>  recv_ns;

        // M3: talk-indicator edge detection (docs/voice.md §7) — updated by push_recv_frame
        // (already off the real-time audio thread), polled by poll_talk_transitions().
        std::atomic<int64_t> last_voice_ms{0};
        bool                  talking = false;

        // ── Decode/playback decoupling ring ─────────────────────────────────────
        // opus_decode() must be called with max_samples == the encoder's fixed frame size
        // (decoder.frame_samples(), e.g. 960 @ 20ms/48kHz) — that's a property of the bitstream,
        // not a choice. miniaudio's playback callback period is a *separate*, independently
        // chosen value (often smaller, e.g. ~480 @ low-latency WASAPI defaults) and must never
        // be passed to opus_decode as max_samples (doing so made decode fail basically every
        // callback — silent playback bug, fixed by this ring). on_playback() tops this ring up
        // by decoding whole Opus frames (decoder's channel count) and drains exactly the
        // hardware-requested sample count from it each callback, padding with silence (PLC) on
        // underrun. Sized once in init_ring() (called off the audio thread); never resized from
        // on_playback (real-time rule).
        std::vector<int16_t> ring;            // capacity = (frame_samples * 8) frames * ring_channels
        size_t                ring_channels = 1;
        size_t                ring_head     = 0;  // next frame (sample-per-channel) to read
        size_t                ring_count    = 0;  // buffered frames (samples-per-channel) ready
        std::vector<int16_t> decode_scratch;       // pre-sized: frame_samples * ring_channels

        void init_ring(int channels, int frame_samples) {
            ring_channels = static_cast<size_t>(std::max(1, channels));
            size_t cap_frames = static_cast<size_t>(std::max(1, frame_samples)) * 8;  // ~160ms @20ms frames
            ring.assign(cap_frames * ring_channels, 0);
            ring_head  = 0;
            ring_count = 0;
            decode_scratch.assign(static_cast<size_t>(std::max(1, frame_samples)) * ring_channels, 0);
        }

        // Appends `n_frames` samples-per-channel (ring_channels each) from `pcm`. Drops the
        // tail (rather than overwriting unread data) if the ring is unexpectedly full — should
        // not happen with the generous 8x sizing above.
        void push_ring(const int16_t* pcm, size_t n_frames) {
            if (ring.empty() || ring_channels == 0) return;
            size_t cap_frames = ring.size() / ring_channels;
            for (size_t i = 0; i < n_frames; ++i) {
                if (ring_count >= cap_frames) return;
                size_t widx = (ring_head + ring_count) % cap_frames;
                for (size_t c = 0; c < ring_channels; ++c)
                    ring[widx * ring_channels + c] = pcm[i * ring_channels + c];
                ++ring_count;
            }
        }

        // Pops one frame (sample-per-channel) into `out` (sized `out_channels`, zero-filled
        // first — covers both a fully-drained ring and ring_channels < out_channels).
        void pop_ring(int16_t* out, size_t out_channels) {
            for (size_t c = 0; c < out_channels; ++c) out[c] = 0;
            if (ring_count == 0 || ring.empty() || ring_channels == 0) return;
            size_t cap_frames = ring.size() / ring_channels;
            for (size_t c = 0; c < ring_channels && c < out_channels; ++c)
                out[c] = ring[ring_head * ring_channels + c];
            ring_head = (ring_head + 1) % cap_frames;
            --ring_count;
        }
    };
    mutable std::mutex                              streams_mu_;
    std::unordered_map<uint32_t, RemoteStream>      streams_;

    int frame_samples_ = 960;  // 20 ms @48 kHz

    static constexpr int64_t kTalkHangoverMs = 300;
};

}  // namespace voicecat::audio

#endif  // VOICECAT_AUDIO_AUDIO_ENGINE_H