#include <stdlib.h>
#include <stdbool.h>
#include <stdio.h>
#include <windows.h>
#include <initguid.h>
#include <mmsystem.h>
#include <mmdeviceapi.h>
#include <audioclient.h>
#include <avrt.h>
#include <functiondiscoverykeys_devpkey.h>

struct nd_pair
{
    wchar_t name[BUFSIZ];
    IMMDevice *device;
};

HANDLE termination = NULL;

BOOL CtrlHandler(DWORD type)
{
    (void)type;

    if (termination != NULL)
    {
        SetEvent(termination);
    }

    return TRUE;
}

int list_devices(struct nd_pair **pairs)
{
    int result = -1;
    HRESULT res = 0;
    IMMDeviceEnumerator *enumerator = NULL;
    IMMDeviceCollection *collection = NULL;
    UINT count = 0;

    if (pairs == NULL)
    {
        goto out;
    }

    if ((res = CoCreateInstance(&CLSID_MMDeviceEnumerator, NULL, CLSCTX_ALL,
            &IID_IMMDeviceEnumerator, (void **)&enumerator)) != S_OK)
    {
        fprintf(stderr, "Failed to create IMMDeviceEnumerator instance! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = enumerator->lpVtbl->EnumAudioEndpoints(enumerator, eRender,
            DEVICE_STATE_ACTIVE, &collection)) != S_OK)
    {
        fprintf(stderr, "Failed to enumerate audio endpoints! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = collection->lpVtbl->GetCount(collection, &count)) != S_OK)
    {
        fprintf(stderr, "Failed to get number of audio endpoints! [0x%08lx]\n", res);
        goto out;
    }

    *pairs = malloc(count * sizeof(struct nd_pair));

    ZeroMemory(*pairs, count * sizeof(struct nd_pair));

    if (*pairs == NULL)
    {
        fprintf(stderr, "Failed to allocate memory for device list!\n");
        goto out;
    }

    for (UINT i = 0; i < count; i++)
    {
        IMMDevice *device = NULL;

        if ((res = collection->lpVtbl->Item(collection, i, &device)) != S_OK)
        {
            fprintf(stderr, "Failed to get audio device! [0x%08lx]\n", res);
            goto out;
        }

        (*pairs)[i].device = device;

        IPropertyStore *props = NULL;

        if ((res = device->lpVtbl->OpenPropertyStore(device, STGM_READ, &props)) != S_OK)
        {
            fprintf(stderr, "Failed to open property store! [0x%08lx]\n", res);
            goto out;
        }

        PROPVARIANT pv;
        PropVariantInit(&pv);

        if ((res = props->lpVtbl->GetValue(props, &PKEY_Device_FriendlyName, &pv)) != S_OK)
        {
            fprintf(stderr, "Failed to get device friendly name! [0x%08lx]\n", res);
            PropVariantClear(&pv);
            props->lpVtbl->Release(props);
            goto out;
        }

        if (pv.vt != VT_LPWSTR)
        {
            fprintf(stderr, "Unexpected property type!\n");
            PropVariantClear(&pv);
            props->lpVtbl->Release(props);
            goto out;
        }

        memset((*pairs)[i].name, 0, BUFSIZ * sizeof(wchar_t));
        wcsncpy((*pairs)[i].name, pv.pwszVal, wcsnlen(pv.pwszVal, BUFSIZ));

        PropVariantClear(&pv);
        props->lpVtbl->Release(props);
    }

    result = count;

out:
    if (result < 0 && *pairs != NULL)
    {
        for (UINT i = 0; i < count; i++)
        {
            if ((*pairs)[i].device != NULL)
            {
                (*pairs)[i].device->lpVtbl->Release((*pairs)[i].device);
            }
        }

        free(*pairs);
    }

    if (collection != NULL)
    {
        collection->lpVtbl->Release(collection);
    }

    if (enumerator != NULL)
    {
        enumerator->lpVtbl->Release(enumerator);
    }

    return result;
}

void clear_devices(struct nd_pair **pairs, int count)
{
    if (pairs == NULL)
    {
        return;
    }
    
    for (int i = 0; i < count; i++)
    {
        (*pairs)[i].device->lpVtbl->Release((*pairs)[i].device);
    }

    free(*pairs);
    *pairs = NULL;
}

bool mirror(IMMDevice *monitor, IMMDevice *output)
{
    bool result = false;
    HRESULT res = 0;
    HANDLE event = NULL;
    DWORD task_index = 0;
    HANDLE task = NULL;
    IAudioClient *mon_client = NULL;
    IAudioClient *out_client = NULL;
    IAudioCaptureClient *capture = NULL;
    IAudioRenderClient *render = NULL;
    WAVEFORMATEX *format = NULL;
    REFERENCE_TIME period = 0;
    UINT32 buffer_size = 0;

    if ((event = CreateEvent(NULL, FALSE, FALSE, NULL)) == NULL)
    {
        fprintf(stderr, "Failed to create event!\n");
        goto out;
    }

    if ((task = AvSetMmThreadCharacteristics(TEXT("Audio"), &task_index)) == NULL)
    {
        fprintf(stderr, "Failed to set thread characteristics!\n");
        goto out;
    }

    if ((res = monitor->lpVtbl->Activate(monitor, &IID_IAudioClient, CLSCTX_ALL,
            NULL, (void **)&mon_client)) != S_OK)
    {
        fprintf(stderr, "Failed to activate monitor device! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = mon_client->lpVtbl->GetDevicePeriod(mon_client, &period, NULL)) != S_OK)
    {
        fprintf(stderr, "Failed to get monitor client period! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = mon_client->lpVtbl->GetMixFormat(mon_client, &format)) != S_OK)
    {
        fprintf(stderr, "Failed to get monitor client mix format! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = mon_client->lpVtbl->Initialize(mon_client, AUDCLNT_SHAREMODE_SHARED,
            AUDCLNT_STREAMFLAGS_LOOPBACK, 0, 0, format, 0)) != S_OK)
    {
        fprintf(stderr, "Failed to initialize monitor client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = mon_client->lpVtbl->GetService(mon_client, &IID_IAudioCaptureClient,
            (void **)&capture)) != S_OK)
    {
        fprintf(stderr, "Failed to get capture service from monitor client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = output->lpVtbl->Activate(output, &IID_IAudioClient, CLSCTX_ALL,
            NULL, (void **)&out_client)) != S_OK)
    {
        fprintf(stderr, "Failed to activate output device! [0x%08lx]\n", res);
        goto out;
    }

    if ((res =out_client->lpVtbl->Initialize(out_client, AUDCLNT_SHAREMODE_SHARED,
            AUDCLNT_STREAMFLAGS_EVENTCALLBACK, period, 0, format, 0)) != S_OK)
    {
        fprintf(stderr, "Failed to initialize output client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = out_client->lpVtbl->GetBufferSize(out_client, &buffer_size)) != S_OK)
    {
        fprintf(stderr, "Failed to get buffer size of output client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = out_client->lpVtbl->GetService(out_client, &IID_IAudioRenderClient,
            (void **)&render)) != S_OK)
    {
        fprintf(stderr, "Failed to get render service from output client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = out_client->lpVtbl->SetEventHandle(out_client, event)) != S_OK)
    {
        fprintf(stderr, "Failed to associate event with output client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = mon_client->lpVtbl->Start(mon_client)) != S_OK)
    {
        fprintf(stderr, "Failed to start monitor client! [0x%08lx]\n", res);
        goto out;
    }

    if ((res = out_client->lpVtbl->Start(out_client)) != S_OK)
    {
        fprintf(stderr, "Failed to start output client! [0x%08lx]\n", res);
        goto out;
    }

    HANDLE events[] = {termination, event};

    for (;;)
    {
        switch (WaitForMultipleObjects(ARRAYSIZE(events), events, FALSE, INFINITE))
        {
            case WAIT_OBJECT_0: // termination
                result = true;
                goto out;

            case WAIT_OBJECT_0 + 1: // event
                break;

            default:
                fprintf(stderr, "Failed to wait for events!\n");
                goto out;
        }
 
        for (;;)
        {
            UINT32 padding = 0;

            if ((res = out_client->lpVtbl->GetCurrentPadding(out_client, &padding)) != S_OK)
            {
                fprintf(stderr, "Failed to get current padding of output client! [0x%08lx]\n", res);
                goto out;
            }

            UINT32 packet_size = 0;

            if ((res = capture->lpVtbl->GetNextPacketSize(capture, &packet_size)) != S_OK)
            {
                fprintf(stderr, "Failed to get next packet size of monitor client! [0x%08lx]\n", res);
                goto out;
            }

            BYTE *mon_data = NULL;
            BYTE *out_data = NULL;
            UINT32 frames = 0;
            DWORD flags = 0;
     
            if (packet_size == 0)
            {
                // no input data, play silence

                if ((res = render->lpVtbl->GetBuffer(render, buffer_size - padding, &out_data)) != S_OK)
                {
                    fprintf(stderr, "Failed to get render buffer of output client! [0x%08lx]\n", res);
                    goto out;
                }

                if ((res = render->lpVtbl->ReleaseBuffer(render, buffer_size - padding,
                        AUDCLNT_BUFFERFLAGS_SILENT)) != S_OK)
                {
                    fprintf(stderr, "Failed to release render buffer of output client! [0x%08lx]\n", res);
                    goto out;
                }

                break;
            }
            else if (packet_size > buffer_size - padding)
            {
                // output buffer is full, we have to wait
                break;
            }
            else
            {
                if ((res = capture->lpVtbl->GetBuffer(capture, &mon_data, &frames, &flags,
                        NULL, NULL)) != S_OK)
                {
                    fprintf(stderr, "Failed to get capture buffer of monitor client! [0x%08lx]\n", res);
                    goto out;
                }

                if ((res = render->lpVtbl->GetBuffer(render, frames, &out_data)) != S_OK)
                {
                    fprintf(stderr, "Failed to get render buffer of output client! [0x%08lx]\n", res);
                    goto out;
                }

                CopyMemory(out_data, mon_data, frames * format->nBlockAlign);

                if ((res = capture->lpVtbl->ReleaseBuffer(capture, frames)) != S_OK)
                {
                    fprintf(stderr, "Failed to release capture buffer of monitor client! [0x%08lx]\n", res);
                    goto out;
                }

                if ((res = render->lpVtbl->ReleaseBuffer(render, frames, 0)) != S_OK)
                {
                    fprintf(stderr, "Failed to release render buffer of output client! [0x%08lx]\n", res);
                    goto out;
                }
            }
        }
    }

out:
    if (render != NULL)
    {
        render->lpVtbl->Release(render);
    }

    if (capture != NULL)
    {
        capture->lpVtbl->Release(capture);
    }

    if (format != NULL)
    {
        CoTaskMemFree(format);
    }

    if (out_client != NULL)
    {
        out_client->lpVtbl->Stop(out_client);
        out_client->lpVtbl->Release(out_client);
    }

    if (mon_client != NULL)
    {
        mon_client->lpVtbl->Stop(mon_client);
        mon_client->lpVtbl->Release(mon_client);
    }

    if (task != NULL)
    {
        AvRevertMmThreadCharacteristics(task);
        CloseHandle(task);
    }

    if (event != NULL)
    {
        CloseHandle(event);
    }

    return result;
}

bool run()
{
    bool result = false;
    struct nd_pair *pairs = NULL;
    int count = list_devices(&pairs);

    if (count <= 0)
    {
        fprintf(stderr, "No playback audio devices detected!\n");
        goto out;
    }

    printf("Detected playback audio devices:\n\n");

    for (int i = 0; i < count; i++)
    {
        printf("[%d] %ls\n", i, pairs[i].name);
    }

    printf("\n");

    char buf[5] = "";

    printf("Select device to monitor: ");

    char *res = fgets(buf, ARRAYSIZE(buf), stdin);
    int mon = strtol(buf, NULL, 10);

    while (res != buf || mon < 0 || mon >= count)
    {
        printf("Select device to monitor: ");

        res = fgets(buf, ARRAYSIZE(buf), stdin);
        mon = strtol(buf, NULL, 10);
    }

    printf("Select output device: ");

    res = fgets(buf, ARRAYSIZE(buf), stdin);
    int out = strtol(buf, NULL, 10);

    while (res != buf || out < 0 || out >= count || out == mon)
    {
        printf("Select output device: ");

        res = fgets(buf, ARRAYSIZE(buf), stdin);
        out = strtol(buf, NULL, 10);
    }

    printf("\nMirroring:\n");
    printf("%ls\n", pairs[mon].name);
    printf("  -> %ls\n\n", pairs[out].name);

    result = mirror(pairs[mon].device, pairs[out].device);

out:
    clear_devices(&pairs, count);
    return result;
}

int main(void)
{
    if ((termination = CreateEvent(NULL, FALSE, FALSE, NULL)) == NULL)
    {
        fprintf(stderr, "Failed to create termination event!\n");
        return EXIT_FAILURE;
    }

    if (SetConsoleCtrlHandler((PHANDLER_ROUTINE)CtrlHandler, TRUE) != TRUE)
    {
        fprintf(stderr, "Failed to set console control handler!\n");
        CloseHandle(termination);
        return EXIT_FAILURE;
    }

    if (CoInitialize(NULL) != S_OK)
    {
        CloseHandle(termination);
        return EXIT_FAILURE;
    }

    if (! run())
    {
        CoUninitialize();
        CloseHandle(termination);
        return EXIT_FAILURE;
    }

    CoUninitialize();
    CloseHandle(termination);
    return EXIT_SUCCESS;
}