/***************************************************
This is the main Grabber code.  It builds the Direct Show graphs
and then inserts SampleGrabbers right before the renderers.  It
then changes all of the renderers to NullRenderers.  This code
supports any number of audio or video streams.  Raw midi streams
are not supported -- I didn't think I should bother.

This code was intended to be used inside of a matlab interface,
but can be used as a generic grabber class for anyone who needs
one.

Written by Micah Richert.
07/14/2005
Modified 09/21/2005, fixed a bug that when no frames are specified it wouldn't capture all frames
**************************************************/

#include "DDGrab.h"
#include <math.h>
#include <assert.h>

GUID NULLGUID = {0};

CSampleGrabberCB::CSampleGrabberCB()
{
	pbFormat = NULL;
	frameNr = 0;
	disabled = false;
	done = false;
	bytesPerWORD = 0;
	rate = 0;
	startTime = 0;
	stopTime = 0;

#ifdef REMOVEMSDLL
	frames = new intListEntry();
	frameBytes = new intListEntry();
	frameNrs = new intListEntry();
#else
	frames = new vector<BYTE*>;
	frameBytes = new vector<int>;
	frameNrs = new vector<int>;
#endif
}

CSampleGrabberCB::~CSampleGrabberCB()
{
	if (pbFormat) free(pbFormat);

#ifdef REMOVEMSDLL
	intListEntry* f=frames, *pf;
	f=f->next; //first is just a place holder
	while(f)
	{
		if(f->data) delete[] ((void*)f->data);
		pf = f;
		f=f->next;
		delete pf;
	}
	frames->next = NULL;
#else
	for (int f=0;f<frames->size();f++) if (frames->at(f)) delete[] (frames->at(f));
#endif
}

// Fake out any COM QI'ing
//
STDMETHODIMP CSampleGrabberCB::QueryInterface(REFIID riid, void ** ppv)
{
	if (!ppv) return E_POINTER;

	if( riid == IID_ISampleGrabberCB || riid == IID_IUnknown )
	{
		*ppv = (void *) static_cast<ISampleGrabberCB*> ( this );
		return NOERROR;
	}

	return E_NOINTERFACE;
}

// The sample grabber is calling us back on its deliver thread.
// This is NOT the main app thread!
//
STDMETHODIMP CSampleGrabberCB::BufferCB( double dblSampleTime, BYTE * pBuffer, long lBufferSize )
{
	if (disabled) return 0;

	if (!pBuffer) return E_POINTER;

	frameNr++;

	int nrWORDs = lBufferSize/bytesPerWORD;

	int offset=0, len=0;

	if (frameNrs->size() > 0)
	{
		//frames are being specified, so we must be a videoCB...
		// check to see if the frame is in our list
		bool foundNr = false;
		int lastFrameNr = 0;
		for (int i=0;i<frameNrs->size();i++)
		{
			if (frameNrs->at(i) == frameNr) foundNr = true;
			if (frameNrs->at(i) > lastFrameNr) lastFrameNr = frameNrs->at(i);
		}

		if (frameNr > lastFrameNr)
		{
			done = true;
			return 0;
		}

		if (foundNr) len = lBufferSize; // since this has to be a video frame, we just capture everything
	} else {
		// either no frames are specified (capture all), or we have time specified
		if (stopTime)
		{
			// time is being used...
			offset = min(max(0,(int)((startTime-dblSampleTime)*rate)*bytesPerWORD),lBufferSize);
			len = min(lBufferSize,(int)((stopTime-dblSampleTime)*rate)*bytesPerWORD)-offset;
			// if we have gone past our stop time...
			if (len < 0)
			{
				done = true;
				return 0;
			}
		} else {
			// capture everything... video or audio
			len = lBufferSize;
		}
	}

	if (len)
	{
		BYTE* tmp = new BYTE[len];
		if (!tmp) return E_OUTOFMEMORY;

		memcpy(tmp,pBuffer+offset,len);
#ifdef REMOVEMSDLL
		new intListEntry((int)tmp,frames);
		new intListEntry(len,frameBytes);
#else
		frames->push_back(tmp);
		frameBytes->push_back(len);
#endif
	}

	return 0;
}

