video-pc/AvRecorder/capturer/audio/wasapi_loopback_capturer.cpp

#include "wasapi_loopback_capturer.h"
#include <audioclient.h>
#include <comdef.h>
#include <mmdeviceapi.h>
#include <windows.h>
#include <QDebug>

#undef min

class WASAPILoopbackCapturerPrivate
{
public:
    WASAPILoopbackCapturerPrivate() { CoInitialize(nullptr); }
    ~WASAPILoopbackCapturerPrivate()
    {
        if (pwfx)
            CoTaskMemFree(pwfx);
        if (pCaptureClient)
            pCaptureClient->Release();
        if (pAudioClient)
            pAudioClient->Release();
        if (pDevice)
            pDevice->Release();
        if (pEnumerator)
            pEnumerator->Release();
        CoUninitialize();
    }

    bool init()
    {
        HRESULT hr;
        hr = CoCreateInstance(__uuidof(MMDeviceEnumerator),
                              nullptr,
                              CLSCTX_ALL,
                              __uuidof(IMMDeviceEnumerator),
                              (void**) &pEnumerator);
        if (FAILED(hr)) {
            qDebug() << "Failed to create MMDeviceEnumerator";
            return false;
        }

        hr = pEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &pDevice);
        if (FAILED(hr)) {
            qDebug() << "Failed to get default audio endpoint";
            return false;
        }

        hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**) &pAudioClient);
        if (FAILED(hr)) {
            qDebug() << "Failed to activate audio client";
            return false;
        }
        return true;
    }

    // 初始化音频客户端
    bool initializeAudioClient()
    {
        AUDCLNT_SHAREMODE shareMode = AUDCLNT_SHAREMODE_SHARED;
        DWORD streamFlags = AUDCLNT_STREAMFLAGS_LOOPBACK | AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM
                            | AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY;
        REFERENCE_TIME hnsBufferDuration = 0;

        HRESULT hr
            = pAudioClient->Initialize(shareMode, streamFlags, hnsBufferDuration, 0, pwfx, nullptr);
        if (FAILED(hr)) {
            qDebug() << "Failed to initialize audio client, HRESULT:" << QString::number(hr, 16);
            return false;
        }
        qDebug() << "Audio client initialized successfully";
        return true;
    }

    bool setupCaptureClient()
    {
        if (!pAudioClient) {
            qDebug() << "Audio client is null, cannot get capture client";
            return false;
        }
        
        HRESULT hr = pAudioClient->GetService(__uuidof(IAudioCaptureClient), (void**) &pCaptureClient);
        if (FAILED(hr)) {
            qDebug() << "Failed to get capture client, HRESULT:" << QString::number(hr, 16) 
                     << "AudioClient:" << pAudioClient << "CaptureClient:" << pCaptureClient;
            return false;
        }
        
        qDebug() << "Capture client obtained successfully:" << pCaptureClient;
        return true;
    }

    bool startCapture()
    {
        if (!pAudioClient) {
            qDebug() << "Audio client is null, cannot start capture";
            return false;
        }
        
        HRESULT hr = pAudioClient->Start();
        if (FAILED(hr)) {
            qDebug() << "Failed to start audio client, HRESULT:" << QString::number(hr, 16);
            return false;
        }
        
        qDebug() << "Audio client started successfully";
        return true;
    }

    void stopCapture()
    {
        if (pAudioClient) {
            pAudioClient->Stop();
        }
    }

    // 获取音频格式并设置到AudioFormat结构中
    bool setupAudioFormat(AudioFormat& audioFormat)
    {
        if (!pwfx) {
            return false;
        }

        // 设置音频格式
        audioFormat.sampleRate = pwfx->nSamplesPerSec;
        audioFormat.channels = pwfx->nChannels;
        audioFormat.bitsPerSample = pwfx->wBitsPerSample;
        audioFormat.blockAlign = pwfx->nBlockAlign;
        audioFormat.avgBytesPerSec = pwfx->nAvgBytesPerSec;

        // 如果是浮点格式，保持32bit，让混音器处理格式转换
        if (pwfx->wFormatTag == WAVE_FORMAT_IEEE_FLOAT) {
            qDebug() << "Keeping 32t float format for mixer processing";
        }

        return true;
    }

    // 处理音频数据
    int processAudioData(std::vector<char>& buffer, std::mutex& mutex)
    {
        UINT32 packetLength = 0;
        HRESULT hr = pCaptureClient->GetNextPacketSize(&packetLength);
        if (FAILED(hr)) {
            qDebug() << "Failed to get next packet size";
            return -1;
        }
        
        if (packetLength == 0) {
            return 0;
        }
        
        BYTE* pData = nullptr;
        UINT32 numFrames = 0;
        DWORD flags = 0;
        hr = pCaptureClient->GetBuffer(&pData, &numFrames, &flags, nullptr, nullptr);
        if (FAILED(hr)) {
            qDebug() << "Failed to get buffer";
            return -1;
        }
        
        int bytes = numFrames * pwfx->nBlockAlign;

        if (pwfx->wFormatTag == WAVE_FORMAT_IEEE_FLOAT) {
            // 32bit浮点格式，直接传递原始数据，让混音器处理格式转换
            std::lock_guard<std::mutex> lock(mutex);
            // 限制缓冲区大小，避免内存无限增长
            const size_t maxBufferSize = 1024 * 1024; // 1MB
            int bytesToAdd = numFrames * pwfx->nBlockAlign;
            if (buffer.size() + bytesToAdd > maxBufferSize) {
                buffer.erase(buffer.begin(), buffer.begin() + (buffer.size() - bytesToAdd));
            }
            buffer.insert(buffer.end(),
                          (char*) pData,
                          (char*) pData + bytesToAdd);
        } else {
            // PCM 直接拷贝
            std::lock_guard<std::mutex> lock(mutex);
            // 限制缓冲区大小，避免内存无限增长
            const size_t maxBufferSize = 1024 * 1024; // 1MB
            int bytesToAdd = numFrames * pwfx->nBlockAlign;
            if (buffer.size() + bytesToAdd > maxBufferSize) {
                buffer.erase(buffer.begin(), buffer.begin() + (buffer.size() - bytesToAdd));
            }
            buffer.insert(buffer.end(),
                          (char*) pData,
                          (char*) pData + bytesToAdd);
        }

        pCaptureClient->ReleaseBuffer(numFrames);
        return bytes;
    }

    IMMDeviceEnumerator* pEnumerator = nullptr;
    IMMDevice* pDevice = nullptr;
    IAudioClient* pAudioClient = nullptr;
    IAudioCaptureClient* pCaptureClient = nullptr;
    WAVEFORMATEX* pwfx = nullptr;

    WAVEFORMATEXTENSIBLE _formatex;
};

