#include "wasapi_player.h"
#include <initguid.h>
#include <mmdeviceapi.h>
//#include <functiondiscoverykeys_devpkey.h>
#include <functional>
#include <audioclient.h>
#include <mmdeviceapi.h>
#include <comdef.h>
#include <comip.h>
#include <avrt.h>
#include <ksmedia.h>
#include <chrono>
#include <assert.h>
#include <codecvt>
#include <locale>
#include "common/snap_exception.hpp"
#include "common/aixlog.hpp"

using namespace std;
using namespace std::chrono;
using namespace std::chrono_literals;

static constexpr auto LOG_TAG = "WASAPI";

template<typename T>
struct COMMemDeleter
{
	void operator() (T* obj)
	{
		if (obj != NULL)
		{
			CoTaskMemFree(obj);
			obj = NULL;
		}
	}
};

template<typename T>
using com_mem_ptr = unique_ptr<T, COMMemDeleter<T> >;

using com_handle = unique_ptr<void, function<BOOL(HANDLE)> >;

const CLSID CLSID_MMDeviceEnumerator = __uuidof(MMDeviceEnumerator);
const IID IID_IMMDeviceEnumerator = __uuidof(IMMDeviceEnumerator);
const IID IID_IAudioClient = __uuidof(IAudioClient);
const IID IID_IAudioRenderClient = __uuidof(IAudioRenderClient);
const IID IID_IAudioClock = __uuidof(IAudioClock);

_COM_SMARTPTR_TYPEDEF(IMMDevice,__uuidof(IMMDevice));
_COM_SMARTPTR_TYPEDEF(IMMDeviceCollection,__uuidof(IMMDeviceCollection));
_COM_SMARTPTR_TYPEDEF(IMMDeviceEnumerator,__uuidof(IMMDeviceEnumerator));
_COM_SMARTPTR_TYPEDEF(IAudioClient,__uuidof(IAudioClient));
_COM_SMARTPTR_TYPEDEF(IPropertyStore,__uuidof(IPropertyStore));
_COM_SMARTPTR_TYPEDEF(IAudioSessionManager, __uuidof(IAudioSessionManager));
_COM_SMARTPTR_TYPEDEF(IAudioSessionControl, __uuidof(IAudioSessionControl));

#define REFTIMES_PER_SEC  10000000
#define REFTIMES_PER_MILLISEC  10000

EXTERN_C const PROPERTYKEY DECLSPEC_SELECTANY PKEY_Device_FriendlyName = { { 0xa45c254e, 0xdf1c, 0x4efd, { 0x80, 0x20, 0x67, 0xd1, 0x46, 0xa8, 0x50, 0xe0 } }, 14 };

#define CHECK_HR(hres) if(FAILED(hres)){stringstream ss;ss<<"HRESULT fault status: "<<hex<<(hres)<<" line "<<dec<<__LINE__<<endl; LOG(FATAL, LOG_TAG) << ss.str();throw SnapException(ss.str());}

WASAPIPlayer::WASAPIPlayer(const PcmDevice& pcmDevice, std::shared_ptr<Stream> stream, ClientSettings::WasapiMode mode) : Player(pcmDevice, stream), mode_(mode)
{
	HRESULT hr = CoInitializeEx(
		NULL,
		COINIT_MULTITHREADED);
	CHECK_HR(hr);

	audioEventListener_ = new AudioSessionEventListener();
}

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

inline PcmDevice convertToDevice(int idx, IMMDevicePtr& device)
{
	HRESULT hr;
	PcmDevice desc;

	LPWSTR id = NULL;
	hr = device->GetId(&id);
	CHECK_HR(hr);

	IPropertyStorePtr properties = nullptr;
	hr = device->OpenPropertyStore(STGM_READ, &properties);

	PROPVARIANT deviceName;
	PropVariantInit(&deviceName);

	hr = properties->GetValue(PKEY_Device_FriendlyName, &deviceName);
	CHECK_HR(hr);

	desc.idx = idx;
	desc.name = wstring_convert<codecvt_utf8<wchar_t>, wchar_t>().to_bytes(id);
	desc.description = wstring_convert<codecvt_utf8<wchar_t>, wchar_t>().to_bytes(deviceName.pwszVal);

	CoTaskMemFree(id);
	
	return desc;
}

