#include "WinMfImageGrabber.h"
#include <mfapi.h>
#include "../log/NLog.h"
#include "../Utils.h"
#include "WinMfFormatReader.h"
#include "WinMfMediaType.h"
#include <shlwapi.h>
#include "WebCamera.h"
#include "../Com.h"

namespace LifeCore {
	WinMfImageGrabber::WinMfImageGrabber(WebCamera* cam)
		: m_cRef(1), ig_pSession(nullptr), camera(cam)
	{

	}
	WinMfImageGrabber::~WinMfImageGrabber() {
		LogInfo() << "~WinMfImageGrabber ";
		/*if(ig_pSession) {
			LogInfo() << "~WinMfImageGrabber shutdown session";
			ig_pSession->Shutdown();
		}*/
	}
	HRESULT WinMfImageGrabber::initImageGrabber(IMFMediaSource *pSource, GUID VideoFormat) {
		IMFActivate *pSinkActivate = NULL;/**/
		IMFMediaType *pType = NULL;/**/
		IMFPresentationDescriptor *pPD = NULL;/**/
		IMFStreamDescriptor *pSD = NULL;/**/
		IMFMediaTypeHandler *pHandler = NULL;/**/
		IMFMediaType *pCurrentType = NULL;/**/
		WinMfMediaType MT;
		auto hr = pSource->CreatePresentationDescriptor(&pPD);
		if(FAILED(hr)) {
			LogError() << "Failed to create presentation descriptor!";
			goto err;
		}

		BOOL fSelected;
		hr = pPD->GetStreamDescriptorByIndex(0, &fSelected, &pSD);
		if(FAILED(hr)) {
			goto err;
		}

		hr = pSD->GetMediaTypeHandler(&pHandler);
		if(FAILED(hr)) {
			goto err;
		}

		DWORD cTypes = 0;
		hr = pHandler->GetMediaTypeCount(&cTypes);
		if(FAILED(hr)) {
			goto err;
		}

		if(cTypes > 0) {
			hr = pHandler->GetCurrentMediaType(&pCurrentType);

			if(FAILED(hr)) {
				goto err;
			}

			MT = WinMfFormatReader::Read(pCurrentType);
		}

	err:
		SafeRelease(&pPD);
		SafeRelease(&pSD);
		SafeRelease(&pHandler);
		SafeRelease(&pCurrentType);

		unsigned int sizeRawImage = 0;
		if(VideoFormat == MFVideoFormat_RGB24) {
			sizeRawImage = MT.MF_MT_FRAME_SIZE * 3;
		} else if(VideoFormat == MFVideoFormat_RGB32) {
			sizeRawImage = MT.MF_MT_FRAME_SIZE * 4;
		}

		buffer_first = new RawBuffer<uint8>(sizeRawImage, MT);
		buffer_second = new RawBuffer<uint8>(sizeRawImage, MT);
		buffer_active = buffer_second;

		hr = MFCreateMediaType(&pType);
		hr = pType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
		hr = pType->SetGUID(MF_MT_SUBTYPE, VideoFormat);

		// Create the sample grabber sink.
		hr = MFCreateSampleGrabberSinkActivate(pType, this, &pSinkActivate);

		// To run as fast as possible, set this attribute (requires Windows 7):
		hr = pSinkActivate->SetUINT32(MF_SAMPLEGRABBERSINK_IGNORE_CLOCK, TRUE);

		// Create the Media Session.

		hr = MFCreateMediaSession(NULL, &ig_pSession);

		// Create the topology.
		hr = CreateTopology(pSource, pSinkActivate, &ig_pTopology);

		SafeRelease(&pSinkActivate);
		SafeRelease(&pType);
		return S_OK;

	}

