openscreen/electron/native/wgc-capture/src/wasapi_loopback_capture.h

#pragma once

#include "mf_encoder.h"

#include <Windows.h>
#include <audioclient.h>
#include <mmdeviceapi.h>
#include <wrl/client.h>

#include <atomic>
#include <cstdint>
#include <functional>
#include <string>
#include <thread>
#include <vector>

enum class WasapiCaptureEndpoint {
    SystemLoopback,
    Microphone,
};

class WasapiLoopbackCapture {
public:
    using AudioCallback = std::function<void(const BYTE* data, DWORD byteCount, int64_t timestampHns, int64_t durationHns)>;

    WasapiLoopbackCapture() = default;
    ~WasapiLoopbackCapture();

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

    bool initializeSystemLoopback();
    bool initializeMicrophone(const std::wstring& deviceId, const std::wstring& deviceName);
    bool start(AudioCallback callback);
    void stop();

    const AudioInputFormat& inputFormat() const;
    const std::wstring& selectedDeviceName() const;

private:
    bool initialize(WasapiCaptureEndpoint endpoint, const std::wstring& deviceId, const std::wstring& deviceName);
    bool resolveMicrophoneByName(const std::wstring& deviceName);
    void captureLoop();
    bool resolveInputFormat(WAVEFORMATEX* mixFormat);

    Microsoft::WRL::ComPtr<IMMDeviceEnumerator> deviceEnumerator_;
    Microsoft::WRL::ComPtr<IMMDevice> device_;
    Microsoft::WRL::ComPtr<IAudioClient> audioClient_;
    Microsoft::WRL::ComPtr<IAudioCaptureClient> captureClient_;
    WAVEFORMATEX* mixFormat_ = nullptr;
    AudioInputFormat inputFormat_{};
    std::wstring selectedDeviceName_;
    AudioCallback callback_;
    std::thread thread_;
    std::atomic<bool> stopRequested_ = false;
    std::vector<BYTE> silenceBuffer_;
    uint64_t writtenFrames_ = 0;
    uint64_t lastDevicePositionEnd_ = 0;
    bool hasLastDevicePosition_ = false;
};