#include "DeckLinkFrameTransfer.h"

#include "OpenGLComposite.h"

#include <initguid.h>

DEFINE_GUID(IID_PinnedMemoryAllocator,
	0xddf921a6, 0x279d, 0x4dcd, 0x86, 0x26, 0x75, 0x7f, 0x58, 0xa8, 0xc4, 0x35);

////////////////////////////////////////////
// PinnedMemoryAllocator
////////////////////////////////////////////

// PinnedMemoryAllocator implements the IDeckLinkVideoBufferAllocator interface to be used instead of the
// built-in buffer allocator.
//
// For this sample application a custom buffer allocator is used to ensure each address
// of buffer memory is aligned on a 4kB boundary required by the OpenGL pinned memory extension.
// If the pinned memory extension is not available, this allocator will still be used and
// demonstrates how to cache buffer allocations for efficiency.
//
// The frame cache delays the releasing of buffers until the cache fills up, thereby avoiding an
// allocate plus pin operation for every frame, followed by an unpin and deallocate on every frame.
PinnedMemoryAllocator::PinnedMemoryAllocator(HDC hdc, HGLRC hglrc, VideoFrameTransfer::Direction direction, unsigned cacheSize, unsigned bufferSize) :
	mHGLDC(hdc),
	mHGLRC(hglrc),
	mRefCount(1),
	mDirection(direction),
	mBufferSize(bufferSize),
	mFrameCacheSize(cacheSize)
{
}

PinnedMemoryAllocator::~PinnedMemoryAllocator()
{
	// Cleanup any unused buffers that remain in the cache
	while (!mFrameCache.empty())
	{
		unPinAddress(mFrameCache.back());
		VirtualFree(mFrameCache.back(), 0, MEM_RELEASE);
		mFrameCache.pop_back();
	}

	for (auto iter = mFrameTransfer.begin(); iter != mFrameTransfer.end(); ++iter)
		delete iter->second;
	mFrameTransfer.clear();
}

bool PinnedMemoryAllocator::transferFrame(void* address, GLuint gpuTexture)
{
	if (mFrameTransfer.count(address) == 0)
	{
		// VideoFrameTransfer prepares and pins address
		mFrameTransfer[address] = new VideoFrameTransfer(mBufferSize, address, mDirection);
	}

	return mFrameTransfer[address]->performFrameTransfer();
}

void PinnedMemoryAllocator::waitForTransferComplete(void* address)
{
	if (mFrameTransfer.count(address))
		mFrameTransfer[address]->waitForTransferComplete();
}

void PinnedMemoryAllocator::unPinAddress(void* address)
{
	// un-pin address only if it has been pinned for transfer
	if (mFrameTransfer.count(address) > 0)
	{
		wglMakeCurrent(mHGLDC, mHGLRC);
		mFrameTransfer.erase(address);
		wglMakeCurrent(NULL, NULL);
	}
}

HRESULT STDMETHODCALLTYPE PinnedMemoryAllocator::QueryInterface(REFIID iid, LPVOID* ppv)
{
	if (!ppv)
		return E_POINTER;

	if (iid == IID_IUnknown || iid == IID_PinnedMemoryAllocator)
	{
		*ppv = this;
	}
	else if (iid == IID_IDeckLinkVideoBufferAllocator)
	{
		*ppv = static_cast<IDeckLinkVideoBufferAllocator*>(this);
	}
	else
	{
		*ppv = nullptr;
		return E_NOINTERFACE;
	}
	AddRef();
	return S_OK;
}

ULONG STDMETHODCALLTYPE PinnedMemoryAllocator::AddRef(void)
{
	return ++mRefCount;
}

ULONG STDMETHODCALLTYPE PinnedMemoryAllocator::Release(void)
{
	int newCount = --mRefCount;
	if (newCount == 0)
		delete this;
	return newCount;
}

