// Minimal Windows Application Loopback Audio Capture with process exclusion.
// Based on Microsoft's ApplicationLoopback sample.
// Outputs raw float32 stereo PCM at 48kHz to stdout.
// Usage: loopback_capture.exe <pid_to_exclude>

#include <windows.h>
#include <mmdeviceapi.h>
#include <audioclient.h>
#include <audioclientactivationparams.h>
#include <mfapi.h>
#include <stdio.h>
#include <fcntl.h>
#include <io.h>

#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mmdevapi.lib")

static IAudioClient *g_pAudioClient = NULL;
static IAudioCaptureClient *g_pCaptureClient = NULL;
static volatile bool g_running = true;

BOOL WINAPI ConsoleHandler(DWORD signal) {
    if (signal == CTRL_C_EVENT || signal == CTRL_BREAK_EVENT || signal == CTRL_CLOSE_EVENT) {
        g_running = false;
        return TRUE;
    }
    return FALSE;
}

// Callback for ActivateAudioInterfaceAsync
class CompletionHandler :
    public IActivateAudioInterfaceCompletionHandler, public IAgileObject {
    LONG m_ref = 1;
    HANDLE m_event;
    HRESULT m_hr = E_FAIL;
    IUnknown *m_pResult = nullptr;
public:
    CompletionHandler(HANDLE event) : m_event(event) {}
    HRESULT GetResult(IUnknown **pp) { *pp = m_pResult; return m_hr; }
    ULONG STDMETHODCALLTYPE AddRef() { return InterlockedIncrement(&m_ref); }
    ULONG STDMETHODCALLTYPE Release() {
        ULONG r = InterlockedDecrement(&m_ref);
        if (r == 0) delete this;
        return r;
    }
    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppv) {
        if (riid == __uuidof(IUnknown) ||
            riid == __uuidof(IActivateAudioInterfaceCompletionHandler) ||
            riid == __uuidof(IAgileObject)) {
            *ppv = static_cast<IActivateAudioInterfaceCompletionHandler*>(this);
            AddRef();
            return S_OK;
        }
        *ppv = nullptr;
        return E_NOINTERFACE;
    }
    HRESULT STDMETHODCALLTYPE ActivateCompleted(IActivateAudioInterfaceAsyncOperation *op) {
        op->GetActivateResult(&m_hr, &m_pResult);
        SetEvent(m_event);
        return S_OK;
    }
};

