#include <windows.h>
#include <functiondiscoverykeys.h>
#include <endpointvolume.h>

#include <vector>
#include <algorithm>

#include "ImguiBase.h"
#include "Util.h"
#include "AudioApi.h"
#include "PolicyConfig.h"

HRESULT getDeviceProperty(IPropertyStore* propertyStore, const PROPERTYKEY propertyKey, PROPVARIANT* outData)
{
	PropVariantInit(outData);
	HRESULT nameResult = propertyStore->GetValue(propertyKey, outData);
	return nameResult;
}

HRESULT getDevicePropertyString(IPropertyStore* propertyStore, const PROPERTYKEY propertyKey, PROPVARIANT* outData, const wchar_t*& outString, const wchar_t* defaultStr)
{
	HRESULT result = getDeviceProperty(propertyStore, propertyKey, outData);
	outString = outData->vt != VT_LPWSTR ? defaultStr : outData->pwszVal;
	return result;
}

void setDefaultAudioDevice(ApplicationData* appData, const wchar_t* deviceId, ERole role)
{
	IPolicyConfigVista* pPolicyConfig;

	HRESULT hr = CoCreateInstance(__uuidof(CPolicyConfigVistaClient), NULL, CLSCTX_ALL, __uuidof(IPolicyConfigVista), (LPVOID*)&pPolicyConfig);
	if (!isError(hr, "Failed to set default audio device: "))
	{
		hr = pPolicyConfig->SetDefaultEndpoint(deviceId, role);
		pPolicyConfig->Release();
		reloadDeviceLists(appData);
	}
}

void loadAudioDevices(std::vector<AudioDevice>& deviceList, EDataFlow deviceType)
{
	deviceList.clear();

	HRESULT err;
	IMMDeviceEnumerator* deviceEnumerator = NULL;
	IMMDeviceCollection* deviceCollection = NULL;

	err = CoCreateInstance(__uuidof(MMDeviceEnumerator), NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&deviceEnumerator));
	if (isError(err, "Failed to set up audio device enumerator: ")) return;

	err = deviceEnumerator->EnumAudioEndpoints(deviceType, DEVICE_STATE_ACTIVE | DEVICE_STATE_DISABLED, &deviceCollection);
	if (isError(err, "Failed to enumerate audio devices: ")) return;

	UINT deviceCount;
	err = deviceCollection->GetCount(&deviceCount);
	if (isError(err, "Failed to count audio devices: ")) return;

	IMMDevice* defaultConsoleDevice = NULL;
	LPWSTR defaultConsoleId = nullptr;
	err = deviceEnumerator->GetDefaultAudioEndpoint(deviceType, ERole::eConsole, &defaultConsoleDevice);
	if (!FAILED(err))
	{
		defaultConsoleDevice->GetId(&defaultConsoleId);
	}

	IMMDevice* defaultMediaOutput = NULL;
	LPWSTR defaultMediaId = nullptr;
	err = deviceEnumerator->GetDefaultAudioEndpoint(deviceType, ERole::eMultimedia, &defaultMediaOutput);
	if (!FAILED(err))
	{
		defaultMediaOutput->GetId(&defaultMediaId);
	}

	IMMDevice* defaultCommunicationOutput = NULL;
	LPWSTR defaultCommunicationId = nullptr;
	err = deviceEnumerator->GetDefaultAudioEndpoint(deviceType, ERole::eCommunications, &defaultCommunicationOutput);
	if (!FAILED(err))
	{
		defaultCommunicationOutput->GetId(&defaultCommunicationId);
	}

	for (UINT i = 0; i < deviceCount; i += 1)
	{
		IMMDevice* device;
		err = deviceCollection->Item(i, &device);
		isError(err, std::stringstream("Failed to get device ") << i << ": ");

		LPWSTR deviceId;
		err = device->GetId(&deviceId);
		isError(err, std::stringstream("Failed to get device id ") << i << ": ");

		IPropertyStore* propertyStore;
		err = device->OpenPropertyStore(STGM_READ, &propertyStore);
		isError(err, std::stringstream("Failed to open device ") << i << "prop store: ");

		PROPVARIANT deviceNameProp;
		const wchar_t* deviceName;
		err = getDevicePropertyString(propertyStore, PKEY_Device_FriendlyName, &deviceNameProp, deviceName);
		isError(err, std::stringstream("Failed to read name of device ") << i << ": ");

		DWORD deviceState;
		err = device->GetState(&deviceState);
		isError(err, std::stringstream("Failed to reat state of device ") << i << ": ");
		
		IAudioEndpointVolume* volumeInterface;
		err = device->Activate(__uuidof(IAudioEndpointVolume), CLSCTX_INPROC_SERVER, NULL, (LPVOID*)&volumeInterface);
		isError(err, "Failed to get audio endpoint volume interface: ");

		IAudioMeterInformation* meterInterface;
		err = device->Activate(__uuidof(IAudioMeterInformation), CLSCTX_INPROC_SERVER, NULL, (LPVOID*)&meterInterface);
		isError(err, "Failed to get audio meter interface: ");

		deviceList.push_back({
			device,
			volumeInterface,
			meterInterface,
			std::wstring(deviceId),
			utf8Encode(deviceName),
			deviceState,
			utf8Encode(defaultConsoleId) == utf8Encode(deviceId),
			utf8Encode(defaultMediaId) == utf8Encode(deviceId),
			utf8Encode(defaultCommunicationId) == utf8Encode(deviceId),
		});

		// Free stuff
		if (propertyStore)
		{
			propertyStore->Release();
		}
		CoTaskMemFree(deviceId);
	}

	if (deviceEnumerator)
	{
		deviceEnumerator->Release();
	}
	if (deviceCollection)
	{
		deviceCollection->Release();
	}

	std::sort(deviceList.begin(), deviceList.end(), [](AudioDevice& a, AudioDevice& b) { return a.state < b.state; });
}

void reloadDeviceLists(ApplicationData* appData)
{
	loadAudioDevices(appData->playbackDevices, EDataFlow::eRender);
	loadAudioDevices(appData->recordingDevices, EDataFlow::eCapture);
}

float getVolume(IAudioEndpointVolume* volumeInterface)
{
	float volume;
	if (FAILED(volumeInterface->GetChannelVolumeLevel(0, &volume)))
	{
		volume = 0.;
	}

	return volume;
}

float getMeterValue(IAudioMeterInformation* meterInterface)
{
	float volume;
	if (FAILED(meterInterface->GetPeakValue(&volume)))
	{
		volume = 0.;
	}

	return volume;
}