audio_monitor.cpp

#include "audio_monitor.h"
#include "logger.h"
#include <Propvarutil.h>

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

AudioMonitor::AudioMonitor() = default;

AudioMonitor::~AudioMonitor() {
    stop();
}

bool AudioMonitor::start() {
    HRESULT hr = CoCreateInstance(
        __uuidof(MMDeviceEnumerator),
        nullptr,
        CLSCTX_ALL,
        __uuidof(IMMDeviceEnumerator),
        reinterpret_cast<void**>(&m_enumerator));

    if (FAILED(hr)) {
        LOG_ERROR(L"Failed to create MMDeviceEnumerator");
        return false;
    }

    hr = m_enumerator->RegisterEndpointNotificationCallback(this);
    if (FAILED(hr)) {
        LOG_ERROR(L"Failed to register audio notification callback");
        m_enumerator->Release();
        m_enumerator = nullptr;
        return false;
    }

    m_started = true;
    return true;
}

void AudioMonitor::stop() {
    if (m_enumerator) {
        if (m_started) {
            m_enumerator->UnregisterEndpointNotificationCallback(this);
        }
        m_enumerator->Release();
        m_enumerator = nullptr;
    }
    m_started = false;
}

std::wstring AudioMonitor::getDeviceFriendlyName(IMMDevice* device) {
    IPropertyStore* props = nullptr;
    HRESULT hr = device->OpenPropertyStore(STGM_READ, &props);
    if (FAILED(hr)) return L"Unknown Device";

    PROPVARIANT varName;
    PropVariantInit(&varName);

    hr = props->GetValue(PKEY_Device_FriendlyName, &varName);
    std::wstring name = L"Unknown Device";
    if (SUCCEEDED(hr) && varName.vt == VT_LPWSTR) {
        name = varName.pwszVal;
    }

    PropVariantClear(&varName);
    props->Release();
    return name;
}

std::wstring AudioMonitor::deviceStateToString(DWORD state) {
    switch (state) {
        case DEVICE_STATE_ACTIVE:     return L"Active";
        case DEVICE_STATE_DISABLED:   return L"Disabled";
        case DEVICE_STATE_NOTPRESENT: return L"Not Present";
        case DEVICE_STATE_UNPLUGGED:  return L"Unplugged";
        default:                       return L"Unknown";
    }
}

std::wstring AudioMonitor::dataFlowToString(EDataFlow flow) {
    switch (flow) {
        case eRender:  return L"[OUT]";
        case eCapture: return L"[IN] ";
        default:       return L"[?]  ";
    }
}

AudioDeviceInfo AudioMonitor::getDeviceInfo(LPCWSTR deviceId, EDataFlow flow) {
    AudioDeviceInfo info;
    info.id = deviceId;
    info.isOutput = (flow == eRender);

    if (!m_enumerator) return info;

    IMMDevice* device = nullptr;
    HRESULT hr = m_enumerator->GetDevice(deviceId, &device);
    if (SUCCEEDED(hr) && device) {
        info.name = getDeviceFriendlyName(device);

        DWORD state = 0;
        device->GetState(&state);
        info.isActive = (state == DEVICE_STATE_ACTIVE);

        device->Release();
    }

    return info;
}

