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

#ifdef VOICECAT_HAS_AUDIO
#define MINIAUDIO_IMPLEMENTATION
#include <miniaudio.h>
#endif

#include "audio/audio_engine.h"

#include <algorithm>
#include <chrono>
#include <cstring>

namespace voicecat::audio {

namespace {
int64_t now_ms() {
    return std::chrono::duration_cast<std::chrono::milliseconds>(
               std::chrono::steady_clock::now().time_since_epoch())
        .count();
}
}  // namespace

// ── JitterBuffer ─────────────────────────────────────────────────────────────

void JitterBuffer::push(Frame f) {
    std::lock_guard lk(mu_);

    uint32_t ts = f.timestamp;

    // Jitter estimation (EWMA of inter-arrival gap vs expected gap).
    if (!first_push_) {
        uint32_t arrived_gap = ts - last_push_ts_;
        uint32_t expected_gap = 960;  // 20 ms @48k; TODO: derive from params
        uint32_t diff = (arrived_gap > expected_gap) ? (arrived_gap - expected_gap)
                                                      : (expected_gap - arrived_gap);
        jitter_est_ = (jitter_est_ * 7 + diff) / 8;
        uint32_t depth = std::clamp(jitter_est_ * 2 + 960u, 960u, 48000u * 200u / 1000u);
        target_depth_ms_.store(depth * 1000u / 48000u, std::memory_order_relaxed);
    }
    last_push_ts_ = ts;
    first_push_   = false;

    buf_.emplace(ts, std::move(f));
}

std::optional<JitterBuffer::Frame> JitterBuffer::pop(uint32_t playout_ts) {
    std::unique_lock lk(mu_, std::try_to_lock);
    if (!lk) return std::nullopt;  // contended — caller does PLC

    if (buf_.empty()) return std::nullopt;

    auto it = buf_.begin();
    uint32_t ts = it->first;

    // Drop frames that are too old (> 500 ms late).
    if (static_cast<int32_t>(playout_ts - ts) > static_cast<int32_t>(kLateDropSamples)) {
        lost_.fetch_add(1, std::memory_order_relaxed);
        buf_.erase(it);
        return std::nullopt;
    }

    // Return frame only when it's due.
    if (static_cast<int32_t>(ts - playout_ts) > 0) return std::nullopt;

    Frame f = std::move(it->second);
    buf_.erase(it);
    return f;
}

void JitterBuffer::reset() {
    std::lock_guard lk(mu_);
    buf_.clear();
    lost_.store(0);
    first_push_ = true;
}

// ── AudioEngine ──────────────────────────────────────────────────────────────

AudioEngine::AudioEngine() = default;

AudioEngine::~AudioEngine() { stop(); }

bool AudioEngine::start(const AudioParams& p, CaptureCallback capture_cb) {
    if (running_.load()) return false;
    params_        = p;
    capture_cb_    = std::move(capture_cb);
    frame_samples_ = static_cast<int>(p.sample_rate / 1000 * p.frame_ms);
    running_.store(true, std::memory_order_release);

#ifdef VOICECAT_HAS_AUDIO
    // ── Capture device ──────────────────────────────────────────────────────
    ma_device_config cap_cfg = ma_device_config_init(ma_device_type_capture);
    cap_cfg.capture.format   = ma_format_s16;
    cap_cfg.capture.channels = p.channels;
    cap_cfg.sampleRate       = p.sample_rate;
    cap_cfg.dataCallback     = capture_data_cb;
    cap_cfg.pUserData        = this;
    cap_cfg.capture.pDeviceID = nullptr;  // always default for now

    if (ma_device_init(nullptr, &cap_cfg, &capture_device_) == MA_SUCCESS) {
        if (ma_device_start(&capture_device_) == MA_SUCCESS) {
            capture_started_ = true;
        } else {
            ma_device_uninit(&capture_device_);
        }
    }

    // ── Playback device ─────────────────────────────────────────────────────
    ma_device_config pb_cfg = ma_device_config_init(ma_device_type_playback);
    pb_cfg.playback.format   = ma_format_s16;
    pb_cfg.playback.channels = p.channels;
    pb_cfg.sampleRate        = p.sample_rate;
    pb_cfg.dataCallback      = playback_data_cb;
    pb_cfg.pUserData         = this;
    pb_cfg.playback.pDeviceID = nullptr;

    if (ma_device_init(nullptr, &pb_cfg, &playback_device_) == MA_SUCCESS) {
        if (ma_device_start(&playback_device_) == MA_SUCCESS) {
            playback_started_ = true;
        } else {
            ma_device_uninit(&playback_device_);
        }
    }
#endif  // VOICECAT_HAS_AUDIO

    return true;
}

void AudioEngine::stop() {
    if (!running_.exchange(false)) return;

#ifdef VOICECAT_HAS_AUDIO
    if (capture_started_) {
        ma_device_stop(&capture_device_);
        ma_device_uninit(&capture_device_);
        capture_started_ = false;
    }
    if (playback_started_) {
        ma_device_stop(&playback_device_);
        ma_device_uninit(&playback_device_);
        playback_started_ = false;
    }
#endif
}

void AudioEngine::inject_capture(int kind, const int16_t* pcm, size_t n) {
    InjectTap* tap;
    {
        std::lock_guard lk(inject_mu_);
        auto& slot = inject_taps_[kind];
        if (!slot) {
            slot = std::make_unique<InjectTap>();
            slot->ring.resize(kInjectCapSamples, 0);
        }
        tap = slot.get();
    }

    size_t w = tap->write.load(std::memory_order_relaxed);
    for (size_t i = 0; i < n; ++i)
        tap->ring[(w + i) % kInjectCapSamples] = pcm[i];
    tap->write.store(w + n, std::memory_order_release);

    // Fire capture_cb_ for each complete frame now available.
    while (true) {
        size_t r     = tap->read.load(std::memory_order_relaxed);
        size_t avail = tap->write.load(std::memory_order_acquire) - r;
        if (avail < static_cast<size_t>(frame_samples_)) break;

        std::vector<int16_t> frame(frame_samples_);
        for (int i = 0; i < frame_samples_; ++i)
            frame[i] = tap->ring[(r + i) % kInjectCapSamples];
        tap->read.store(r + frame_samples_, std::memory_order_release);

        if (capture_cb_) capture_cb_(kind, frame.data(), frame_samples_);
    }
}

void AudioEngine::push_recv_frame(uint32_t ssrc, JitterBuffer::Frame f) {
    std::lock_guard lk(streams_mu_);
    auto& s = streams_[ssrc];
    s.last_voice_ms.store(now_ms(), std::memory_order_relaxed);
    s.jitter.push(std::move(f));
}

void AudioEngine::set_stream_gain(uint32_t ssrc, float gain) {
    std::lock_guard lk(streams_mu_);
    streams_[ssrc].gain = gain;
}

void AudioEngine::set_stream_mute(uint32_t ssrc, bool mute) {
    std::lock_guard lk(streams_mu_);
    streams_[ssrc].mute = mute;
}

void AudioEngine::set_stream_noise_reduction(uint32_t ssrc, bool enable) {
    std::lock_guard lk(streams_mu_);
    auto& s = streams_[ssrc];
    s.noise_reduction_enabled = enable;
    if (enable) {
        if (!s.recv_ns) s.recv_ns = ApmProcessor::create();
    } else {
        s.recv_ns.reset();
    }
}

void AudioEngine::remove_stream(uint32_t ssrc) {
    std::lock_guard lk(streams_mu_);
    streams_.erase(ssrc);
}

std::vector<std::pair<uint32_t, bool>> AudioEngine::poll_talk_transitions() {
    std::vector<std::pair<uint32_t, bool>> edges;
    std::lock_guard lk(streams_mu_);
    int64_t now = now_ms();
    for (auto& [ssrc, stream] : streams_) {
        bool now_talking = (now - stream.last_voice_ms.load(std::memory_order_relaxed)) <
                            kTalkHangoverMs;
        if (now_talking != stream.talking) {
            stream.talking = now_talking;
            edges.emplace_back(ssrc, now_talking);
        }
    }
    return edges;
}

uint32_t AudioEngine::stream_packets_lost(uint32_t ssrc) const {
    std::lock_guard lk(streams_mu_);
    auto it = streams_.find(ssrc);
    return (it != streams_.end()) ? it->second.jitter.packets_lost() : 0;
}

uint32_t AudioEngine::stream_target_depth_ms(uint32_t ssrc) const {
    std::lock_guard lk(streams_mu_);
    auto it = streams_.find(ssrc);
    return (it != streams_.end()) ? it->second.jitter.target_depth_ms() : 40;
}

#ifdef VOICECAT_HAS_OPUS
void AudioEngine::init_recv_stream(uint32_t ssrc, const codec::OpusParams& p) {
    std::lock_guard lk(streams_mu_);
    streams_[ssrc].decoder.init(p);
}
#endif

#ifdef VOICECAT_HAS_AUDIO

void AudioEngine::capture_data_cb(ma_device* dev, void* /*out*/,
                                   const void* in, ma_uint32 frame_count) {
    auto* self = static_cast<AudioEngine*>(dev->pUserData);
    self->on_capture(static_cast<const int16_t*>(in), frame_count);
}

void AudioEngine::on_capture(const int16_t* pcm, ma_uint32 frames) {
    // The real hardware capture device is always the "primary" tap (kind 0 / MIC). A second
    // concurrent local stream (e.g. SCREEN_AUDIO) is fed via inject_capture() in M3 — there is
    // only one real capture device.
    if (capture_cb_) capture_cb_(0, pcm, static_cast<int>(frames));
}

void AudioEngine::playback_data_cb(ma_device* dev, void* out,
                                    const void* /*in*/, ma_uint32 frame_count) {
    auto* self = static_cast<AudioEngine*>(dev->pUserData);
    self->on_playback(static_cast<int16_t*>(out), frame_count);
}

void AudioEngine::on_playback(int16_t* out, ma_uint32 frames) {
    std::memset(out, 0, frames * params_.channels * sizeof(int16_t));

    std::unique_lock lk(streams_mu_, std::try_to_lock);
    if (!lk) return;  // contended: emit silence this period

#ifdef VOICECAT_HAS_OPUS
    std::vector<int32_t> mix(frames * params_.channels, 0);

    for (auto& [ssrc, stream] : streams_) {
        if (stream.mute || !stream.decoder.valid()) continue;

        // M3: a stream's Opus channel count (mono/stereo, per-channel AudioConfig) may differ
        // from the engine-wide playback channel count (always mono in M3 — see PROGRESS.md).
        // Decode into a buffer sized for the *decoder's* channel count (opus_decode's
        // frame_size parameter is samples-per-channel, not total samples — pass `frames`,
        // not `frames * channels`), then convert at this boundary.
        int dec_channels = std::max(1, stream.decoder.channels());
        auto maybe_frame = stream.jitter.pop(stream.playout_ts);
        std::vector<int16_t> pcm(frames * static_cast<ma_uint32>(dec_channels));
        int n;

        if (maybe_frame) {
            n = stream.decoder.decode(
                maybe_frame->payload.data(),
                static_cast<int>(maybe_frame->payload.size()),
                pcm.data(), static_cast<int>(frames));
        } else {
            n = stream.decoder.decode(nullptr, 0, pcm.data(), static_cast<int>(frames));
        }

        if (n > 0) {
            // `n` is samples-per-channel (matches the frame_samples convention used by
            // OpusEncoder::encode elsewhere in the codebase).
            if (stream.recv_ns)
                stream.recv_ns->process_capture(pcm.data(), n,
                                                static_cast<int>(params_.sample_rate));

            float g = stream.gain;
            for (int i = 0; i < n; ++i) {
                // Downmix decoder output to the engine's mono accumulator if needed
                // (average L/R); upmix is unnecessary since the mix buffer is per-channel.
                int32_t sample = (dec_channels == 2)
                                      ? (static_cast<int32_t>(pcm[i * 2]) +
                                         static_cast<int32_t>(pcm[i * 2 + 1])) / 2
                                      : static_cast<int32_t>(pcm[i]);
                for (uint32_t c = 0; c < params_.channels; ++c)
                    mix[i * params_.channels + c] += static_cast<int32_t>(sample * g);
            }
        }
        stream.playout_ts += frames;
    }

    for (ma_uint32 i = 0; i < frames * params_.channels; ++i)
        out[i] = static_cast<int16_t>(std::clamp(mix[i], -32768, 32767));
#else
    (void)out;
    (void)frames;
#endif
}

#endif  // VOICECAT_HAS_AUDIO

}  // namespace voicecat::audio