//---------------------------------------------------------------------------
//	AVFile - library for audio and video files processing
//	2002 - 2007 Stanislav Sumec <sumec@fit.vutbr.cz>
//
//	dvfile.cpp
//	DV camera input 
//
//	10.2.2004
//		- first release
//---------------------------------------------------------------------------

#include "dshow_hack.h"
#include "BaseClasses_hack.h"
#include "dvfile.h"
#include "dsfilei.h"

//---------------------------------------------------------------------------
TDVFile::TDVFile()
{
	Internal = new TDSFileI();
	Description = "DirectShow camera input";
	Implemented = iGetFrameA + iGetFrameA2 + iDIBGetFrameA;
	CoInitialized = false;
	Bitmap = NULL;
	FrameBuffer = NULL;
	g_dwGraphRegister = 0;
}

//---------------------------------------------------------------------------
TDVFile::~TDVFile()
{
	Close();
	delete Internal;
}

//---------------------------------------------------------------------------
void TDVFile::OpenA(const char *FileName, int Mode, int Type, bool ThrowExceptions)
{
	USES_CONVERSION;
	HRESULT hr;
	Close();

	if (Mode == omRead)
	{
		if (Type&otVideo)
		{
			char *Error = NULL;
			if (Error == NULL)
			{
				if (!CoInitialized)
				{
					if (FAILED(CoInitializeEx(NULL, COINIT_APARTMENTTHREADED)))
					{
						Error = "Could not CoInitializeEx!";
					}
					else
					{
						CoInitialized = true;
					}
				}
			}

			if (Error == NULL)
			{
				Internal->pGraph.CoCreateInstance(CLSID_FilterGraph);
				if (!Internal->pGraph)
				{
					Error = "Could not create the FilterGraph!";
				}
			}

			if (Error == NULL)
			{
				Internal->pBilder.CoCreateInstance(CLSID_CaptureGraphBuilder2);
				if (!Internal->pBilder)
				{
					Error = "Could not create the FilterGraph!";
				}
				else
				{
					hr = Internal->pBilder->SetFiltergraph(Internal->pGraph);
					if (hr != S_OK)
						Error = "Could not set FilterGraph!";
				}
			}

			if (Error == NULL)
			{
				g_dwGraphRegister = 0;
				Internal->AddGraphToRot(Internal->pGraph, &g_dwGraphRegister);
			}

			CComPtr<IBaseFilter> pDVCapture;
			if (FileName != NULL && strncmp(FileName, "rtsp://", 7) == 0)
			{
				CComPtr<IBaseFilter> pSource;
				pSource.CoCreateInstance(CLSID_RTSPSource);
				if (!pSource)
				{
					Error = "Could not find RTSP source!";
				}

				if (Error == NULL)
				{
					hr = Internal->pGraph->AddFilter(pSource, L"RTSPSource");
					if (FAILED(hr))
					{
						Error = "Could not add filter!";
					}
				}
/*
				if (Error == NULL)
				{
					pDVCapture.CoCreateInstance(CLSID_MoonlightDemuxer);
					if (!pDVCapture)
					{
						Error = "Could not find MPEG demuxer!";
					}
				}

				if (Error == NULL)
				{
					hr = Internal->pGraph->AddFilter(pDVCapture, L"MPEG demuxer");
					if (FAILED(hr))
					{
						Error = "Could not add filter!";
					}
				}

				if (Error == NULL)
				{
					hr = Internal->pBilder->RenderStream(NULL, NULL, pSource, NULL, pDVCapture);
					if (FAILED(hr))
					{
						Error = "Could not connect source filter to demuxer!";
					}
				}
*/
				if (Error == NULL)
				{
					CComQIPtr<IFileSourceFilter, &IID_IFileSourceFilter> pLoad(pSource);
					if (pLoad)
					{
						hr = pLoad->Load(T2W(LPTSTR(FileName)), NULL);
						if (FAILED(hr))
						{
							Error = "Could not connect to server!";
						}
					}
				}

				pDVCapture = pSource;

				if (Error == NULL)
				{
					IPin *Pin = Internal->WaitForPin(pDVCapture, PINDIR_OUTPUT, 0, 30000);
					if (Pin)
					{
						Pin->Release();
					}
				}
			}
			else
			{
				// Create the System Device Enumerator.
				ICreateDevEnum *pDevEnum = NULL;
				if (Error == NULL)
				{
					hr = CoCreateInstance(CLSID_SystemDeviceEnum, NULL, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, (void**)(&pDevEnum));
					if (FAILED(hr))
					{
						Error = "Could not create System Device Enumerator!";
					}
				}

				// Create an enumerator for the video capture category.
				IEnumMoniker *pEnum = NULL;
				if (Error == NULL)
				{
					hr = pDevEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &pEnum, 0);
					if (pEnum == NULL)
					{
						pDevEnum->Release();
						Error = "Could not create an enumerator for the video capture category!";
					}
				}

				if (Error == NULL)
				{
					IMoniker *pMoniker = NULL;
					int UnknownCount = 0;
					while (pEnum->Next(1, &pMoniker, NULL) == S_OK)
					{
						IPropertyBag *pPropBag = NULL;
						hr = pMoniker->BindToStorage(0, 0, IID_IPropertyBag, (void**)(&pPropBag));
						if (FAILED(hr))
						{
								pMoniker->Release();
								continue;  // Skip this one, maybe the next one will work.
						}

						TDevice Device = Internal->GetDevice(pPropBag, UnknownCount);
						pPropBag->Release();
						pMoniker->Release();

						if (FileName == NULL || Device == FileName)
						{
							hr = pMoniker->BindToObject(0, 0, IID_IBaseFilter, (void**)&pDVCapture);
							if (SUCCEEDED(hr))
							{
								break;
							}
						}
					}

					pEnum->Release();
					pDevEnum->Release();

					if (!pDVCapture)
					{
						Error = "Could not find DV device!";
					}

					if (Error == NULL)
					{
						hr = Internal->pGraph->AddFilter(pDVCapture, L"DVCapture");
						if (FAILED(hr))
						{
							Error = "Could not add filter!";
						}
					}
				}
			}

			if (Error == NULL)
			{
				Internal->pGrabber.CoCreateInstance(CLSID_SampleGrabber);
				if (!Internal->pGrabber)
				{
					Error = "Could not create SampleGrabber!";
				}
			}
			CComQIPtr<IBaseFilter, &IID_IBaseFilter> pGrabberBase(Internal->pGrabber);

			if (Error == NULL)
			{
				hr = Internal->pGraph->AddFilter(pGrabberBase, L"Grabber");
				if (FAILED(hr))
				{
					Error = "Could not add filter!";
				}
			}

			if (Error == NULL)
			{
				// Tell the grabber to grab 24-bit video. Must do this
				// before connecting it
				CMediaType GrabType;
				GrabType.SetType(&MEDIATYPE_Video);
				GrabType.SetSubtype(&MEDIASUBTYPE_RGB32);
				hr = Internal->pGrabber->SetMediaType(&GrabType);
				if (FAILED(hr))
				{
					Error = "Could not set media type!";
				}
			}

			if (Error == NULL)
			{
				if (FAILED(Internal->pGrabber->SetBufferSamples(TRUE)) || FAILED(Internal->pGrabber->SetOneShot(FALSE)))
				{
					Error = "Could not setup grabber!";
				}
			}

			if (Error == NULL)
			{
				hr = Internal->pBilder->RenderStream(NULL, NULL, pDVCapture, NULL, pGrabberBase);
				if (FAILED(hr))
				{
					Error = "Could not connect source filter to grabber!";
				}
			}

			CComPtr<IBaseFilter> pNullRenderer;
			if (Error == NULL)
			{
				pNullRenderer.CoCreateInstance(CLSID_NullRenderer);
				if (!pNullRenderer)
				{
					Error = "Could not create NullRenderer!";
				}
			}

			if (Error == NULL)
			{
				hr = Internal->pGraph->AddFilter(pNullRenderer, L"NullRenderer");
				if (FAILED(hr))
				{
					Error = "Could not add filter!";
				}
			}

			if (Error == NULL)
			{
				hr = Internal->pBilder->RenderStream(NULL, NULL, pGrabberBase, NULL, pNullRenderer);
				if (FAILED(hr))
				{
					Error = "Could not connect null renderer!";
				}
			}

			if (Error == NULL)
			{
				AM_MEDIA_TYPE mt;
				hr = Internal->pGrabber->GetConnectedMediaType(&mt);
				if (FAILED(hr))
				{
					Error = "Could not get media type!";
				}
				else
				{
					VIDEOINFOHEADER *vih = (VIDEOINFOHEADER*)mt.pbFormat;

					VideoInfo.Width = vih->bmiHeader.biWidth;
					VideoInfo.Height = vih->bmiHeader.biHeight;
					VideoInfo.FrameRate.Rate = 25;
					VideoInfo.FrameRate.Scale = 1;

					ScanLineWidth = VideoInfo.Width * 4;
					AlignedScanLineWidth = DIB_SCAN_LINE(ScanLineWidth);

					BitmapDataSize = VideoInfo.Height * AlignedScanLineWidth;
					Bitmap = (BITMAPINFO*)GetMem(sizeof(BITMAPINFOHEADER) + BitmapDataSize);
					if (Bitmap == NULL)
					{
						Error = "Could not allocate memory!";
					}
					else
					{
						BitmapData = (BYTE*)Bitmap + sizeof(BITMAPINFOHEADER);
						Bitmap->bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
						Bitmap->bmiHeader.biWidth = VideoInfo.Width;
						Bitmap->bmiHeader.biHeight = VideoInfo.Height;

						Bitmap->bmiHeader.biPlanes = 1;
						Bitmap->bmiHeader.biBitCount = 32;
						Bitmap->bmiHeader.biCompression = BI_RGB;
						Bitmap->bmiHeader.biSizeImage = BitmapDataSize;
						Bitmap->bmiHeader.biXPelsPerMeter = 0;
						Bitmap->bmiHeader.biYPelsPerMeter = 0;
						Bitmap->bmiHeader.biClrUsed = 0;
						Bitmap->bmiHeader.biClrImportant = 0;
					}
					FreeMediaType(mt);
				}
			}

			if (Error == NULL)
			{
				FrameBuffer = (BYTE*)GetMem(VideoInfo.Width * 4 * VideoInfo.Height);
				if (FrameBuffer == NULL)
				{
					Error = "Could not allocate memory!";
				}
			}
			
			if (Error == NULL)
			{
				CComQIPtr<IMediaControl, &IID_IMediaControl> pControl(Internal->pGraph);
				hr = pControl->Run();
				if (hr == S_FALSE)
				{
					OAFilterState State;
					pControl->GetState(INFINITE, &State);
				}
			}

			if (Error != NULL)
			{
				Close();
				if (ThrowExceptions)
				{
					throw AVFileErrorOpenVideo(Error);
				}
			}
			else
			{
				Length = 1;
				FrameInBuffer = ~0;
				VideoPos = 0;
				OpenMode = Mode;
				OpenType |= otVideo;
			}
		}
	}

	if (ThrowExceptions)
	{
		if (Mode != OpenMode)
		{
			Close();
			throw AVFileErrorOpenInvalidMode();
		}
		if ((Type&otVideo) && !(OpenType&otVideo))
		{
			Close();
			throw AVFileErrorOpenVideo();
		}
		if ((Type&otAudio) && !(OpenType&otAudio))
		{
			Close();
			throw AVFileErrorOpenAudio();
		}
	}
}

