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

#include <vector>
#include <algorithm>

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

void initAudio(ApplicationData& appData)
{
	HRESULT audioResult;
	audioResult = CoInitializeEx(NULL, COINIT_MULTITHREADED);
	isError(audioResult, "Failed to initialize COM: ");

	audioResult = CoCreateInstance(__uuidof(MMDeviceEnumerator), NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&appData.audioData->deviceEnumerator));
	isError(audioResult, "Failed to set up audio device enumerator: ");

	appData.audioData->audioNotificationListener = new AudioNotificationListener(appData.audioData);
	audioResult = appData.audioData->deviceEnumerator->RegisterEndpointNotificationCallback(appData.audioData->audioNotificationListener);
	isError(audioResult, "Failed to register audio notification listener: ");

	reloadDeviceLists(*appData.audioData);
}

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

	HRESULT err;
	IMMDeviceCollection* deviceCollection = nullptr;

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

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

	IMMDevice* defaultCommunicationOutput = nullptr;
	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)
	{
		IMMDevice* device;
		err = deviceCollection->Item(i, &device);
		if (isError(err, std::stringstream("Failed to get device ") << i << ": "))
		{
			continue;
		}

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

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

			deviceList.push_back(std::move(audioDevice));
		}

		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);
}

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 = nullptr;

	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);
	}
}

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;
}