//Don't forget to change project settings:
//1. C++: add include path to DirectShow include folder (such as c:\dxsdk\include)
//2. Link: add link path to DirectShow lib folder (such as c:\dxsdk\lib).
//3. Link: add strmiids.lib and quartz.lib

#include "stdafx.h"
#include <DShow.h>
#include <atlbase.h>
#include <initguid.h>
#include <dvdmedia.h>
#include <qedit.h>
#include <stdio.h>

BOOL hrcheck(HRESULT hr, TCHAR* errtext)
{
    if (hr >= S_OK)
        return FALSE;
    TCHAR szErr[MAX_ERROR_TEXT_LEN];
    DWORD res = AMGetErrorText(hr, szErr, MAX_ERROR_TEXT_LEN);
    if (res)
        printf("Error %x: %s\n%s\n",hr, errtext,szErr);
    else
        printf("Error %x: %s\n", hr, errtext);
    return TRUE;
}

//change this macro to fit your style of error handling
#define CHECK_HR(hr, msg) if (hrcheck(hr, msg)) return hr;

CComPtr<IPin> GetPin(IBaseFilter *pFilter, LPCOLESTR pinname)
{
    CComPtr<IEnumPins>  pEnum;
    CComPtr<IPin>       pPin;

    HRESULT hr = pFilter->EnumPins(&pEnum);
    if (hrcheck(hr, "Can't enumerate pins."))
        return NULL;

    while(pEnum->Next(1, &pPin, 0) == S_OK)
    {
        PIN_INFO pinfo;
        pPin->QueryPinInfo(&pinfo);
        BOOL found = !wcsicmp(pinname, pinfo.achName);
        if (pinfo.pFilter) pinfo.pFilter->Release();
        if (found)
            return pPin;
        pPin.Release();
    }
    printf("Pin not found!\n");
    return NULL;  
}

class CMyCB : public ISampleGrabberCB {
	public: 
		CMyCB(const char * fname, int nframe, int width, int sx, int sy, int _w, int _h) 
			: dstname(fname), m_width(width), m_sx(sx), m_sy(sy), w(_w), h(_h), req_fn(nframe), fn(0) { }

		STDMETHODIMP QueryInterface(REFIID riid, void **ppv) {      
		if (NULL == ppv) return E_POINTER;
			if (riid == __uuidof(IUnknown))
			{
				*ppv = static_cast<IUnknown*>(this);
				return S_OK;
			}
			if (riid == __uuidof(ISampleGrabberCB))
			{
				*ppv = static_cast<ISampleGrabberCB*>(this);
				return S_OK;
			}
			return E_NOINTERFACE;
		};                                                         
		STDMETHODIMP_(ULONG) AddRef() {                             
			return S_OK; 
		};                                                          
		STDMETHODIMP_(ULONG) Release() {                            
			return S_OK;
		};

		//ISampleGrabberCB
		STDMETHODIMP SampleCB(double SampleTime, IMediaSample *pSample);
		STDMETHODIMP BufferCB(double SampleTime, BYTE *pBuffer, long BufferLen) { return S_OK; };	

	protected:
		int m_width, m_sx, m_sy, w, h, fn, req_fn;
		const char* dstname;
};

volatile bool done = false;

STDMETHODIMP CMyCB::SampleCB(double SampleTime, IMediaSample *pSample)
{
	if (!pSample) return E_POINTER;
	if (fn==req_fn) {
		long sz = pSample->GetActualDataLength();
		BYTE *pBuf = NULL;
		pSample->GetPointer(&pBuf);

		int dsz = w * h * 4;
		BITMAPFILEHEADER hdr;
		memset(&hdr, 0, sizeof(BITMAPFILEHEADER));
		hdr.bfType = 0x4d42;
		hdr.bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + dsz;
		hdr.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
		BITMAPINFOHEADER bmi;
		memset(&bmi, 0, sizeof(bmi));
		bmi.biBitCount = 32;
		bmi.biHeight = h;
		bmi.biPlanes = 1;
		bmi.biSize = sizeof(bmi);
		bmi.biSizeImage = dsz;
		bmi.biWidth = w;

		FILE *of = fopen(dstname, "wb");
		if (!of) {
			printf("error writing to %s\n", dstname);
			return E_FAIL;
		}
		fwrite(&hdr, 1, sizeof(BITMAPFILEHEADER), of);
		fwrite(&bmi, 1, sizeof(BITMAPINFOHEADER), of);
		//fwrite(pDst, 1, sz, of);
		for(int y=m_sy; y<m_sy + h; y++)
			fwrite(&pBuf[(y * m_width + m_sx)*4], 4, w, of); 

		fclose(of);
		done = true;
	}
	fn++;
	return S_OK;
}

