// sb-send: Sonobus audio sender

// Functions dealing with Windows audio capture with WASAPI.

#include "sb-send.h"

#ifdef _WIN32
int wasapi_list_devices(void) {
    IMMDeviceEnumerator *enumerator = NULL;
    IMMDeviceCollection *collection = NULL;

    HRESULT hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
    if (FAILED(hr)) { fprintf(stderr, "COM init failed: 0x%08lx\n", hr); return -1; }

    hr = CoCreateInstance(&CLSID_MMDeviceEnumerator, NULL, CLSCTX_ALL,
                           &IID_IMMDeviceEnumerator, (void**)&enumerator);
    if (FAILED(hr)) { fprintf(stderr, "Device enumerator failed: 0x%08lx\n", hr); CoUninitialize(); return -1; }

    hr = enumerator->lpVtbl->EnumAudioEndpoints(enumerator, eCapture,
                                                  DEVICE_STATE_ACTIVE, &collection);
    if (FAILED(hr)) { fprintf(stderr, "EnumAudioEndpoints failed: 0x%08lx\n", hr);
        enumerator->lpVtbl->Release(enumerator); CoUninitialize(); return -1; }

    UINT count = 0;
    collection->lpVtbl->GetCount(collection, &count);
    fprintf(stderr, "Audio capture devices (%u):\n", count);

    for (UINT i = 0; i < count; i++) {
        IMMDevice *device = NULL;
        collection->lpVtbl->Item(collection, i, &device);

        IPropertyStore *props = NULL;
        device->lpVtbl->OpenPropertyStore(device, STGM_READ, &props);

        PROPVARIANT name;
        PropVariantInit(&name);
        props->lpVtbl->GetValue(props, &PKEY_Device_FriendlyName, &name);

        if (name.pwszVal) {
            char utf8[256];
            WideCharToMultiByte(CP_UTF8, 0, name.pwszVal, -1, utf8, sizeof(utf8), NULL, NULL);
            fprintf(stderr, "  [%u] %s\n", i, utf8);
        }

        PropVariantClear(&name);
        props->lpVtbl->Release(props);
        device->lpVtbl->Release(device);
    }

    collection->lpVtbl->Release(collection);
    enumerator->lpVtbl->Release(enumerator);
    return 0;
}