void AudioMonitor::listCurrentDevices() {
    if (!m_enumerator) return;

    LOG_SECTION(L"AUDIO DEVICES");

    // List output devices (speakers, headphones)
    IMMDeviceCollection* renderDevices = nullptr;
    HRESULT hr = m_enumerator->EnumAudioEndpoints(eRender, DEVICE_STATEMASK_ALL, &renderDevices);
    if (SUCCEEDED(hr) && renderDevices) {
        UINT count = 0;
        renderDevices->GetCount(&count);

        // Get default device
        IMMDevice* defaultDevice = nullptr;
        std::wstring defaultId;
        if (SUCCEEDED(m_enumerator->GetDefaultAudioEndpoint(eRender, eConsole, &defaultDevice)) && defaultDevice) {
            LPWSTR id = nullptr;
            if (SUCCEEDED(defaultDevice->GetId(&id))) {
                defaultId = id;
                CoTaskMemFree(id);
            }
            defaultDevice->Release();
        }

        for (UINT i = 0; i < count; i++) {
            IMMDevice* device = nullptr;
            if (SUCCEEDED(renderDevices->Item(i, &device)) && device) {
                AudioDeviceInfo info;
                info.isOutput = true;
                info.name = getDeviceFriendlyName(device);

                DWORD state = 0;
                device->GetState(&state);
                info.isActive = (state == DEVICE_STATE_ACTIVE);

                LPWSTR id = nullptr;
                if (SUCCEEDED(device->GetId(&id))) {
                    info.id = id;
                    info.isDefault = (info.id == defaultId);
                    CoTaskMemFree(id);
                }

                LOG_AUDIO(info.getDisplayString());
                device->Release();
            }
        }
        renderDevices->Release();
    }

    // List input devices (microphones)
    IMMDeviceCollection* captureDevices = nullptr;
    hr = m_enumerator->EnumAudioEndpoints(eCapture, DEVICE_STATEMASK_ALL, &captureDevices);
    if (SUCCEEDED(hr) && captureDevices) {
        UINT count = 0;
        captureDevices->GetCount(&count);

        // Get default capture device
        IMMDevice* defaultDevice = nullptr;
        std::wstring defaultId;
        if (SUCCEEDED(m_enumerator->GetDefaultAudioEndpoint(eCapture, eConsole, &defaultDevice)) && defaultDevice) {
            LPWSTR id = nullptr;
            if (SUCCEEDED(defaultDevice->GetId(&id))) {
                defaultId = id;
                CoTaskMemFree(id);
            }
            defaultDevice->Release();
        }

        for (UINT i = 0; i < count; i++) {
            IMMDevice* device = nullptr;
            if (SUCCEEDED(captureDevices->Item(i, &device)) && device) {
                AudioDeviceInfo info;
                info.isOutput = false;
                info.name = getDeviceFriendlyName(device);

                DWORD state = 0;
                device->GetState(&state);
                info.isActive = (state == DEVICE_STATE_ACTIVE);

                LPWSTR id = nullptr;
                if (SUCCEEDED(device->GetId(&id))) {
                    info.id = id;
                    info.isDefault = (info.id == defaultId);
                    CoTaskMemFree(id);
                }

                LOG_AUDIO(info.getDisplayString());
                device->Release();
            }
        }
        captureDevices->Release();
    }
}

// IUnknown implementation
ULONG STDMETHODCALLTYPE AudioMonitor::AddRef() {
    return ++m_refCount;
}

ULONG STDMETHODCALLTYPE AudioMonitor::Release() {
    LONG ref = --m_refCount;
    if (ref == 0) {
        delete this;
    }
    return ref;
}

HRESULT STDMETHODCALLTYPE AudioMonitor::QueryInterface(REFIID riid, void** ppvObject) {
    if (riid == IID_IUnknown || riid == __uuidof(IMMNotificationClient)) {
        *ppvObject = static_cast<IMMNotificationClient*>(this);
        AddRef();
        return S_OK;
    }
    *ppvObject = nullptr;
    return E_NOINTERFACE;
}

// IMMNotificationClient implementation
HRESULT STDMETHODCALLTYPE AudioMonitor::OnDeviceStateChanged(LPCWSTR pwstrDeviceId, DWORD dwNewState) {
    if (!m_enumerator) return S_OK;

    IMMDevice* device = nullptr;
    HRESULT hr = m_enumerator->GetDevice(pwstrDeviceId, &device);
    if (SUCCEEDED(hr) && device) {
        std::wstring name = getDeviceFriendlyName(device);
        std::wstring state = deviceStateToString(dwNewState);

        std::wstring message = name + L" -> " + state;
        LOG_AUDIO_CHANGE(message);

        device->Release();
    }
    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioMonitor::OnDeviceAdded(LPCWSTR pwstrDeviceId) {
    if (!m_enumerator) return S_OK;

    IMMDevice* device = nullptr;
    HRESULT hr = m_enumerator->GetDevice(pwstrDeviceId, &device);
    if (SUCCEEDED(hr) && device) {
        std::wstring name = getDeviceFriendlyName(device);
        LOG_AUDIO_CHANGE(L"Device Added: " + name);
        device->Release();
    }
    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioMonitor::OnDeviceRemoved(LPCWSTR pwstrDeviceId) {
    // Note: Device is already removed, so we can't get its name
    LOG_AUDIO_CHANGE(L"Device Removed: " + std::wstring(pwstrDeviceId));
    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioMonitor::OnDefaultDeviceChanged(EDataFlow flow, ERole role, LPCWSTR pwstrDefaultDeviceId) {
    if (role != eConsole) return S_OK;  // Only care about console role
    if (!m_enumerator || !pwstrDefaultDeviceId) return S_OK;

    IMMDevice* device = nullptr;
    HRESULT hr = m_enumerator->GetDevice(pwstrDefaultDeviceId, &device);
    if (SUCCEEDED(hr) && device) {
        std::wstring name = getDeviceFriendlyName(device);
        std::wstring flowStr = dataFlowToString(flow);

        LOG_AUDIO_CHANGE(L"Default " + flowStr + L" changed to: " + name);

        device->Release();
    }
    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioMonitor::OnPropertyValueChanged(LPCWSTR /*pwstrDeviceId*/, const PROPERTYKEY /*key*/) {
    // We don't log property changes to reduce noise
    return S_OK;
}