HRESULT BuildGraph(IGraphBuilder *pGraph, LPCOLESTR srcFile1, const char *outfile, int fn, int sx, int sy)
{
	HRESULT hr = S_OK;

    //graph builder
    CComPtr<ICaptureGraphBuilder2> pBuilder;
    hr = pBuilder.CoCreateInstance(CLSID_CaptureGraphBuilder2);
    CHECK_HR(hr, "Can't create Capture Graph Builder");
    hr = pBuilder->SetFiltergraph(pGraph);
    CHECK_HR(hr, "Can't SetFiltergraph");

    //add File Source (Async.)
    CComPtr<IBaseFilter> pFileSourceAsync;
    hr = pFileSourceAsync.CoCreateInstance(CLSID_AsyncReader);
    CHECK_HR(hr, "Can't create File Source (Async.)");
    hr = pGraph->AddFilter(pFileSourceAsync, L"File Source (Async.)");
    CHECK_HR(hr, "Can't add File Source (Async.) to graph");
    //set source filename
    CComQIPtr<IFileSourceFilter, &IID_IFileSourceFilter> pFileSourceAsync_src(pFileSourceAsync);
    if (!pFileSourceAsync_src)
        CHECK_HR(E_NOINTERFACE, "Can't get IFileSourceFilter");
    hr = pFileSourceAsync_src->Load(srcFile1, NULL);
    CHECK_HR(hr, "Can't load file");


    //add SampleGrabber
    CComPtr<IBaseFilter> pSampleGrabber;
    hr = pSampleGrabber.CoCreateInstance(CLSID_SampleGrabber);
    CHECK_HR(hr, "Can't create SampleGrabber");
    hr = pGraph->AddFilter(pSampleGrabber, L"SampleGrabber");
    CHECK_HR(hr, "Can't add SampleGrabber to graph");
    AM_MEDIA_TYPE pSampleGrabber_pmt;
    ZeroMemory(&pSampleGrabber_pmt, sizeof(AM_MEDIA_TYPE));
    pSampleGrabber_pmt.majortype = MEDIATYPE_Video;
    pSampleGrabber_pmt.subtype = MEDIASUBTYPE_RGB32;

    CComQIPtr<ISampleGrabber, &IID_ISampleGrabber> pSampleGrabber_isg(pSampleGrabber);
    hr = pSampleGrabber_isg->SetMediaType(&pSampleGrabber_pmt);
    CHECK_HR(hr, "Can't set media type to sample grabber");

	hr = pBuilder->RenderStream(NULL, NULL, pFileSourceAsync, NULL, pSampleGrabber);
    CHECK_HR(hr, "Can't connect File Source (Async.) and grabber");

	AM_MEDIA_TYPE mt;
	pSampleGrabber_isg->GetConnectedMediaType(&mt);
	VIDEOINFOHEADER* vih = (VIDEOINFOHEADER*)mt.pbFormat;

	//CMyCB(const char * fname, int nframe, int width, int sx, int sy, int _w, int _h)
	int width = vih->bmiHeader.biWidth;
	int height = abs(vih->bmiHeader.biHeight);
	if (width < 600 || height < 400) {
		printf("video dimensions less than 600x400\n");
		return E_INVALIDARG;
	}
	CMyCB* pCB = new CMyCB(outfile, fn, width, sx, height-400-sy, 600, 400);
	pSampleGrabber_isg->SetCallback(pCB, 0); 

    //add Null Renderer
    CComPtr<IBaseFilter> pNullRenderer;
    hr = pNullRenderer.CoCreateInstance(CLSID_NullRenderer);
    CHECK_HR(hr, "Can't create Null Renderer");
    hr = pGraph->AddFilter(pNullRenderer, L"Null Renderer");
    CHECK_HR(hr, "Can't add Null Renderer to graph");


    //connect SampleGrabber and Null Renderer
    hr = pBuilder->RenderStream(NULL, &MEDIATYPE_Video, pSampleGrabber, NULL, pNullRenderer);
    CHECK_HR(hr, "Can't connect SampleGrabber and Null Renderer");

	CComQIPtr<IMediaFilter> pIMF(pGraph);
	if (!!pIMF) pIMF->SetSyncSource(NULL);

    return S_OK;
}

//int _tmain(int argc, _TCHAR* argv[]) //use this line in VS2008
int main(int argc, char* argv[])
{
	if (argc<6) {
		printf("usage: grab video.avi video.bmp frame_number X Y\n");
		return 0;
	}
	unsigned short fname[1024]={0};
	MultiByteToWideChar(CP_ACP, 0, argv[1], strlen(argv[1]), fname, 1020);

    CoInitialize(NULL);
    CComPtr<IGraphBuilder> graph;
    graph.CoCreateInstance(CLSID_FilterGraph);

    printf("Building graph...\n");
    HRESULT hr = BuildGraph(graph, fname, argv[2], atoi(argv[3]), atoi(argv[4]), atoi(argv[5]));
    if (hr==S_OK) {
        printf("Running");
        CComQIPtr<IMediaControl, &IID_IMediaControl> mediaControl(graph);
        hr = mediaControl->Run();
        CHECK_HR(hr, "Can't run the graph");
        CComQIPtr<IMediaEvent, &IID_IMediaEvent> mediaEvent(graph);
        BOOL stop = FALSE;
        MSG msg;
        while(!stop) 
        {
            long ev=0, p1=0, p2=0;
            Sleep(500);
            printf(".");
            while(PeekMessage(&msg, NULL, 0,0, PM_REMOVE))
                DispatchMessage(&msg);
            while (mediaEvent->GetEvent(&ev, &p1, &p2, 0)==S_OK)
            {
                if (ev == EC_COMPLETE || ev == EC_USERABORT)
                {
                    printf("Done!\n");
                    stop = TRUE;
                }
                else
                if (ev == EC_ERRORABORT)
                {
                    printf("An error occured: HRESULT=%x\n", p1);
                    mediaControl->Stop();
                    stop = TRUE;
                }
                mediaEvent->FreeEventParams(ev, p1, p2);
            }
			if (done) {
				mediaControl->Stop();
                stop = TRUE;
			}
        }
    }
    CoUninitialize();
    return 0;
}