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

#include <vector>
#include <algorithm>

#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(AudioData& audioData, 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(audioData);
	}
}

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

	HRESULT err;
	IMMDeviceCollection* deviceCollection = NULL;

	err = audioData.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 = audioData.deviceEnumerator->GetDefaultAudioEndpoint(deviceType, ERole::eConsole, &defaultConsoleDevice);
	if (!FAILED(err))
	{
		defaultConsoleDevice->GetId(&defaultConsoleId);
	}

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

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

	for (UINT i = 0; i < deviceCount; i += 1)
	{
		AudioDevice deviceData{};

		err = deviceCollection->Item(i, &deviceData.device);
		if (isError(err, std::stringstream("Failed to get device ") << i << ": "))
		{
			continue;
		}

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

		IPropertyStore* propertyStore;
		err = deviceData.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 << ": ");
		deviceData.name = utf8Encode(deviceName);

		err = deviceData.device->GetState(&deviceData.state);
		isError(err, std::stringstream("Failed to reat state of device ") << i << ": ");

		err = deviceData.device->Activate(__uuidof(IAudioEndpointVolume), CLSCTX_INPROC_SERVER, NULL, (LPVOID*)&deviceData.volumeInterface);
		isError(err, "Failed to get audio endpoint volume interface: ");

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

		if (defaultConsoleId)
		{
			deviceData.isDefaultConsole = wcscmp(defaultConsoleId, deviceId) == 0;
		}
		if (defaultMediaId)
		{
			deviceData.isDefaultMedia = wcscmp(defaultMediaId, deviceId) == 0;
		}
		if (defaultCommunicationId)
		{
			deviceData.isDefaultCommunication = wcscmp(defaultCommunicationId, deviceId) == 0;
		}

		deviceList.push_back(std::move(deviceData));

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

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

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

void reloadDeviceLists(AudioData& audioData)
{
	loadAudioDevices(audioData, audioData.playbackDevices, EDataFlow::eRender);
	loadAudioDevices(audioData, audioData.recordingDevices, EDataFlow::eCapture);
}

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

	return volume;
}

void setVolume(IAudioEndpointVolume* volumeInterface, float newVolume)
{
	HRESULT hr = volumeInterface->SetMasterVolumeLevelScalar(newVolume, NULL);
	isError(hr, "Failed to set volume level: ");
}

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

	return volume;
}