HRESULT STDMETHODCALLTYPE PinnedMemoryAllocator::AllocateVideoBuffer(IDeckLinkVideoBuffer** allocatedBuffer)
{
	std::shared_ptr<void> sharedMemBuffer;

	// Manage caching of allocated buffers via shared_ptr deleter.
	auto deleter = [this](void* buffer) mutable {
		if (mFrameCache.size() < mFrameCacheSize)
		{
			mFrameCache.push_back(buffer);
		}
		else
		{
			// No room left in cache, so un-pin (if it was pinned) and free this buffer
			unPinAddress(buffer);
			VirtualFree(buffer, 0, MEM_RELEASE);
		}
		// We AddRef this class once the deleter is used because this class owns the mem
		Release();
	};

	if (mFrameCache.empty())
	{
		// Allocate memory on a page boundary
		void* memBuffer = VirtualAlloc(NULL, mBufferSize, MEM_COMMIT | MEM_RESERVE | MEM_WRITE_WATCH, PAGE_READWRITE);
		if (!memBuffer)
			return E_OUTOFMEMORY;

		sharedMemBuffer = std::shared_ptr<void>(memBuffer, deleter);
	}
	else
	{
		// Re-use most recently released address
		sharedMemBuffer = std::shared_ptr<void>(mFrameCache.back(), deleter);
		mFrameCache.pop_back();
	}

	// This class owns the mem so the buffer we return needs to AddRef() this, and Release() in the deleter
	AddRef();

	*allocatedBuffer = new DeckLinkVideoBuffer(sharedMemBuffer, this);
	return S_OK;
}

////////////////////////////////////////////
// InputAllocatorPool Class
////////////////////////////////////////////
InputAllocatorPool::InputAllocatorPool(HDC hdc, HGLRC hglrc)
{
	mHDC = hdc;
	mHGLRC = hglrc;
}

HRESULT InputAllocatorPool::QueryInterface(REFIID iid, void** ppv)
{
	if (!ppv)
		return E_POINTER;

	if (iid == IID_IUnknown)
	{
		*ppv = this;
	}
	else if (iid == IID_IDeckLinkVideoBufferAllocatorProvider)
	{
		*ppv = static_cast<IDeckLinkVideoBufferAllocatorProvider*>(this);
	}
	else
	{
		*ppv = nullptr;
		return E_NOINTERFACE;
	}
	AddRef();
	return S_OK;
}

ULONG InputAllocatorPool::AddRef(void)
{
	return ++mRefCount;
}

ULONG InputAllocatorPool::Release(void)
{
	int newCount = --mRefCount;
	if (newCount == 0)
		delete this;
	return newCount;
}

HRESULT InputAllocatorPool::GetVideoBufferAllocator(
	/* [in] */ unsigned int bufferSize,
	/* [in] */ unsigned int,
	/* [in] */ unsigned int,
	/* [in] */ unsigned int,
	/* [in] */ BMDPixelFormat,
	/* [out] */ IDeckLinkVideoBufferAllocator** allocator)
{
	if (!allocator)
		return E_POINTER;

	auto existing = mAllocatorBySize.find(bufferSize);
	if (existing != mAllocatorBySize.end())
	{
		*allocator = &*existing->second;
		(*allocator)->AddRef();
		return S_OK;
	}

	CComPtr<PinnedMemoryAllocator> newAllocator;
	newAllocator.Attach(new (std::nothrow) PinnedMemoryAllocator(mHDC, mHGLRC, VideoFrameTransfer::CPUtoGPU, 3, bufferSize));
	if (!newAllocator)
		return E_OUTOFMEMORY;

	mAllocatorBySize.emplace(std::make_pair(bufferSize, newAllocator));
	*allocator = newAllocator.Detach();
	return S_OK;
}

////////////////////////////////////////////
// DeckLink Video Buffer Class
////////////////////////////////////////////
DeckLinkVideoBuffer::DeckLinkVideoBuffer(std::shared_ptr<void>& buffer, PinnedMemoryAllocator* parent) :
	mParentAllocator(parent),
	mRefCount(1),
	mBuffer(buffer)
{
}

HRESULT STDMETHODCALLTYPE DeckLinkVideoBuffer::QueryInterface(REFIID riid, void** ppvObject)
{
	HRESULT result = S_OK;

	if (ppvObject == nullptr)
		return E_POINTER;

	if (riid == IID_IUnknown)
	{
		*ppvObject = this;
		AddRef();
	}
	else if (riid == IID_IDeckLinkVideoBuffer)
	{
		*ppvObject = static_cast<IDeckLinkVideoBuffer*>(this);
		AddRef();
	}
	else if (riid == IID_PinnedMemoryAllocator)
	{
		result = mParentAllocator->QueryInterface(riid, ppvObject);
	}
	else
	{
		*ppvObject = nullptr;
		result = E_NOINTERFACE;
	}

	return result;
}