vector<PcmDevice> WASAPIPlayer::pcm_list()
{
	HRESULT hr;
	IMMDeviceCollectionPtr devices = nullptr;
	IMMDeviceEnumeratorPtr deviceEnumerator = nullptr;
	
	hr = CoInitializeEx(
		NULL,
		COINIT_MULTITHREADED);
	if (hr != CO_E_ALREADYINITIALIZED)
		CHECK_HR(hr);
	
	hr = CoCreateInstance(
		CLSID_MMDeviceEnumerator, NULL,
		CLSCTX_SERVER, IID_IMMDeviceEnumerator,
		(void**)&deviceEnumerator);
	CHECK_HR(hr);

	hr = deviceEnumerator->EnumAudioEndpoints(eRender, DEVICE_STATE_ACTIVE, &devices);
	CHECK_HR(hr);

	UINT deviceCount;
	devices->GetCount(&deviceCount);

	if (deviceCount == 0)
		throw SnapException("no valid devices");
	
	vector<PcmDevice> deviceList;
	
	{
		IMMDevicePtr defaultDevice = nullptr;
		hr = deviceEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &defaultDevice);
		CHECK_HR(hr);

		auto dev = convertToDevice(0, defaultDevice);
		dev.name = "default";
		deviceList.push_back(dev);
	}
	
	for (UINT i = 0; i < deviceCount; ++i)
	{
		IMMDevicePtr device = nullptr;

		hr = devices->Item(i, &device);
		CHECK_HR(hr);
		deviceList.push_back(convertToDevice(i, device));
	}

	return deviceList;
}