int wasapi_init_capture(AppState *state) {
    IMMDeviceEnumerator *enumerator = NULL;
    IMMDevice *device = NULL;

    HRESULT hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
    if (FAILED(hr)) { fprintf(stderr, "COM init failed: 0x%08lx\n", hr); return -1; }

    hr = CoCreateInstance(&CLSID_MMDeviceEnumerator, NULL, CLSCTX_ALL,
                           &IID_IMMDeviceEnumerator, (void**)&enumerator);
    if (FAILED(hr)) { fprintf(stderr, "Device enumerator failed: 0x%08lx\n", hr); return -1; }

    if (state->audio_device[0] && strcmp(state->audio_device, "default") != 0) {
        IMMDeviceCollection *coll = NULL;
        enumerator->lpVtbl->EnumAudioEndpoints(enumerator, eCapture,
                                                 DEVICE_STATE_ACTIVE, &coll);
        UINT count = 0;
        coll->lpVtbl->GetCount(coll, &count);
        for (UINT i = 0; i < count; i++) {
            IMMDevice *dev = NULL;
            coll->lpVtbl->Item(coll, i, &dev);
            IPropertyStore *props = NULL;
            dev->lpVtbl->OpenPropertyStore(dev, STGM_READ, &props);
            PROPVARIANT pname;
            PropVariantInit(&pname);
            props->lpVtbl->GetValue(props, &PKEY_Device_FriendlyName, &pname);
            if (pname.pwszVal) {
                char utf8[256];
                WideCharToMultiByte(CP_UTF8, 0, pname.pwszVal, -1, utf8, sizeof(utf8), NULL, NULL);
                if (_stricmp(utf8, state->audio_device) == 0) {
                    device = dev;
                    PropVariantClear(&pname);
                    props->lpVtbl->Release(props);
                    break;
                }
            }
            PropVariantClear(&pname);
            props->lpVtbl->Release(props);
            dev->lpVtbl->Release(dev);
        }
        coll->lpVtbl->Release(coll);
    } else {
        enumerator->lpVtbl->GetDefaultAudioEndpoint(enumerator, eCapture, eConsole, &device);
    }

    enumerator->lpVtbl->Release(enumerator);

    if (!device) { fprintf(stderr, "No audio capture device found\n"); return -1; }

    IAudioClient *client = NULL;
    hr = device->lpVtbl->Activate(device, &IID_IAudioClient, CLSCTX_ALL, NULL, (void**)&client);
    device->lpVtbl->Release(device);
    if (FAILED(hr)) { fprintf(stderr, "AudioClient activate failed: 0x%08lx\n", hr); return -1; }

    WAVEFORMATEX *mix_format = NULL;
    hr = client->lpVtbl->GetMixFormat(client, &mix_format);
    if (FAILED(hr)) { fprintf(stderr, "GetMixFormat failed: 0x%08lx\n", hr); return -1; }

    fprintf(stderr, "WASAPI mix format: %lu Hz %d ch %d bit\n",
           (unsigned long)mix_format->nSamplesPerSec, mix_format->nChannels, mix_format->wBitsPerSample);

    state->wasapi_channels = mix_format->nChannels;

    WAVEFORMATEX wf = {0};
    wf.wFormatTag = WAVE_FORMAT_IEEE_FLOAT;
    wf.nChannels = state->channels;
    wf.nSamplesPerSec = SAMPLE_RATE;
    wf.wBitsPerSample = 32;
    wf.nBlockAlign = wf.nChannels * wf.wBitsPerSample / 8;
    wf.nAvgBytesPerSec = wf.nSamplesPerSec * wf.nBlockAlign;

    REFERENCE_TIME buffer_duration = 50000;
    hr = client->lpVtbl->Initialize(client, AUDCLNT_SHAREMODE_SHARED,
                                      AUDCLNT_STREAMFLAGS_EVENTCALLBACK,
                                      buffer_duration, buffer_duration, &wf, NULL);
    if (SUCCEEDED(hr)) {
        state->wasapi_channels = state->channels;
        fprintf(stderr, "WASAPI: negotiated %d Hz %d ch IEEE_FLOAT\n", SAMPLE_RATE, state->channels);
    } else {
        WAVEFORMATEX wf2 = {0};
        wf2.wFormatTag = mix_format->wFormatTag;
        wf2.nChannels = mix_format->nChannels;
        wf2.nSamplesPerSec = SAMPLE_RATE;
        wf2.wBitsPerSample = mix_format->wBitsPerSample;
        wf2.nBlockAlign = wf2.nChannels * wf2.wBitsPerSample / 8;
        wf2.nAvgBytesPerSec = SAMPLE_RATE * wf2.nBlockAlign;
        hr = client->lpVtbl->Initialize(client, AUDCLNT_SHAREMODE_SHARED,
                                          AUDCLNT_STREAMFLAGS_EVENTCALLBACK,
                                          buffer_duration, buffer_duration, &wf2, NULL);
        if (SUCCEEDED(hr)) {
            fprintf(stderr, "WASAPI: negotiated %d Hz %d ch (mix format tag)\n",
                   SAMPLE_RATE, mix_format->nChannels);
        } else {
            hr = client->lpVtbl->Initialize(client, AUDCLNT_SHAREMODE_SHARED,
                                              AUDCLNT_STREAMFLAGS_EVENTCALLBACK,
                                              buffer_duration, buffer_duration, mix_format, NULL);
            if (FAILED(hr)) {
                fprintf(stderr, "AudioClient init failed: 0x%08lx\n", hr);
                CoTaskMemFree(mix_format);
                return -1;
            }
            if (mix_format->nSamplesPerSec != SAMPLE_RATE) {
                fprintf(stderr, "WASAPI: WARNING device at %lu Hz, expected %d Hz - audio will be mis-timed\n",
                       (unsigned long)mix_format->nSamplesPerSec, SAMPLE_RATE);
            }
            state->wasapi_channels = mix_format->nChannels;
            fprintf(stderr, "WASAPI: using device native format %lu Hz %d ch\n",
                   (unsigned long)mix_format->nSamplesPerSec, mix_format->nChannels);
        }
    }
    CoTaskMemFree(mix_format);

    if (state->wasapi_channels > 2) {
        fprintf(stderr, "WASAPI: device has %d channels, max supported is 2\n", state->wasapi_channels);
        client->lpVtbl->Release(client);
        return -1;
    }

    IAudioCaptureClient *capture = NULL;
    hr = client->lpVtbl->GetService(client, &IID_IAudioCaptureClient, (void**)&capture);
    if (FAILED(hr)) { fprintf(stderr, "CaptureClient failed: 0x%08lx\n", hr); return -1; }

    state->wasapi_client = client;
    state->wasapi_capture = capture;
    state->use_wasapi = true;
    return 0;
}

void *wasapi_reader_thread(void *arg) {
    AppState *state = (AppState *)arg;
    AudioRing *ring = &g_ring;
    IAudioClient *client = state->wasapi_client;
    IAudioCaptureClient *capture = state->wasapi_capture;
    int wch = state->wasapi_channels;

    HANDLE event = CreateEvent(NULL, FALSE, FALSE, NULL);
    if (!event) { ring_signal_done(ring); return NULL; }

    client->lpVtbl->SetEventHandle(client, event);
    client->lpVtbl->Start(client);

    while (g_running && !atomic_load(&ring->done)) {
        DWORD wr = WaitForSingleObject(event, 200);
        if (wr == WAIT_FAILED || wr == WAIT_TIMEOUT) {
            continue;
        }

        BYTE *data = NULL;
        UINT32 frames = 0;
        DWORD flags = 0;
        HRESULT hr = capture->lpVtbl->GetBuffer(capture, &data, &frames, &flags, NULL, NULL);
        if (FAILED(hr) || !data || frames == 0) {
            continue;
        }

        if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
            float zeros[state->frame_size * 2];
            for (int zi = 0; zi < state->frame_size * 2; zi++) {
                zeros[zi] = 0.0f;
            }
            for (UINT32 remaining = frames; remaining > 0; ) {
                int chunk;
                if (remaining > (UINT32)state->frame_size) {
                    chunk = state->frame_size;
                } else {
                    chunk = (int)remaining;
                }
                ring_put_float(ring, zeros, chunk * wch);
                remaining -= chunk;
            }
        } else {
            float *fdata = (float *)data;
            for (UINT32 off = 0; off < frames; ) {
                int chunk;
                if ((frames - off) > (UINT32)state->frame_size) {
                    chunk = state->frame_size;
                } else {
                    chunk = (int)(frames - off);
                }
                ring_put_float(ring, fdata + off * wch, chunk * wch);
                off += chunk;
            }
        }

        capture->lpVtbl->ReleaseBuffer(capture, frames);
    }

    client->lpVtbl->Stop(client);
    CloseHandle(event);
    ring_signal_done(ring);
    return NULL;
}
#endif
// EOF
