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 <atomic>
#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);

    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    channels           = 1;
    uint32_t    frame_ms           = 20;
    std::string capture_device_id;   // "" = default
    std::string playback_device_id;  // "" = default
};

// ── 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). 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)>;

    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); }

    // 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

   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;
#endif

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

    // 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;

        // 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;
    };
    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