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

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

#include <initguid.h>
#include <mmdeviceapi.h>

#include <array>
#include <vector>
#include <algorithm>
#include <iostream>

void initAudio(ApplicationData& appData)
{
	HRESULT 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;
			}

			if (audioDevice.isInitialized)
			{
				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);
	}
}

// @slow (.2ms)
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: ");
}

// @slow (.4ms)
UINT getMeterValues(IAudioMeterInformation* meterInterface, std::array<float, 2>& outLevels)
{
	UINT channelCount;
	if (FAILED(meterInterface->GetMeteringChannelCount(&channelCount)) || channelCount > 2) return 0;

	if (FAILED(meterInterface->GetChannelsPeakValues(channelCount, &outLevels[0]))) return 0;

	return channelCount;
}

void getVolumeLimit(IAudioEndpointVolume* volumeInterface, float* outMin, float* outMax)
{
	float dummy;
	volumeInterface->GetVolumeRange(outMin, outMax, &dummy);
}

// @slow (.3ms)
bool isMuted(IAudioEndpointVolume* volumeInterface)
{
	BOOL result = false;
	if (FAILED(volumeInterface->GetMute(&result))) return false;

	return static_cast<bool>(result);
}

void setMuted(IAudioEndpointVolume* volumeInterface, bool newState)
{
	volumeInterface->SetMute(static_cast<BOOL>(newState), NULL);
}

void CreateLoudnessEqualizationKey(PROPERTYKEY& key)
{
	// Realtek: const wchar_t* guid = L"E0A941A0-88A2-4df5-8D6B-DD20BB06E8FB";
	const wchar_t* guid = L"FC52A749-4BE9-4510-896E-966BA6525980";
	UuidFromString((RPC_WSTR)guid, &key.fmtid);
	key.pid = 3;
}

AudioProcessingState getAudioProcessing(AudioDevice& device)
{
	return AudioProcessingState::Unknown;

	LPWSTR pwstrEndpointId = NULL;
	HRESULT hr = device.device->GetId(&pwstrEndpointId);

	IPolicyConfig* pPolicyConfig;
	hr = CoCreateInstance(__uuidof(CPolicyConfigClient), NULL, CLSCTX_ALL, IID_PPV_ARGS(&pPolicyConfig));
	if (SUCCEEDED(hr))
	{
		PROPVARIANT var;
		PropVariantInit(&var);

		PROPERTYKEY key{};
		CreateLoudnessEqualizationKey(key);

		hr = pPolicyConfig->GetPropertyValue(pwstrEndpointId, TRUE, key, &var);
		pPolicyConfig->Release();
		return var.boolVal ? AudioProcessingState::Disabled : AudioProcessingState::Enabled;
	}
	return AudioProcessingState::Unknown;
}

void setAudioProcessing(AudioDevice& device, bool newVal)
{
	return;

	LPWSTR pwstrEndpointId = NULL;
	HRESULT hr = device.device->GetId(&pwstrEndpointId);

	IPolicyConfig* pPolicyConfig;
	hr = CoCreateInstance(__uuidof(CPolicyConfigClient), NULL, CLSCTX_ALL, IID_PPV_ARGS(&pPolicyConfig));
	if (SUCCEEDED(hr))
	{
		PROPVARIANT var;
		PropVariantInit(&var);
		var.vt = VT_UI4;
		var.uintVal = newVal ? ENDPOINT_SYSFX_ENABLED : ENDPOINT_SYSFX_DISABLED;
		PROPERTYKEY key{};
		CreateLoudnessEqualizationKey(key);
		hr = pPolicyConfig->SetPropertyValue(pwstrEndpointId, TRUE, key, &var);
		pPolicyConfig->Release();
	}
}