WASAPILoopbackCapturer::WASAPILoopbackCapturer(QObject* parent)
    : QObject(parent)
    , d(new WASAPILoopbackCapturerPrivate)
{

}

WASAPILoopbackCapturer::~WASAPILoopbackCapturer()
{
    Stop();
    delete d;
}

bool WASAPILoopbackCapturer::Init(Type deviceType)
{
    // 只支持扬声器
    if (deviceType != Type::Speaker)
        return false;

    if (!d->init()) {
        qDebug() << "Failed to initialize WASAPI components";
        return false;
    }

    // 尝试使用默认格式
    AudioFormat defaultFormat = IAudioCapturer::GetDefaultFormat();
    WAVEFORMATEX* formatToUse = (WAVEFORMATEX*)CoTaskMemAlloc(sizeof(WAVEFORMATEX));
    if (!formatToUse) {
        qDebug() << "Failed to allocate memory for WAVEFORMATEX";
        return false;
    }
    formatToUse->wFormatTag = WAVE_FORMAT_PCM;
    formatToUse->nChannels = defaultFormat.channels;
    formatToUse->nSamplesPerSec = defaultFormat.sampleRate;
    formatToUse->wBitsPerSample = defaultFormat.bitsPerSample;
    formatToUse->nBlockAlign = (formatToUse->nChannels * formatToUse->wBitsPerSample) / 8;
    formatToUse->nAvgBytesPerSec = formatToUse->nSamplesPerSec * formatToUse->nBlockAlign;
    formatToUse->cbSize = 0;

    d->pwfx = formatToUse;

    // 尝试使用默认格式初始化，如果失败则回退
    if (!d->initializeAudioClient()) {
        qWarning("Default audio format not supported, falling back to GetMixFormat.");

        CoTaskMemFree(d->pwfx);
        d->pwfx = nullptr;

        // 使用GetMixFormat获取系统首选格式
        HRESULT hr = d->pAudioClient->GetMixFormat(&d->pwfx);
        if (FAILED(hr)) {
            qDebug() << "Failed to GetMixFormat on fallback, HRESULT:" << QString::number(hr, 16);
            return false;
        }

        // 再次尝试初始化
        if (!d->initializeAudioClient()) {
            qDebug() << "Failed to initialize audio client on fallback.";
            CoTaskMemFree(d->pwfx);
            d->pwfx = nullptr;
            return false;
        }
    }

    // 设置音频格式
    if (!d->setupAudioFormat(m_audioFormat)) {
        qDebug() << "Failed to setup audio format";
        return false;
    }

    qDebug() << "WASAPI Loopback Capturer initialized successfully";
    return true;
}

bool WASAPILoopbackCapturer::Start()
{
    if (m_running)
        return false;
        
    if (!d->setupCaptureClient()) {
        return false;
    }
    
    if (!d->startCapture()) {
        return false;
    }
    
    m_running = true;
    m_captureThread = std::thread(&WASAPILoopbackCapturer::captureThreadFunc, this);
    return true;
}

void WASAPILoopbackCapturer::Stop()
{
    if (!m_running)
        return;
    m_running = false;
    if (m_captureThread.joinable())
        m_captureThread.join();
    d->stopCapture();
}

const AudioFormat& WASAPILoopbackCapturer::GetFormat() const
{
    return m_audioFormat;
}

int WASAPILoopbackCapturer::readAudioData(char* buf, int maxLen)
{
    std::unique_lock<std::mutex> lock(m_mutex);
    
    // 按帧对齐读取，确保不会破坏音频帧
    int blockAlign = m_audioFormat.blockAlign;
    if (blockAlign <= 0) {
        return 0;
    }
    
    // 计算可以读取的完整帧数
    int availableFrames = m_buffer.size() / blockAlign;
    int requestedFrames = maxLen / blockAlign;
    int framesToRead = std::min(availableFrames, requestedFrames);
    
    if (framesToRead > 0) {
        int bytesToRead = framesToRead * blockAlign;
        memcpy(buf, m_buffer.data(), bytesToRead);
        m_buffer.erase(m_buffer.begin(), m_buffer.begin() + bytesToRead);
        return bytesToRead;
    }
    
    return 0;
}

void WASAPILoopbackCapturer::captureThreadFunc()
{
    qDebug() << "WASAPI Loopback capture started successfully";

    while (m_running) {
        int result = d->processAudioData(m_buffer, m_mutex);
        if (result < 0) {
            qDebug() << "Error processing audio data";
            break;
        }
        if (result == 0) {
            Sleep(10);
        }
    }

    qDebug() << "WASAPI Loopback capture stopped";
}