ULONG STDMETHODCALLTYPE DeckLinkVideoBuffer::AddRef()
{
	return ++mRefCount;
}

ULONG STDMETHODCALLTYPE DeckLinkVideoBuffer::Release()
{
	int newValue = --mRefCount;
	if (newValue == 0)
		delete this;

	return newValue;
}

HRESULT STDMETHODCALLTYPE DeckLinkVideoBuffer::GetBytes(void** buffer)
{
	if (buffer == nullptr)
		return E_POINTER;

	*buffer = mBuffer.get();
	return S_OK;
}

HRESULT STDMETHODCALLTYPE DeckLinkVideoBuffer::GetSize(uint64_t* size)
{
	if (size == nullptr)
		return E_POINTER;

	*size = mParentAllocator->bufferSize();
	return S_OK;
}

HRESULT STDMETHODCALLTYPE DeckLinkVideoBuffer::StartAccess(BMDBufferAccessFlags)
{
	return S_OK;
}

HRESULT STDMETHODCALLTYPE DeckLinkVideoBuffer::EndAccess(BMDBufferAccessFlags)
{
	return S_OK;
}

////////////////////////////////////////////
// DeckLink Capture Delegate Class
////////////////////////////////////////////
CaptureDelegate::CaptureDelegate(OpenGLComposite* pOwner) :
	m_pOwner(pOwner),
	mRefCount(1)
{
}

HRESULT CaptureDelegate::QueryInterface(REFIID, LPVOID* ppv)
{
	*ppv = NULL;
	return E_NOINTERFACE;
}

ULONG CaptureDelegate::AddRef()
{
	return InterlockedIncrement(&mRefCount);
}

ULONG CaptureDelegate::Release()
{
	int newCount = InterlockedDecrement(&mRefCount);
	if (newCount == 0)
		delete this;
	return newCount;
}

HRESULT CaptureDelegate::VideoInputFrameArrived(IDeckLinkVideoInputFrame* inputFrame, IDeckLinkAudioInputPacket*)
{
	if (!inputFrame)
	{
		// It's possible to receive a NULL inputFrame, but a valid audioPacket. Ignore audio-only frame.
		return S_OK;
	}

	bool hasNoInputSource = (inputFrame->GetFlags() & bmdFrameHasNoInputSource) == bmdFrameHasNoInputSource;
	m_pOwner->VideoFrameArrived(inputFrame, hasNoInputSource);
	return S_OK;
}

HRESULT CaptureDelegate::VideoInputFormatChanged(BMDVideoInputFormatChangedEvents, IDeckLinkDisplayMode*, BMDDetectedVideoInputFormatFlags)
{
	return S_OK;
}

////////////////////////////////////////////
// DeckLink Playout Delegate Class
////////////////////////////////////////////
PlayoutDelegate::PlayoutDelegate(OpenGLComposite* pOwner) :
	m_pOwner(pOwner),
	mRefCount(1)
{
}

HRESULT PlayoutDelegate::QueryInterface(REFIID, LPVOID* ppv)
{
	*ppv = NULL;
	return E_NOINTERFACE;
}

ULONG PlayoutDelegate::AddRef()
{
	return InterlockedIncrement(&mRefCount);
}

ULONG PlayoutDelegate::Release()
{
	int newCount = InterlockedDecrement(&mRefCount);
	if (newCount == 0)
		delete this;
	return newCount;
}

HRESULT PlayoutDelegate::ScheduledFrameCompleted(IDeckLinkVideoFrame* completedFrame, BMDOutputFrameCompletionResult result)
{
	switch (result)
	{
	case bmdOutputFrameDisplayedLate:
		OutputDebugStringA("ScheduledFrameCompleted() frame did not complete: Frame Displayed Late\n");
		break;
	case bmdOutputFrameDropped:
		OutputDebugStringA("ScheduledFrameCompleted() frame did not complete: Frame Dropped\n");
		break;
	case bmdOutputFrameCompleted:
	case bmdOutputFrameFlushed:
		// Don't log bmdOutputFrameFlushed result since it is expected when Stop() is called
		break;
	default:
		OutputDebugStringA("ScheduledFrameCompleted() frame did not complete: Unknown error\n");
	}

	m_pOwner->PlayoutFrameCompleted(completedFrame, result);
	return S_OK;
}

HRESULT PlayoutDelegate::ScheduledPlaybackHasStopped()
{
	return S_OK;
}