int wmain(int argc, wchar_t *argv[]) {
    if (argc < 2) {
        fwprintf(stderr, L"Usage: loopback_capture.exe <pid_to_exclude> [sample_rate]\n");
        return 1;
    }
    DWORD excludePid = (DWORD)_wtoi(argv[1]);
    if (excludePid == 0) {
        fwprintf(stderr, L"Invalid PID\n");
        return 1;
    }
    DWORD targetRate = (argc >= 3) ? (DWORD)_wtoi(argv[2]) : 0;

    _setmode(_fileno(stdout), _O_BINARY);
    SetConsoleCtrlHandler(ConsoleHandler, TRUE);

    HRESULT hr;
    hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (FAILED(hr)) { fwprintf(stderr, L"CoInit failed\n"); return 1; }
    hr = MFStartup(MF_VERSION);
    if (FAILED(hr)) { fwprintf(stderr, L"MFStartup failed\n"); return 1; }

    // Set up activation params for process-exclude loopback
    AUDIOCLIENT_ACTIVATION_PARAMS acParams = {};
    acParams.ActivationType = AUDIOCLIENT_ACTIVATION_TYPE_PROCESS_LOOPBACK;
    acParams.ProcessLoopbackParams.TargetProcessId = excludePid;
    acParams.ProcessLoopbackParams.ProcessLoopbackMode =
        PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE;

    PROPVARIANT activateParams = {};
    activateParams.vt = VT_BLOB;
    activateParams.blob.cbSize = sizeof(acParams);
    activateParams.blob.pBlobData = (BYTE*)&acParams;

    HANDLE completionEvent = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    auto *handler = new CompletionHandler(completionEvent);
    IActivateAudioInterfaceAsyncOperation *asyncOp = nullptr;

    hr = ActivateAudioInterfaceAsync(
        VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK,
        __uuidof(IAudioClient),
        &activateParams,
        handler,
        &asyncOp
    );
    if (FAILED(hr)) {
        fwprintf(stderr, L"ActivateAudioInterfaceAsync failed: 0x%08x\n", hr);
        return 1;
    }

    DWORD waitResult = WaitForSingleObject(completionEvent, 5000);
    if (waitResult != WAIT_OBJECT_0) {
        fwprintf(stderr, L"Activation timed out\n");
        return 1;
    }

    IUnknown *punkResult = nullptr;
    hr = handler->GetResult(&punkResult);
    if (FAILED(hr) || !punkResult) {
        fwprintf(stderr, L"Activation failed: 0x%08x\n", hr);
        return 1;
    }
    hr = punkResult->QueryInterface(__uuidof(IAudioClient), (void**)&g_pAudioClient);
    if (FAILED(hr)) {
        fwprintf(stderr, L"QueryInterface IAudioClient failed\n");
        return 1;
    }

    // Process loopback doesn't support GetMixFormat directly.
    // Get the mix format from the default audio endpoint instead.
    WAVEFORMATEX *pwfx = nullptr;
    {
        IMMDeviceEnumerator *pEnum = nullptr;
        IMMDevice *pDev = nullptr;
        IAudioClient *pTmpClient = nullptr;
        hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
                              __uuidof(IMMDeviceEnumerator), (void**)&pEnum);
        if (SUCCEEDED(hr))
            hr = pEnum->GetDefaultAudioEndpoint(eRender, eConsole, &pDev);
        if (SUCCEEDED(hr))
            hr = pDev->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**)&pTmpClient);
        if (SUCCEEDED(hr))
            hr = pTmpClient->GetMixFormat(&pwfx);
        if (pTmpClient) pTmpClient->Release();
        if (pDev) pDev->Release();
        if (pEnum) pEnum->Release();
    }
    if (FAILED(hr) || !pwfx) {
        fwprintf(stderr, L"GetMixFormat failed: 0x%08x\n", hr);
        return 1;
    }

    // Override sample rate if a target was specified by the renderer.
    // WASAPI shared mode will resample automatically.
    if (targetRate > 0 && targetRate != pwfx->nSamplesPerSec) {
        fwprintf(stderr, L"overriding rate %d -> %d\n", pwfx->nSamplesPerSec, targetRate);
        pwfx->nSamplesPerSec = targetRate;
        pwfx->nAvgBytesPerSec = targetRate * pwfx->nBlockAlign;
    }

    fwprintf(stderr, L"format: ch=%d rate=%d bits=%d\n",
        pwfx->nChannels, pwfx->nSamplesPerSec, pwfx->wBitsPerSample);

    // Initialize in shared mode, polling (no event callback — simpler, no deadlocks)
    hr = g_pAudioClient->Initialize(
        AUDCLNT_SHAREMODE_SHARED,
        AUDCLNT_STREAMFLAGS_LOOPBACK,
        200000, // 20ms buffer
        0,
        pwfx,
        nullptr
    );
    if (FAILED(hr)) {
        fwprintf(stderr, L"Initialize failed: 0x%08x\n", hr);
        CoTaskMemFree(pwfx);
        return 1;
    }

    hr = g_pAudioClient->GetService(__uuidof(IAudioCaptureClient), (void**)&g_pCaptureClient);
    if (FAILED(hr)) {
        fwprintf(stderr, L"GetService failed\n");
        CoTaskMemFree(pwfx);
        return 1;
    }

    hr = g_pAudioClient->Start();
    if (FAILED(hr)) {
        fwprintf(stderr, L"Start failed: 0x%08x\n", hr);
        CoTaskMemFree(pwfx);
        return 1;
    }

    fwprintf(stderr, L"capturing (excluding pid %u)...\n", excludePid);

    // Allocate a silence buffer once (reused for silent packets)
    UINT32 maxSilenceBytes = pwfx->nSamplesPerSec * pwfx->nBlockAlign; // 1s worth
    BYTE *silenceBuf = (BYTE*)calloc(maxSilenceBytes, 1);

    // Capture loop — poll every 2ms for low latency
    while (g_running) {
        Sleep(2);
        UINT32 numFrames = 0;

        hr = g_pCaptureClient->GetNextPacketSize(&numFrames);
        if (FAILED(hr)) break;

        while (numFrames > 0) {
            BYTE *pData = nullptr;
            UINT32 framesAvailable = 0;
            DWORD flags = 0;
            hr = g_pCaptureClient->GetBuffer(&pData, &framesAvailable, &flags, nullptr, nullptr);
            if (FAILED(hr)) break;

            UINT32 bytes = framesAvailable * pwfx->nBlockAlign;
            if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
                if (bytes <= maxSilenceBytes)
                    fwrite(silenceBuf, 1, bytes, stdout);
            } else {
                fwrite(pData, 1, bytes, stdout);
            }

            g_pCaptureClient->ReleaseBuffer(framesAvailable);
            g_pCaptureClient->GetNextPacketSize(&numFrames);
        }
        fflush(stdout);
    }

    free(silenceBuf);

    g_pAudioClient->Stop();
    g_pCaptureClient->Release();
    g_pAudioClient->Release();
    if (asyncOp) asyncOp->Release();
    handler->Release();
    CloseHandle(completionEvent);
    CoTaskMemFree(pwfx);
    MFShutdown();
    CoUninitialize();
    return 0;
}