void WASAPIPlayer::worker()
{
  assert(sizeof(char) == sizeof(BYTE));
	
	HRESULT hr;

	// Create the format specifier
	com_mem_ptr<WAVEFORMATEX> waveformat((WAVEFORMATEX*)(CoTaskMemAlloc(sizeof(WAVEFORMATEX))));
	waveformat->wFormatTag      = WAVE_FORMAT_PCM;
	waveformat->nChannels       = stream_->getFormat().channels();
	waveformat->nSamplesPerSec  = stream_->getFormat().rate();
	waveformat->wBitsPerSample  = stream_->getFormat().bits();

	waveformat->nBlockAlign     = waveformat->nChannels * waveformat->wBitsPerSample / 8;
	waveformat->nAvgBytesPerSec = waveformat->nSamplesPerSec * waveformat->nBlockAlign;

	waveformat->cbSize          = 0;

	com_mem_ptr<WAVEFORMATEXTENSIBLE> waveformatExtended((WAVEFORMATEXTENSIBLE*)(CoTaskMemAlloc(sizeof(WAVEFORMATEXTENSIBLE))));
	waveformatExtended->Format                      = *waveformat;
	waveformatExtended->Format.wFormatTag           = WAVE_FORMAT_EXTENSIBLE;
	waveformatExtended->Format.cbSize               = 22;
	waveformatExtended->Samples.wValidBitsPerSample = waveformat->wBitsPerSample;
	waveformatExtended->dwChannelMask               = SPEAKER_FRONT_LEFT | SPEAKER_FRONT_RIGHT;
	waveformatExtended->SubFormat                   = KSDATAFORMAT_SUBTYPE_PCM;


	// Retrieve the device enumerator
	IMMDeviceEnumeratorPtr deviceEnumerator = nullptr;
	hr = CoCreateInstance(
		CLSID_MMDeviceEnumerator, NULL, 
		CLSCTX_SERVER, IID_IMMDeviceEnumerator,
		(void**)&deviceEnumerator);
	CHECK_HR(hr);

	// Register the default playback device (eRender for playback)
	IMMDevicePtr device = nullptr;
	if (pcmDevice_.idx == 0)
	{
		hr = deviceEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &device);
		CHECK_HR(hr);
	}
	else
	{
		IMMDeviceCollectionPtr devices = nullptr;
		hr = deviceEnumerator->EnumAudioEndpoints(eRender, DEVICE_STATE_ACTIVE, &devices);
		CHECK_HR(hr);

		devices->Item(pcmDevice_.idx, &device);
	}

	IPropertyStorePtr properties = nullptr;
	hr = device->OpenPropertyStore(STGM_READ, &properties);
	CHECK_HR(hr);

	PROPVARIANT format;
	hr = properties->GetValue(PKEY_AudioEngine_DeviceFormat, &format);
	CHECK_HR(hr);

	PWAVEFORMATEX formatEx = (PWAVEFORMATEX)format.blob.pBlobData;
	LOG(INFO, LOG_TAG) << "Device accepts format: " << formatEx->nSamplesPerSec << ":" << formatEx->wBitsPerSample << ":" << formatEx->nChannels << "\n";    
	// Activate the device
	IAudioClientPtr audioClient = nullptr;
	hr = device->Activate(IID_IAudioClient, CLSCTX_SERVER, NULL, (void**)&audioClient);
	CHECK_HR(hr);

	if(mode_ == ClientSettings::WasapiMode::EXCLUSIVE)
	{
		hr = audioClient->IsFormatSupported(
			AUDCLNT_SHAREMODE_EXCLUSIVE,
			&(waveformatExtended->Format),
			NULL);
		CHECK_HR(hr);
	}

	IAudioSessionManagerPtr sessionManager = nullptr;
    // Get the session manager for the endpoint device.
    hr = device->Activate(__uuidof(IAudioSessionManager), CLSCTX_INPROC_SERVER, NULL, (void**)&sessionManager);
    CHECK_HR(hr);

    // Get the control interface for the process-specific audio
    // session with session GUID = GUID_NULL. This is the session
    // that an audio stream for a DirectSound, DirectShow, waveOut,
    // or PlaySound application stream belongs to by default.
    IAudioSessionControlPtr control = nullptr;
    hr = sessionManager->GetAudioSessionControl(NULL, 0, &control);
    CHECK_HR(hr);

    // register
    hr = control->RegisterAudioSessionNotification(audioEventListener_);
    CHECK_HR(hr);
	
	// Get the device period
	REFERENCE_TIME hnsRequestedDuration = REFTIMES_PER_SEC;
	hr = audioClient->GetDevicePeriod(NULL, &hnsRequestedDuration);
	CHECK_HR(hr);
	
	LOG(INFO, LOG_TAG) << "Initializing WASAPI in " << (mode_ == ClientSettings::WasapiMode::SHARED ? "shared" : "exclusive") << " mode\n";

	_AUDCLNT_SHAREMODE share_mode = mode_ == ClientSettings::WasapiMode::SHARED ? AUDCLNT_SHAREMODE_SHARED : AUDCLNT_SHAREMODE_EXCLUSIVE;
    DWORD stream_flags = mode_ == ClientSettings::WasapiMode::SHARED
                                 ? AUDCLNT_STREAMFLAGS_EVENTCALLBACK | AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM | AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY
                                 : AUDCLNT_STREAMFLAGS_EVENTCALLBACK;

	// Initialize the client at minimum latency
	hr = audioClient->Initialize(
		share_mode,
		stream_flags,
		hnsRequestedDuration,
		hnsRequestedDuration,
		&(waveformatExtended->Format),
		NULL);

	if (hr == AUDCLNT_E_BUFFER_SIZE_NOT_ALIGNED)
	{
		UINT32 alignedBufferSize;
		hr = audioClient->GetBufferSize(&alignedBufferSize);
		CHECK_HR(hr);
		audioClient.Attach(NULL, false);
		hnsRequestedDuration = (REFERENCE_TIME)((10000.0 * 1000 / waveformatExtended->Format.nSamplesPerSec * alignedBufferSize) + 0.5);
		hr = device->Activate(IID_IAudioClient, CLSCTX_SERVER, NULL, (void**)&audioClient);
		CHECK_HR(hr);
		hr = audioClient->Initialize(
			share_mode,
			stream_flags,
			hnsRequestedDuration,
			hnsRequestedDuration,
			&(waveformatExtended->Format),
			NULL);
	}
	CHECK_HR(hr);

	// Register an event to refill the buffer
	com_handle eventHandle(CreateEvent(NULL, FALSE, FALSE, NULL), &::CloseHandle);
	if (eventHandle == NULL)
		CHECK_HR(E_FAIL);
	hr = audioClient->SetEventHandle(HANDLE(eventHandle.get()));
	CHECK_HR(hr);

	// Get size of buffer
	UINT32 bufferFrameCount;
	hr = audioClient->GetBufferSize(&bufferFrameCount);
	CHECK_HR(hr);

	// Get the rendering service
	IAudioRenderClient* renderClient = NULL;
	hr = audioClient->GetService(
		IID_IAudioRenderClient,
		(void**)&renderClient);
	CHECK_HR(hr);

	// Grab the clock service
	IAudioClock* clock = NULL;
	hr = audioClient->GetService(
		IID_IAudioClock,
		(void**)&clock);
	CHECK_HR(hr);

	// Boost our priority
	DWORD taskIndex = 0;
	com_handle taskHandle(AvSetMmThreadCharacteristics(TEXT("Pro Audio"), &taskIndex),
		&::AvRevertMmThreadCharacteristics);
	if (taskHandle == NULL)
		CHECK_HR(E_FAIL);

	// And, action!
	hr = audioClient->Start();
	CHECK_HR(hr);
	
	size_t bufferSize = bufferFrameCount * waveformatExtended->Format.nBlockAlign;
	BYTE* buffer;
	unique_ptr<char[]> queueBuffer(new char[bufferSize]);
	UINT64 position = 0, bufferPosition = 0, frequency;
	clock->GetFrequency(&frequency);
	
	while (active_)
	{
	  DWORD returnVal = WaitForSingleObject(eventHandle.get(), 2000);
		if (returnVal != WAIT_OBJECT_0)
		{
            //stop();
            LOG(INFO, LOG_TAG) << "Got timeout waiting for audio device callback\n";
            CHECK_HR(ERROR_TIMEOUT);
			
			hr = audioClient->Stop();
            CHECK_HR(hr);
            hr = audioClient->Reset();
            CHECK_HR(hr);

            while (active_ && !stream_->waitForChunk(std::chrono::milliseconds(100)))
                LOG(INFO, LOG_TAG) << "Waiting for chunk\n";

            hr = audioClient->Start();
            CHECK_HR(hr);
            bufferPosition = 0;
            break;
		}

		// Thread was sleeping above, double check that we are still running
		if (!active_)
			break;

		// update our volume from IAudioControl
        volCorrection_ = audioEventListener_->getVolume();

		clock->GetPosition(&position, NULL);
        
		UINT32 padding = 0;
		if(mode_ == ClientSettings::WasapiMode::SHARED)
        {
			hr = audioClient->GetCurrentPadding(&padding);
			CHECK_HR(hr);
		}

		int available = bufferFrameCount - padding;

		if (stream_->getPlayerChunk(queueBuffer.get(), microseconds(
		                                                      ((bufferPosition * 1000000) / waveformat->nSamplesPerSec) -
		                                                      ((position * 1000000) / frequency)),
                                            available))
		{
			if (available > 0)
			{
				adjustVolume(queueBuffer.get(), available);
				hr = renderClient->GetBuffer(available, &buffer);
				CHECK_HR(hr);
				memcpy(buffer, queueBuffer.get(), bufferSize);
				hr = renderClient->ReleaseBuffer(available, 0);
				CHECK_HR(hr);
			
				bufferPosition += available;
			}
		}
		else
		{
			std::clog << static_cast<AixLog::Severity>(INFO) << AixLog::Tag(LOG_TAG);
			LOG(INFO, LOG_TAG) << "Failed to get chunk\n";
			
			hr = audioClient->Stop();
			CHECK_HR(hr);
			hr = audioClient->Reset();
			CHECK_HR(hr);
			
			while (active_ && !stream_->waitForChunk(std::chrono::milliseconds(100)))
				LOG(INFO, LOG_TAG) << "Waiting for chunk\n";
			
			hr = audioClient->Start();
			CHECK_HR(hr);
			bufferPosition = 0;
		}
	}
}

