#define UNICODE
#define _UNICODE

#include <audioclient.h>
#include <comdef.h>
#include <comip.h>
#include <fcntl.h>
#include <io.h>
#include <mmdeviceapi.h>
#include <stdio.h>
#include <stdlib.h>
#include <windows.h>

#define REFTIMES_PER_SEC 10000000
#define REFTIMES_PER_MILLISEC 10000

class CCoInitialize {
    HRESULT m_hr;

  public:
    CCoInitialize() : m_hr(CoInitialize(NULL)) {}
    ~CCoInitialize() {
        if (SUCCEEDED(m_hr))
            CoUninitialize();
    }
    operator HRESULT() const { return m_hr; }
};

_COM_SMARTPTR_TYPEDEF(IUnknown, __uuidof(IUnknown));
_COM_SMARTPTR_TYPEDEF(IMMDeviceEnumerator, __uuidof(IMMDeviceEnumerator));
_COM_SMARTPTR_TYPEDEF(IMMDevice, __uuidof(IMMDevice));
_COM_SMARTPTR_TYPEDEF(IAudioClient, __uuidof(IAudioClient));
_COM_SMARTPTR_TYPEDEF(IAudioCaptureClient, __uuidof(IAudioCaptureClient));

class DeviceChangeNotification : public IMMNotificationClient {
    volatile ULONG ref;
    volatile bool &changed;

  public:
    DeviceChangeNotification(volatile bool &changed) : changed(changed) {}

    // This is meant to be allocated on stack, so we don't actually free.
    STDMETHODIMP_(ULONG) AddRef() { return InterlockedIncrement(&ref); }
    STDMETHODIMP_(ULONG) Release() { return InterlockedDecrement(&ref); }

    STDMETHODIMP QueryInterface(REFIID iid, void **ppv) {
        if (iid == IID_IUnknown) {
            AddRef();
            *ppv = (IUnknown *)this;
        } else if (iid == __uuidof(IMMNotificationClient)) {
            AddRef();
            *ppv = (IMMNotificationClient *)this;
        } else {
            *ppv = nullptr;
            return E_NOINTERFACE;
        }
        return S_OK;
    }

    STDMETHODIMP OnDefaultDeviceChanged(EDataFlow flow, ERole role, LPCWSTR) {
        if (flow == eRender && role == eConsole) {
            changed = true;
        }
        return S_OK;
    }

    STDMETHODIMP OnDeviceAdded(LPCWSTR) { return S_OK; }
    STDMETHODIMP OnDeviceRemoved(LPCWSTR) { return S_OK; }
    STDMETHODIMP OnDeviceStateChanged(LPCWSTR, DWORD) { return S_OK; }
    STDMETHODIMP OnPropertyValueChanged(LPCWSTR, const PROPERTYKEY) { return S_OK; }
};

class EnsureCaptureStop {
    IAudioClientPtr m_client;

  public:
    EnsureCaptureStop(IAudioClientPtr pClient) : m_client(pClient) {}
    ~EnsureCaptureStop() { m_client->Stop(); }
};

class EnsureFree {
    void *m_memory;

  public:
    EnsureFree(void *memory) : m_memory(memory) {}
    ~EnsureFree() { free(m_memory); }
};

struct {
    HRESULT hr;
    LPWSTR error;
} error_table[] = {
    {AUDCLNT_E_UNSUPPORTED_FORMAT, L"Requested sound format unsupported"},
};

#define ensure(hr) ensure_(__FILE__, __LINE__, hr)
void ensure_(const char *file, int line, HRESULT hr) {
    if (FAILED(hr)) {
        _com_error err(hr);
        LPCWSTR msg = err.ErrorMessage();

        for (int i = 0; i < sizeof error_table / sizeof error_table[0]; ++i) {
            if (error_table[i].hr == hr) {
                msg = error_table[i].error;
            }
        }

        fwprintf(stderr, L"Error at %S:%d (0x%08x): %s\n", file, line, hr, msg);
        exit(1);
    }
}