//---------------------------------------------------------------------------
void TDVFile::Close()
{
	OpenMode = omClosed;
	OpenType = otNone;
	FirstFrame = true;

	if (Internal->pBilder)
	{
		Internal->pBilder.Release();
	}

	if (Internal->pGraph)
	{
		if (g_dwGraphRegister != 0)
		{
			Internal->RemoveGraphFromRot(g_dwGraphRegister);
		}
		Internal->pGraph.Release();
	}

	if (Internal->pGrabber)
	{
		Internal->pGrabber.Release();
	}

	if (CoInitialized)
	{
		CoUninitialize();
		CoInitialized = false;
	}

	if (Bitmap != NULL)
	{
		FreeMem(Bitmap);
		Bitmap = NULL;
	}

	if (FrameBuffer != NULL)
	{
		FreeMem(FrameBuffer);
		FrameBuffer = NULL;
	}
}

//---------------------------------------------------------------------------
unsigned int TDVFile::GetVideoLength()
{
	CheckOpened();
	CheckVideoStream();
	return Length;
}

//---------------------------------------------------------------------------
BITMAPINFO* TDVFile::DIBGetFrameA()
{
	CheckRead();
	CheckVideoStream();

	HRESULT hr;
	int Count = 10;

	while (true)
	{
		hr = Internal->pGrabber->GetCurrentBuffer((long*)&BitmapDataSize, (long*)BitmapData);
		if (FAILED(hr))
		{
			if (!FirstFrame || Count == 0)
			{
				throw AVFileErrorGetVideo();
			}
			else
			{
				Count--;
				Sleep(200);
			}
		}
		else
		{
			break;
		}
	} 

	FirstFrame = false;
	VideoPos = 0;
	return Bitmap;
}