HRESULT STDMETHODCALLTYPE AudioSessionEventListener::QueryInterface(REFIID riid, VOID** ppvInterface)
{
    if (IID_IUnknown == riid)
    {
        AddRef();
        *ppvInterface = (IUnknown*)this;
    }
    else if (__uuidof(IAudioSessionEvents) == riid)
    {
        AddRef();
        *ppvInterface = (IAudioSessionEvents*)this;
    }
    else
    {
        *ppvInterface = NULL;
        return E_NOINTERFACE;
    }
    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioSessionEventListener::OnSimpleVolumeChanged(float NewVolume, BOOL NewMute, LPCGUID EventContext)
{
    volume_ = NewVolume;
    if (NewMute)
    {
        LOG(DEBUG, LOG_TAG) << ("MUTE\n");
    }
    else
    {
        LOG(DEBUG, LOG_TAG) << "Volume = " << (UINT32)(100 * NewVolume + 0.5) << " percent\n";
    }

    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioSessionEventListener::OnStateChanged(AudioSessionState NewState)
{
    char* pszState = "?????";

    switch (NewState)
    {
        case AudioSessionStateActive:
            pszState = "active";
            break;
        case AudioSessionStateInactive:
            pszState = "inactive";
            break;
    }
    LOG(DEBUG, LOG_TAG) << "New session state = " << pszState << "\n";

    return S_OK;
}

HRESULT STDMETHODCALLTYPE AudioSessionEventListener::OnSessionDisconnected(AudioSessionDisconnectReason DisconnectReason)
{
    char* pszReason = "?????";

    switch (DisconnectReason)
    {
        case DisconnectReasonDeviceRemoval:
            pszReason = "device removed";
            break;
        case DisconnectReasonServerShutdown:
            pszReason = "server shut down";
            break;
        case DisconnectReasonFormatChanged:
            pszReason = "format changed";
            break;
        case DisconnectReasonSessionLogoff:
            pszReason = "user logged off";
            break;
        case DisconnectReasonSessionDisconnected:
            pszReason = "session disconnected";
            break;
        case DisconnectReasonExclusiveModeOverride:
            pszReason = "exclusive-mode override";
            break;
    }
    LOG(INFO, LOG_TAG) << "Audio session disconnected (reason: " << pszReason << ")";

    return S_OK;
}