#pragma once

#include <windows.h>
#include <gl/gl.h>

#include <atlbase.h>
#include <atomic>
#include <map>
#include <memory>
#include <vector>

#include "DeckLinkAPI_h.h"
#include "VideoFrameTransfer.h"

extern "C" const IID IID_PinnedMemoryAllocator;

class OpenGLComposite;

////////////////////////////////////////////
// PinnedMemoryAllocator
////////////////////////////////////////////
class PinnedMemoryAllocator : public IDeckLinkVideoBufferAllocator
{
public:
	PinnedMemoryAllocator(HDC hdc, HGLRC hglrc, VideoFrameTransfer::Direction direction, unsigned cacheSize, unsigned bufferSize);
	virtual ~PinnedMemoryAllocator();

	bool transferFrame(void* address, GLuint gpuTexture);
	void waitForTransferComplete(void* address);
	unsigned bufferSize() { return mBufferSize; }

	// IUnknown methods
	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, LPVOID* ppv) override;
	virtual ULONG STDMETHODCALLTYPE AddRef(void) override;
	virtual ULONG STDMETHODCALLTYPE Release(void) override;

	// IDeckLinkVideoBufferAllocator methods
	virtual HRESULT STDMETHODCALLTYPE AllocateVideoBuffer(IDeckLinkVideoBuffer** allocatedBuffer) override;

private:
	void unPinAddress(void* address);

private:
	HDC mHGLDC;
	HGLRC mHGLRC;
	std::atomic<ULONG> mRefCount;
	VideoFrameTransfer::Direction mDirection;
	std::map<void*, VideoFrameTransfer*> mFrameTransfer;
	unsigned mBufferSize;
	std::vector<void*> mFrameCache;
	unsigned mFrameCacheSize;
};

////////////////////////////////////////////
// InputAllocatorPool
////////////////////////////////////////////
class InputAllocatorPool : public IDeckLinkVideoBufferAllocatorProvider
{
public:
	InputAllocatorPool(HDC hdc, HGLRC hglrc);

	// IUnknown interface
	ULONG STDMETHODCALLTYPE AddRef() override;
	ULONG STDMETHODCALLTYPE Release() override;
	HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, void** ppv) override;

	// IDeckLinkVideoBufferAllocatorProvider interface
	HRESULT STDMETHODCALLTYPE GetVideoBufferAllocator(
		/* [in] */ unsigned int bufferSize,
		/* [in] */ unsigned int width,
		/* [in] */ unsigned int height,
		/* [in] */ unsigned int rowBytes,
		/* [in] */ BMDPixelFormat pixelFormat,
		/* [out] */ IDeckLinkVideoBufferAllocator** allocator) override;

private:
	std::atomic<ULONG> mRefCount;
	std::map<unsigned int, CComPtr<PinnedMemoryAllocator> > mAllocatorBySize;
	HDC mHDC;
	HGLRC mHGLRC;
};

////////////////////////////////////////////
// DeckLinkVideoBuffer
////////////////////////////////////////////
class DeckLinkVideoBuffer : public IDeckLinkVideoBuffer
{
public:
	explicit DeckLinkVideoBuffer(std::shared_ptr<void>& buffer, PinnedMemoryAllocator* parent);
	virtual ~DeckLinkVideoBuffer() = default;

	// IUnknown interface
	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppvObject) override;
	virtual ULONG STDMETHODCALLTYPE AddRef(void) override;
	virtual ULONG STDMETHODCALLTYPE Release(void) override;

	// IDeckLinkVideoBuffer interface
	virtual HRESULT STDMETHODCALLTYPE GetBytes(void** buffer) override;
	virtual HRESULT STDMETHODCALLTYPE GetSize(uint64_t* size) override;
	virtual HRESULT STDMETHODCALLTYPE StartAccess(BMDBufferAccessFlags flags) override;
	virtual HRESULT STDMETHODCALLTYPE EndAccess(BMDBufferAccessFlags flags) override;

private:
	CComPtr<PinnedMemoryAllocator> mParentAllocator;
	std::atomic<ULONG> mRefCount;
	std::shared_ptr<void> mBuffer;
};

////////////////////////////////////////////
// Capture Delegate Class
////////////////////////////////////////////
class CaptureDelegate : public IDeckLinkInputCallback
{
	OpenGLComposite* m_pOwner;
	LONG mRefCount;

public:
	CaptureDelegate(OpenGLComposite* pOwner);

	// IUnknown needs only a dummy implementation
	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, LPVOID* ppv);
	virtual ULONG STDMETHODCALLTYPE AddRef();
	virtual ULONG STDMETHODCALLTYPE Release();

	virtual HRESULT STDMETHODCALLTYPE VideoInputFrameArrived(IDeckLinkVideoInputFrame* videoFrame, IDeckLinkAudioInputPacket* audioPacket);
	virtual HRESULT STDMETHODCALLTYPE VideoInputFormatChanged(BMDVideoInputFormatChangedEvents notificationEvents, IDeckLinkDisplayMode* newDisplayMode, BMDDetectedVideoInputFormatFlags detectedSignalFlags);
};

////////////////////////////////////////////
// Render Delegate Class
////////////////////////////////////////////
class PlayoutDelegate : public IDeckLinkVideoOutputCallback
{
	OpenGLComposite* m_pOwner;
	LONG mRefCount;

public:
	PlayoutDelegate(OpenGLComposite* pOwner);

	// IUnknown needs only a dummy implementation
	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, LPVOID* ppv);
	virtual ULONG STDMETHODCALLTYPE AddRef();
	virtual ULONG STDMETHODCALLTYPE Release();

	virtual HRESULT STDMETHODCALLTYPE ScheduledFrameCompleted(IDeckLinkVideoFrame* completedFrame, BMDOutputFrameCompletionResult result);
	virtual HRESULT STDMETHODCALLTYPE ScheduledPlaybackHasStopped();
};