LoopbackCapture.cpp

#include <windows.h>
#include <mmdeviceapi.h>
#include <audioclient.h>
#include <avrt.h>
#include <fstream>
#include <iostream>
#include <cstdint>
#include <string>
#include <direct.h>  // For _getcwd
#include <vector>    // Added for std::vector

#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "avrt.lib")

// WAV file header structure
struct WAVHeader {
    char riff[4] = {'R', 'I', 'F', 'F'};
    uint32_t chunkSize;
    char wave[4] = {'W', 'A', 'V', 'E'};
    char fmt[4] = {'f', 'm', 't', ' '};
    uint32_t subchunk1Size = 16;
    uint16_t audioFormat = 1; // PCM
    uint16_t numChannels;
    uint32_t sampleRate;
    uint32_t byteRate;
    uint16_t blockAlign;
    uint16_t bitsPerSample;
    char data[4] = {'d', 'a', 't', 'a'};
    uint32_t dataSize;
};

// Function to write WAV header
void WriteWAVHeader(std::ofstream& outFile, WAVEFORMATEX* pwfx, uint32_t dataSize) {
    WAVHeader header;
    header.numChannels = pwfx->nChannels;
    header.sampleRate = pwfx->nSamplesPerSec;
    header.bitsPerSample = pwfx->wBitsPerSample;
    header.byteRate = pwfx->nAvgBytesPerSec;
    header.blockAlign = pwfx->nBlockAlign;
    header.dataSize = dataSize;
    header.chunkSize = 36 + dataSize;

    outFile.write(reinterpret_cast<const char*>(&header), sizeof(WAVHeader));
}

// Function to get full path of output file
std::string GetOutputPath(const std::string& fileName) {
    char currentDir[MAX_PATH];
    if (_getcwd(currentDir, MAX_PATH) != nullptr) {
        return std::string(currentDir) + "\\" + fileName;
    }
    return fileName; // Fallback to just the filename
}