	RawBuffer<uint8>* WinMfImageGrabber::GetImage() {
		return buffer_active;
	}

#define CHECK_HR(x) if (FAILED(x)) { goto done; }
	HRESULT WinMfImageGrabber::CreateTopology(IMFMediaSource *pSource, IMFActivate *pSinkActivate, IMFTopology **ppTopo) {
		IMFTopology *pTopology = NULL;
		IMFPresentationDescriptor *pPD = NULL;
		IMFStreamDescriptor *pSD = NULL;
		IMFMediaTypeHandler *pHandler = NULL;
		IMFTopologyNode *pNode1 = NULL;
		IMFTopologyNode *pNode2 = NULL;

		HRESULT hr = S_OK;
		DWORD cStreams = 0;

		CHECK_HR(hr = MFCreateTopology(&pTopology));
		CHECK_HR(hr = pSource->CreatePresentationDescriptor(&pPD));
		CHECK_HR(hr = pPD->GetStreamDescriptorCount(&cStreams));

		for(DWORD i = 0; i < cStreams; i++) {
			// In this example, we look for audio streams and connect them to the sink.

			BOOL fSelected = FALSE;
			GUID majorType;

			CHECK_HR(hr = pPD->GetStreamDescriptorByIndex(i, &fSelected, &pSD));
			CHECK_HR(hr = pSD->GetMediaTypeHandler(&pHandler));
			CHECK_HR(hr = pHandler->GetMajorType(&majorType));

			if(majorType == MFMediaType_Video && fSelected) {
				CHECK_HR(hr = AddSourceNode(pTopology, pSource, pPD, pSD, &pNode1));
				CHECK_HR(hr = AddOutputNode(pTopology, pSinkActivate, 0, &pNode2));
				CHECK_HR(hr = pNode1->ConnectOutput(0, pNode2, 0));
				break;
			} else {
				CHECK_HR(hr = pPD->DeselectStream(i));
			}
			SafeRelease(&pSD);
			SafeRelease(&pHandler);
		}

		*ppTopo = pTopology;
		(*ppTopo)->AddRef();

	done:
		SAFE_RELEASE(pTopology);
		SAFE_RELEASE(pNode1);
		SAFE_RELEASE(pNode2);
		SAFE_RELEASE(pPD);
		SAFE_RELEASE(pSD);
		SAFE_RELEASE(pHandler);
		return hr;
	}

	HRESULT WinMfImageGrabber::AddSourceNode(
		IMFTopology *pTopology,           // Topology.
		IMFMediaSource *pSource,          // Media source.
		IMFPresentationDescriptor *pPD,   // Presentation descriptor.
		IMFStreamDescriptor *pSD,         // Stream descriptor.
		IMFTopologyNode **ppNode)         // Receives the node pointer.
	{
		IMFTopologyNode *pNode = NULL;

		HRESULT hr = S_OK;
		CHECK_HR(hr = MFCreateTopologyNode(MF_TOPOLOGY_SOURCESTREAM_NODE, &pNode));
		CHECK_HR(hr = pNode->SetUnknown(MF_TOPONODE_SOURCE, pSource));
		CHECK_HR(hr = pNode->SetUnknown(MF_TOPONODE_PRESENTATION_DESCRIPTOR, pPD));
		CHECK_HR(hr = pNode->SetUnknown(MF_TOPONODE_STREAM_DESCRIPTOR, pSD));
		CHECK_HR(hr = pTopology->AddNode(pNode));

		// Return the pointer to the caller.
		*ppNode = pNode;
		(*ppNode)->AddRef();

	done:
		SafeRelease(&pNode);
		return hr;
	}

	HRESULT WinMfImageGrabber::AddOutputNode(
		IMFTopology *pTopology,     // Topology.
		IMFActivate *pActivate,     // Media sink activation object.
		DWORD dwId,                 // Identifier of the stream sink.
		IMFTopologyNode **ppNode)   // Receives the node pointer.
	{
		IMFTopologyNode *pNode = NULL;

		HRESULT hr = S_OK;
		CHECK_HR(hr = MFCreateTopologyNode(MF_TOPOLOGY_OUTPUT_NODE, &pNode));
		CHECK_HR(hr = pNode->SetObject(pActivate));
		CHECK_HR(hr = pNode->SetUINT32(MF_TOPONODE_STREAMID, dwId));
		CHECK_HR(hr = pNode->SetUINT32(MF_TOPONODE_NOSHUTDOWN_ON_REMOVE, FALSE));
		CHECK_HR(hr = pTopology->AddNode(pNode));

		// Return the pointer to the caller.
		*ppNode = pNode;
		(*ppNode)->AddRef();

	done:
		SafeRelease(&pNode);
		return hr;
	}



	HRESULT WinMfImageGrabber::startGrabbing(void) {
		PROPVARIANT var;
		PropVariantInit(&var);

		HRESULT hr = S_OK;
		hr = ig_pSession->SetTopology(0, ig_pTopology);
		hr = ig_pSession->Start(&GUID_NULL, &var);
		return hr;
	}