[[noreturn]] void usage(char *argv0) {
    fprintf(stderr, "Usage: %s <channels> <sample rate> <bits per sample>\n", argv0);
    exit(2);
}

int parse_int_arg(int argc, char *argv[], int argn) {
    if (argn < argc) {
        int result = atoi(argv[argn]);
        if (result)
            return result;
    }
    usage(argv[0]);
}

int main(int argc, char *argv[]) {
    int channels = parse_int_arg(argc, argv, 1);
    int rate = parse_int_arg(argc, argv, 2);
    int bits = parse_int_arg(argc, argv, 3);

    if (_setmode(_fileno(stdout), _O_BINARY) == -1) {
        perror("_setmode failed");
        exit(1);
    }

    CCoInitialize comInit;
    ensure(comInit);

    WAVEFORMATEX wfx;
    wfx.wFormatTag = WAVE_FORMAT_PCM;
    wfx.nChannels = (WORD)channels;
    wfx.nSamplesPerSec = (DWORD)rate;
    wfx.wBitsPerSample = (WORD)bits;
    wfx.nBlockAlign = wfx.nChannels * wfx.wBitsPerSample / 8;
    wfx.nAvgBytesPerSec = wfx.nSamplesPerSec * wfx.nBlockAlign;

    IMMDeviceEnumeratorPtr pEnumerator;
    ensure(pEnumerator.CreateInstance(__uuidof(MMDeviceEnumerator)));

    volatile bool deviceChanged = false;
    DeviceChangeNotification deviceChangeNotification(deviceChanged);
    pEnumerator->RegisterEndpointNotificationCallback(&deviceChangeNotification);

    for (;;) {
        IMMDevicePtr pDevice;
        ensure(pEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &pDevice));

        IAudioClientPtr pClient;
        ensure(pDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void **)&pClient));

        ensure(pClient->Initialize(AUDCLNT_SHAREMODE_SHARED, AUDCLNT_STREAMFLAGS_LOOPBACK,
                                   16 * REFTIMES_PER_MILLISEC, 0, &wfx, nullptr));

        UINT32 bufferFrameCount;
        ensure(pClient->GetBufferSize(&bufferFrameCount));

        IAudioCaptureClientPtr pCapture;
        ensure(pClient->GetService(__uuidof(IAudioCaptureClient), (void **)&pCapture));

        DWORD dwDelay = (DWORD)(((double)REFTIMES_PER_SEC * bufferFrameCount / wfx.nSamplesPerSec) /
                                REFTIMES_PER_MILLISEC / 2);

        LPBYTE pSilence = (LPBYTE)malloc(bufferFrameCount * wfx.nBlockAlign);
        EnsureFree freeSilence(pSilence);
        ZeroMemory(pSilence, bufferFrameCount * wfx.nBlockAlign);

        ensure(pClient->Start());
        EnsureCaptureStop autoStop(pClient);

        while (!deviceChanged) {
            Sleep(dwDelay);

            UINT32 packetLength;
            HRESULT hrNext = pCapture->GetNextPacketSize(&packetLength);
            if (hrNext == AUDCLNT_E_DEVICE_INVALIDATED) {
                while (!deviceChanged)
                    ;
                break;
            } else {
                ensure(hrNext);
            }

            while (packetLength) {
                LPBYTE pData;
                UINT32 numFramesAvailable;
                DWORD flags;

                ensure(pCapture->GetBuffer(&pData, &numFramesAvailable, &flags, nullptr, nullptr));

                if (flags & AUDCLNT_BUFFERFLAGS_SILENT)
                    pData = pSilence;

                _write(_fileno(stdout), pData, wfx.nBlockAlign * numFramesAvailable);
                ensure(pCapture->ReleaseBuffer(numFramesAvailable));
                ensure(pCapture->GetNextPacketSize(&packetLength));
            }
        }
        deviceChanged = false;
    }
}