DDGrabber::DDGrabber()
{
#ifdef REMOVEMSDLL
	VideoCBs = new intListEntry();
	AudioCBs = new intListEntry();
#else
	VideoCBs = new vector<CSampleGrabberCB*>;
	AudioCBs = new vector<CSampleGrabberCB*>;
#endif
	stopForced = false;
}

void DDGrabber::cleanUp()
{
	IEnumFilters* filterList;
	IBaseFilter* filt;
	ULONG tmp;
	int i;

#ifdef REMOVEMSDLL
	intListEntry* e=VideoCBs, *pe;
	e=e->next; //first is just a place holder
	while(e)
	{
		if(e->data) delete((CSampleGrabberCB*)e->data);
		pe = e;
		e=e->next;
		delete pe;
	}
	VideoCBs->next = NULL;

	e=AudioCBs;
	e=e->next; //first is just a place holder
	while(e)
	{
		if(e->data) delete((CSampleGrabberCB*)e->data);
		pe = e;
		e=e->next;
		delete pe;
	}
	AudioCBs->next = NULL;
#else
	for (i=0; i<VideoCBs->size(); i++) delete VideoCBs->at(i);
	for (i=0; i<AudioCBs->size(); i++) delete AudioCBs->at(i);

	VideoCBs->clear();
	AudioCBs->clear();
#endif

	if (!FAILED(pGraphBuilder->EnumFilters(&filterList)))
	{
		filterList->Reset();
		while (filterList->Next(1, &filt, &tmp) == S_OK)
		{
			filt->Release();
		}
		filterList->Release();
	}
	pGraphBuilder = NULL;
}

void DDGrabber::MyFreeMediaType(AM_MEDIA_TYPE& mt)
{
	if (mt.cbFormat != 0)
	{
		CoTaskMemFree((PVOID)mt.pbFormat);
		mt.cbFormat = 0;
		mt.pbFormat = NULL;
	}
	if (mt.pUnk != NULL)
	{
		// Unecessary because pUnk should not be used, but safest.
		mt.pUnk->Release();
		mt.pUnk = NULL;
	}
}

PIN_INFO DDGrabber::getPinInfo(IPin* pin)
{
	PIN_INFO	info = {0};

	if (pin)
	{
		if (!FAILED(pin->QueryPinInfo(&info)))
		{
			info.pFilter->Release();
		}
	}

	return info;
}

IPin* DDGrabber::getInputPin(IBaseFilter* filt)
{
	IPin* pin = NULL;
	IEnumPins* pinList;
	ULONG tmp;

	if (!filt) return NULL;

	//get the input
	if (!FAILED(filt->EnumPins(&pinList)))
	{
		pinList->Reset();
		while (pinList->Next(1, &pin, &tmp) == S_OK && getPinInfo(pin).dir != PINDIR_INPUT);
		pinList->Release();

		if (getPinInfo(pin).dir != PINDIR_INPUT) return NULL;
	}

	return pin;
}

IPin* DDGrabber::getOutputPin(IBaseFilter* filt)
{
	IPin* pin = NULL;
	IEnumPins* pinList;
	ULONG tmp;

	if (!filt) return NULL;

	//get the output
	if (!FAILED(filt->EnumPins(&pinList)))
	{
		pinList->Reset();
		while (pinList->Next(1, &pin, &tmp) == S_OK && getPinInfo(pin).dir != PINDIR_OUTPUT);
		pinList->Release();

		if (getPinInfo(pin).dir != PINDIR_OUTPUT) return NULL;
	}

	return pin;
}

bool DDGrabber::isRenderer(IBaseFilter* filt)
{
	if (!filt) return false;

	IEnumPins*	pinList;
	int nrOutput = 0;
	int nrInput = 0;
	IPin*		pin = NULL;
	ULONG		tmp;

	if (FAILED(filt->EnumPins(&pinList))) return false;
	pinList->Reset();
	while (pinList->Next(1, &pin, &tmp) == S_OK)
	{
		if (getPinInfo(pin).dir == PINDIR_OUTPUT) nrOutput++;
		else nrInput++;
	}
	pinList->Release();

	#ifdef _DEBUG
		FILTER_INFO info;
		filt->QueryFilterInfo(&info);
		char str[100];
		WideCharToMultiByte( CP_ACP, 0, info.achName, -1, str, 100, NULL, NULL );
		_RPT0(_CRT_WARN,str);
		_RPT2(_CRT_WARN," %d %d\n", nrOutput, nrInput);
	#endif

	return nrOutput == 0 && nrInput == 1;  // the only filters that have no outputs are renderers
}