	bool WinMfImageGrabber::processGrabbing() {
		HRESULT hrStatus = S_OK;
		MediaEventType met;
		HRESULT hr;
		IMFMediaEvent *pEvent = NULL;

		if(!ig_pSession) return false;
		hr = ig_pSession->GetEvent(0, &pEvent);
		if(!SUCCEEDED(hr)) {
			hr = S_OK;

			return false;
		}

		hr = pEvent->GetStatus(&hrStatus);
		if(!SUCCEEDED(hr)) {
			hr = S_OK;

			return false;
		}

		hr = pEvent->GetType(&met);
		if(!SUCCEEDED(hr)) {
			hr = S_OK;

			return false;
		}

		if(met == MESessionEnded) {
			LogInfo() << "IMAGEGRABBER VIDEODEVICE %i: MESessionEnded ";

			ig_pSession->Stop();

			return false;
		}

		if(met == MESessionStopped) {
			LogInfo() << "IMAGEGRABBER VIDEODEVICE %i: MESessionStopped";

			return false;
		}


		if(met == MEVideoCaptureDeviceRemoved) {
			LogInfo() << "IMAGEGRABBER VIDEODEVICE %i: MEVideoCaptureDeviceRemoved";

			return false;
		}

		SafeRelease(&pEvent);
		return true;
	}

	void WinMfImageGrabber::stopGrabbing() {
		if(ig_pSession) {
			ig_pSession->Stop();
			ig_pSession->Shutdown();
			SAFE_RELEASE(ig_pSession);
			SAFE_RELEASE(ig_pTopology);
		}

		/*IMFMediaSource *ig_pSource;
		IMFMediaSession *ig_pSession;
		IMFTopology *ig_pTopology;*/

	}

	// IUnknown methods
	STDMETHODIMP WinMfImageGrabber::QueryInterface(REFIID iid, void** ppv) {
		// Creation tab of shifting interfaces from start of this class
		static const QITAB qit[] = {

			QITABENT(WinMfImageGrabber, IMFSampleGrabberSinkCallback),
			QITABENT(WinMfImageGrabber, IMFClockStateSink), { 0 }
		};
		return QISearch(this, qit, iid, ppv);
	}
	STDMETHODIMP_(ULONG) WinMfImageGrabber::AddRef() {
		return InterlockedIncrement(&m_cRef);
	}
	STDMETHODIMP_(ULONG) WinMfImageGrabber::Release() {
		ULONG cRef = InterlockedDecrement(&m_cRef);
		if(cRef == 0) {
			delete this;
		}
		return cRef;
	}

	// IMFClockStateSink methods
	STDMETHODIMP WinMfImageGrabber::OnClockStart(MFTIME hnsSystemTime, LONGLONG llClockStartOffset) {
		return S_OK;
	}
	STDMETHODIMP WinMfImageGrabber::OnClockStop(MFTIME hnsSystemTime) {
		return S_OK;
	}
	STDMETHODIMP WinMfImageGrabber::OnClockPause(MFTIME hnsSystemTime) {
		return S_OK;
	}
	STDMETHODIMP WinMfImageGrabber::OnClockRestart(MFTIME hnsSystemTime) {
		return S_OK;
	}
	STDMETHODIMP WinMfImageGrabber::OnClockSetRate(MFTIME hnsSystemTime, float flRate) {
		return S_OK;
	}

	// IMFSampleGrabberSinkCallback methods
	STDMETHODIMP WinMfImageGrabber::OnSetPresentationClock(IMFPresentationClock* pClock) {
		return S_OK;
	}
	STDMETHODIMP WinMfImageGrabber::OnProcessSample(REFGUID guidMajorMediaType, DWORD dwSampleFlags,
		LONGLONG llSampleTime, LONGLONG llSampleDuration, const BYTE * pSampleBuffer,
		DWORD dwSampleSize) 
	{
		LogInfo() << "Process sample";
		buffer_active->SetData((void*)pSampleBuffer);
		camera->OnImageGrabbed(this, buffer_active);
		return 0;
	}
	STDMETHODIMP WinMfImageGrabber::OnShutdown() {
		return S_OK;
	}
}