//---------------------------------------------------------------------------
BYTE* TDVFile::GetFrameA()
{
	BITMAPINFO *Bitmap = DIBGetFrameA();
	BYTE *Src = (BYTE*)Bitmap + Bitmap->bmiHeader.biSize + Bitmap->bmiHeader.biSizeImage;
	BYTE *Dst = FrameBuffer;
	for (int i = 0; i < VideoInfo.Height; i++)
	{
		Src -= AlignedScanLineWidth;
		memcpy(Dst, Src, ScanLineWidth);
		Dst += ScanLineWidth;
	}
	return FrameBuffer;
}

//---------------------------------------------------------------------------
void TDVFile::GetFrameA2(BYTE *Buffer)
{
	BITMAPINFO *Bitmap = DIBGetFrameA();
	BYTE *Src = (BYTE*)Bitmap + Bitmap->bmiHeader.biSize + Bitmap->bmiHeader.biSizeImage;
	BYTE *Dst = Buffer;
	for (int i = 0; i < VideoInfo.Height; i++)
	{
		Src -= AlignedScanLineWidth;
		memcpy(Dst, Src, ScanLineWidth);
		Dst += ScanLineWidth;
	}
}

//---------------------------------------------------------------------------
TDVFile::TDevices TDVFile::GetDevices()
{
	TDevices Devices;

	if (FAILED(CoInitializeEx(NULL, COINIT_APARTMENTTHREADED)))
	{
		throw AVFileError("Could not CoInitializeEx!");
	}

	HRESULT hr;
	ICreateDevEnum *pDevEnum = NULL;
	IEnumMoniker *pEnum = NULL;

	hr = CoCreateInstance(CLSID_SystemDeviceEnum, NULL, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, (void**)(&pDevEnum));
	if (hr != S_OK)
	{
		CoUninitialize();
		throw AVFileError("Could not create System Device Enumerator!");
	}

	hr = pDevEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &pEnum, 0);
	if (hr != S_OK)
	{
		pDevEnum->Release();
		CoUninitialize();
		throw AVFileError("Could not create an enumerator for the video capture category!");
	}

	IMoniker *pMoniker = NULL;
	int UnknownCount = 0;
	while (pEnum->Next(1, &pMoniker, NULL) == S_OK)
	{
		IPropertyBag *pPropBag;
		hr = pMoniker->BindToStorage(0, 0, IID_IPropertyBag, (void**)(&pPropBag));
		if (hr != S_OK)
		{
			pMoniker->Release();
			continue;
		}
		Devices.push_back(Internal->GetDevice(pPropBag, UnknownCount));
		pPropBag->Release();
		pMoniker->Release();
	}

	pEnum->Release();
	pDevEnum->Release();
	CoUninitialize();
	return Devices;
}