int main() {
    // Print current directory
    char currentDir[MAX_PATH];
    if (_getcwd(currentDir, MAX_PATH) != nullptr) {
        std::cout << "Current directory: " << currentDir << std::endl;
    }

    HRESULT hr;
    hr = CoInitialize(nullptr);
    if (FAILED(hr)) {
        std::cerr << "CoInitialize failed: " << std::hex << hr << std::endl;
        return -1;
    }

    IMMDeviceEnumerator* pEnumerator = nullptr;
    hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
                          __uuidof(IMMDeviceEnumerator), (void**)&pEnumerator);
    if (FAILED(hr)) {
        std::cerr << "CoCreateInstance failed: " << std::hex << hr << std::endl;
        CoUninitialize();
        return -1;
    }

    IMMDevice* pDevice = nullptr;
    hr = pEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &pDevice);
    pEnumerator->Release();
    if (FAILED(hr)) {
        std::cerr << "GetDefaultAudioEndpoint failed: " << std::hex << hr << std::endl;
        CoUninitialize();
        return -1;
    }

    IAudioClient* pAudioClient = nullptr;
    hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**)&pAudioClient);
    pDevice->Release();
    if (FAILED(hr)) {
        std::cerr << "Device Activate failed: " << std::hex << hr << std::endl;
        CoUninitialize();
        return -1;
    }

    WAVEFORMATEX* pwfx = nullptr;
    hr = pAudioClient->GetMixFormat(&pwfx);
    if (FAILED(hr)) {
        std::cerr << "GetMixFormat failed: " << std::hex << hr << std::endl;
        pAudioClient->Release();
        CoUninitialize();
        return -1;
    }

    // Print audio format information for debugging
    std::cout << "Audio Format: " << std::endl;
    std::cout << "  Format Tag: " << pwfx->wFormatTag << std::endl;
    std::cout << "  Channels: " << pwfx->nChannels << std::endl;
    std::cout << "  Sample Rate: " << pwfx->nSamplesPerSec << std::endl;
    std::cout << "  Bits Per Sample: " << pwfx->wBitsPerSample << std::endl;
    std::cout << "  Block Align: " << pwfx->nBlockAlign << std::endl;
    std::cout << "  Bytes Per Second: " << pwfx->nAvgBytesPerSec << std::endl;

    // Check if we need to convert the format to PCM
    // WAVE_FORMAT_IEEE_FLOAT = 3
    bool isFloat = (pwfx->wFormatTag == 3);
    
    // If format is extensible, check if it's float
    if (pwfx->wFormatTag == WAVE_FORMAT_EXTENSIBLE && pwfx->cbSize >= 22) {
        // Extensible format
        WAVEFORMATEXTENSIBLE* pwfxext = reinterpret_cast<WAVEFORMATEXTENSIBLE*>(pwfx);
        // Check first 4 bytes of GUID to see if it's float format
        // KSDATAFORMAT_SUBTYPE_IEEE_FLOAT first DWORD is 0x00000003
        DWORD formatType = pwfxext->SubFormat.Data1;
        isFloat = (formatType == 3);
    }

    // If format is not PCM, create a compatible PCM format
    bool customFormat = false;
    
    if (pwfx->wFormatTag != WAVE_FORMAT_PCM) {
        // Try to find a compatible PCM format
        std::cout << "Not a standard PCM format, trying to find compatible format..." << std::endl;
        
        // Create a PCM format with the same number of channels and sample rate
        WAVEFORMATEX pcmFormat;
        pcmFormat.wFormatTag = WAVE_FORMAT_PCM;
        pcmFormat.nChannels = pwfx->nChannels;
        pcmFormat.nSamplesPerSec = pwfx->nSamplesPerSec;
        pcmFormat.wBitsPerSample = 16; // Standard 16-bit PCM
        pcmFormat.nBlockAlign = pcmFormat.nChannels * (pcmFormat.wBitsPerSample / 8);
        pcmFormat.nAvgBytesPerSec = pcmFormat.nSamplesPerSec * pcmFormat.nBlockAlign;
        pcmFormat.cbSize = 0;
        
        // Use the new format
        CoTaskMemFree(pwfx);
        pwfx = (WAVEFORMATEX*)CoTaskMemAlloc(sizeof(WAVEFORMATEX));
        memcpy(pwfx, &pcmFormat, sizeof(WAVEFORMATEX));
        customFormat = true;
        
        std::cout << "Using custom PCM format:" << std::endl;
        std::cout << "  Channels: " << pwfx->nChannels << std::endl;
        std::cout << "  Sample Rate: " << pwfx->nSamplesPerSec << std::endl;
        std::cout << "  Bits Per Sample: " << pwfx->wBitsPerSample << std::endl;
    }

    hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_SHARED,
                                  AUDCLNT_STREAMFLAGS_LOOPBACK,
                                  0, 0, pwfx, nullptr);
    if (FAILED(hr)) {
        std::cerr << "AudioClient Initialize failed: " << std::hex << hr << std::endl;
        if (customFormat) {
            std::cout << "Failed with custom format, trying original format..." << std::endl;
            CoTaskMemFree(pwfx);
            hr = pAudioClient->GetMixFormat(&pwfx);
            if (SUCCEEDED(hr)) {
                hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_SHARED,
                                             AUDCLNT_STREAMFLAGS_LOOPBACK,
                                             0, 0, pwfx, nullptr);
            }
        }
        
        // If still failing, exit
        if (FAILED(hr)) {
            CoTaskMemFree(pwfx);
            pAudioClient->Release();
            CoUninitialize();
            return -1;
        }
    }

    IAudioCaptureClient* pCaptureClient = nullptr;
    hr = pAudioClient->GetService(__uuidof(IAudioCaptureClient), (void**)&pCaptureClient);
    if (FAILED(hr)) {
        std::cerr << "GetService failed: " << std::hex << hr << std::endl;
        pAudioClient->Release();
        CoTaskMemFree(pwfx);
        CoUninitialize();
        return -1;
    }

    hr = pAudioClient->Start();
    if (FAILED(hr)) {
        std::cerr << "AudioClient Start failed: " << std::hex << hr << std::endl;
        pCaptureClient->Release();
        pAudioClient->Release();
        CoTaskMemFree(pwfx);
        CoUninitialize();
        return -1;
    }

    // Use full path for output file
    std::string outputPath = GetOutputPath("output.wav");
    std::cout << "Will save to: " << outputPath << std::endl;
    
    std::ofstream outFile(outputPath, std::ios::binary);
    if (!outFile) {
        std::cerr << "Failed to open " << outputPath << " for writing." << std::endl;
        std::cerr << "Error code: " << GetLastError() << std::endl;
        pAudioClient->Stop();
        pCaptureClient->Release();
        pAudioClient->Release();
        CoTaskMemFree(pwfx);
        CoUninitialize();
        return -1;
    }

    // Write placeholder for WAV header
    WAVHeader placeholderHeader = {};
    outFile.write(reinterpret_cast<const char*>(&placeholderHeader), sizeof(WAVHeader));

    if (!outFile) {
        std::cerr << "Failed to write WAV header placeholder" << std::endl;
        outFile.close();
        pAudioClient->Stop();
        pCaptureClient->Release();
        pAudioClient->Release();
        CoTaskMemFree(pwfx);
        CoUninitialize();
        return -1;
    }

    UINT32 packetLength = 0;
    BYTE* pData;
    DWORD flags;
    UINT32 numFrames;
    uint32_t totalDataSize = 0;
    bool isRecording = true;
    
    std::cout << "Recording... Press Ctrl+C to stop." << std::endl;
    std::cout << "If Ctrl+C doesn't work, close this window after a few seconds." << std::endl;

    // Setup signal handler to gracefully exit on Ctrl+C
    SetConsoleCtrlHandler([](DWORD ctrlType) -> BOOL {
        if (ctrlType == CTRL_C_EVENT) {
            return TRUE; // We handled the event
        }
        return FALSE;
    }, TRUE);

    // Record for 30 seconds maximum (as a safety measure)
    DWORD startTime = GetTickCount();
    DWORD maxDuration = 30 * 1000; // 30 seconds in milliseconds

    while (isRecording && (GetTickCount() - startTime < maxDuration)) {
        hr = pCaptureClient->GetNextPacketSize(&packetLength);
        if (FAILED(hr)) {
            std::cerr << "GetNextPacketSize failed: " << std::hex << hr << std::endl;
            break;
        }

        if (packetLength == 0) {
            Sleep(10);
            continue;
        }

        hr = pCaptureClient->GetBuffer(&pData, &numFrames, &flags, nullptr, nullptr);
        if (FAILED(hr)) {
            std::cerr << "GetBuffer failed: " << std::hex << hr << std::endl;
            break;
        }

        if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
            // Silent data, fill with zeros
            std::vector<BYTE> silentData(numFrames * pwfx->nBlockAlign, 0);
            outFile.write(reinterpret_cast<const char*>(silentData.data()), silentData.size());
            totalDataSize += silentData.size();
        } else {
            // Handle float audio format conversion if needed
            if (isFloat && pwfx->wBitsPerSample == 32) {
                float* floatData = reinterpret_cast<float*>(pData);
                UINT32 numSamples = numFrames * pwfx->nChannels;
                std::vector<int16_t> pcmData(numSamples);
                
                // Convert float (-1.0 to 1.0) to int16 (-32768 to 32767)
                for (UINT32 i = 0; i < numSamples; i++) {
                    float sample = floatData[i];
                    // Clamp the float value between -1.0 and 1.0
                    if (sample > 1.0f) sample = 1.0f;
                    if (sample < -1.0f) sample = -1.0f;
                    // Convert to int16
                    pcmData[i] = static_cast<int16_t>(sample * 32767.0f);
                }
                
                // Write the converted PCM data
                outFile.write(reinterpret_cast<const char*>(pcmData.data()), numSamples * sizeof(int16_t));
                totalDataSize += numSamples * sizeof(int16_t);
            } else {
                // Just write the original data if it's already PCM
                UINT32 bytesToWrite = numFrames * pwfx->nBlockAlign;
                outFile.write(reinterpret_cast<const char*>(pData), bytesToWrite);
                totalDataSize += bytesToWrite;
            }
            
            if (!outFile) {
                std::cerr << "Failed to write audio data!" << std::endl;
                break;
            }
        }

        hr = pCaptureClient->ReleaseBuffer(numFrames);
        if (FAILED(hr)) {
            std::cerr << "ReleaseBuffer failed: " << std::hex << hr << std::endl;
            break;
        }
        
        // Check if file is still good
        if (!outFile) {
            std::cerr << "Output file error occurred!" << std::endl;
            break;
        }
    }

    // Stop recording
    hr = pAudioClient->Stop();
    if (FAILED(hr)) {
        std::cerr << "AudioClient Stop failed: " << std::hex << hr << std::endl;
    }

    // Make sure we have valid data
    if (totalDataSize > 0 && outFile) {
        // Write the actual WAV header now that we know the data size
        outFile.seekp(0, std::ios::beg);
        if (!outFile) {
            std::cerr << "Failed to seek to beginning of file!" << std::endl;
        } else {
            WriteWAVHeader(outFile, pwfx, totalDataSize);
            if (!outFile) {
                std::cerr << "Failed to write WAV header!" << std::endl;
            }
        }
    } else {
        std::cerr << "No audio data was captured or file error occurred!" << std::endl;
    }

    outFile.flush();
    outFile.close();

    pCaptureClient->Release();
    pAudioClient->Release();
    CoTaskMemFree(pwfx);
    CoUninitialize();

    std::cout << "Recording stopped. Total data size: " << totalDataSize << " bytes" << std::endl;
    std::cout << "Saved to " << outputPath << std::endl;
    
    // Test if the file is accessible and valid
    std::ifstream testFile(outputPath, std::ios::binary);
    if (testFile) {
        WAVHeader testHeader;
        testFile.read(reinterpret_cast<char*>(&testHeader), sizeof(WAVHeader));
        if (testFile) {
            std::cout << "WAV file successfully created and verified." << std::endl;
            std::cout << "  RIFF header: " << testHeader.riff[0] << testHeader.riff[1] 
                      << testHeader.riff[2] << testHeader.riff[3] << std::endl;
            std::cout << "  File size: " << (testHeader.chunkSize + 8) << " bytes" << std::endl;
            std::cout << "  Channels: " << testHeader.numChannels << std::endl;
            std::cout << "  Sample rate: " << testHeader.sampleRate << " Hz" << std::endl;
        } else {
            std::cerr << "Created file appears to be invalid - couldn't read header." << std::endl;
        }
        testFile.close();
    } else {
        std::cerr << "Unable to open the created file for verification." << std::endl;
    }

    return 0;
}