IPin* DDGrabber::connectedToInput(IBaseFilter* filt)
{
	if (!filt) return NULL;

	IPin* inPin;
	IPin* outPin;

	inPin = getInputPin(filt);
	if (!inPin) return NULL;

	if (FAILED(inPin->ConnectedTo(&outPin))) return NULL;
	return outPin;
}

GUID DDGrabber::getMajorType(IBaseFilter* filt)
{
	if (!filt) return NULLGUID;

	IPin* inPin;
	inPin = getInputPin(filt);
	if (!inPin) return NULLGUID;

	AM_MEDIA_TYPE mt = {0};
	if (FAILED(inPin->ConnectionMediaType(&mt))) return NULLGUID;

	GUID ret = mt.majortype;
	MyFreeMediaType(mt);

	return ret;
}

HRESULT DDGrabber::getVideoInfo(unsigned int id, int* width, int* height, int* rate, int* nrFramesCaptured, int* nrFramesTotal)
{
	if (!width || !height || !nrFramesCaptured || !nrFramesTotal) return E_POINTER;

	if (id >= VideoCBs->size()) return E_NOINTERFACE;
	CSampleGrabberCB* CB = (CSampleGrabberCB*)VideoCBs->at(id);

	if (!CB) return E_POINTER;

	if (!CB->pbFormat)
	{
		*width = 0;
		*height = 0;
		*rate = 0;
	} else {
		VIDEOINFOHEADER * h = (VIDEOINFOHEADER*) CB->pbFormat;
		if (!h) return E_POINTER;
		*width  = h->bmiHeader.biWidth;
		*height = h->bmiHeader.biHeight;
		if (h->AvgTimePerFrame == 0) *rate = 0; // can't calc rate correctly...
		else *rate = 10000000.0/h->AvgTimePerFrame; // make it samples per second.
	}
	*nrFramesCaptured = CB->frames->size();
	*nrFramesTotal = CB->frameNr;

	// see if we can get more reliable nrFramesTotal and rate information
	CComPtr<IMediaSeeking> pMediaSeeking;
	pGraphBuilder->QueryInterface(IID_IMediaSeeking, (void **)&pMediaSeeking);
	if (pMediaSeeking)
	{
		LONGLONG duration = 0;
		LONGLONG durationus = 0;

		if (SUCCEEDED(pMediaSeeking->SetTimeFormat(&TIME_FORMAT_MEDIA_TIME))) pMediaSeeking->GetDuration(&durationus);

		if (SUCCEEDED(pMediaSeeking->SetTimeFormat(&TIME_FORMAT_FRAME)) && SUCCEEDED(pMediaSeeking->GetDuration(&duration)))
		{
			if (stopForced) *nrFramesTotal = duration; // if we stopped early, calculate nrFramesTotal
			if (*rate == 0 && durationus) *rate = duration/(durationus/10000000.0);
		} else {
			 // if we stopped early, calculate nrFramesTotal
			if (stopForced && durationus && *rate) *nrFramesTotal = -(*rate)*(durationus/10000000.0);
		}
	}

	// some things don't work right if rate is 0
	if (*rate == 0) *rate = 1;

	return S_OK;
}

HRESULT DDGrabber::getAudioInfo(unsigned int id, int* nrChannels, int* rate, int* bits, int* nrFramesCaptured, int* nrFramesTotal)
{
	if (!nrChannels || !rate || !bits || !nrFramesCaptured || !nrFramesTotal) return E_POINTER;

	if (id >= AudioCBs->size()) return E_NOINTERFACE;
	CSampleGrabberCB* CB = (CSampleGrabberCB*)AudioCBs->at(id);

	if (!CB) return E_POINTER;

	if (!CB->pbFormat)
	{
		*nrChannels = 0;
		*rate = 0;