#include <iostream>
#include <string>
#include <vector>
#include <atlbase.h>
#include <mfapi.h>
#include <mfidl.h>
#include <wmcodecdsp.h>
#include <wmcontainer.h>
#include <comdef.h>
#include <dmo.h>
#include <dshow.h>

#pragma comment(lib, "mf.lib")
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mfuuid.lib")
#pragma comment(lib, "strmiids.lib")
#pragma comment(lib, "winmm.lib")
#pragma comment(lib, "msdmo.lib")

class ComException : public std::exception
{
public:
	ComException(HRESULT errorCode, const std::string& message) :
		errorCode(errorCode),
		message(message)
	{}
	friend std::ostream& operator<<(std::ostream& s, const ComException& e)
	{
		_com_error comError(e.errorCode);
		return s << "Error in '" << e.message << "': 0x" << std::hex << e.errorCode << " - " << comError.ErrorMessage();
	}
private:
	HRESULT errorCode;
	std::string message;
};


#define THROW_ON_FAIL(expr) {                                      \
	HRESULT _errorCode = expr;                                     \
	if(FAILED(_errorCode)) throw ComException(_errorCode, #expr);  \
}


static const GUID Codec = MFVideoFormat_MSS2;
static const int FrameWidth = 1280;
static const int FrameHeight = 720;
static const int FrameRate = 25;
static const int Bitrate = 1000000;

static CComPtr<IMFMediaType> createInputMediaType()
{
	CComPtr<IMFMediaType> mediaType;
	THROW_ON_FAIL(MFCreateMediaType(&mediaType));
	THROW_ON_FAIL(mediaType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video));
	THROW_ON_FAIL(mediaType->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_RGB32));
	THROW_ON_FAIL(mediaType->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive));
	THROW_ON_FAIL(MFSetAttributeSize(mediaType, MF_MT_FRAME_SIZE, FrameWidth, FrameHeight));
	THROW_ON_FAIL(MFSetAttributeRatio(mediaType, MF_MT_FRAME_RATE, FrameRate, 1));
	return mediaType;
}

static CComPtr<IMFMediaType> createOutputMediaType()
{
	CComPtr<IMFMediaType> mediaType;
	THROW_ON_FAIL(MFCreateMediaType(&mediaType));
	THROW_ON_FAIL(mediaType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video));
	THROW_ON_FAIL(mediaType->SetGUID(MF_MT_SUBTYPE, Codec));
	THROW_ON_FAIL(mediaType->SetUINT32(MF_MT_AVG_BITRATE, Bitrate));
	THROW_ON_FAIL(mediaType->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive));
	THROW_ON_FAIL(MFSetAttributeSize(mediaType, MF_MT_FRAME_SIZE, FrameWidth, FrameHeight));
	THROW_ON_FAIL(MFSetAttributeRatio(mediaType, MF_MT_FRAME_RATE, FrameRate, 1));
	return mediaType;
}

static CComPtr<IMFTransform> findRegisteredEncoderTransform(IMFMediaType* outputType)
{
	MFT_REGISTER_TYPE_INFO outputTypeInfo = {};
	THROW_ON_FAIL(outputType->GetGUID(MF_MT_MAJOR_TYPE, &outputTypeInfo.guidMajorType));
	THROW_ON_FAIL(outputType->GetGUID(MF_MT_SUBTYPE, &outputTypeInfo.guidSubtype));

	IMFActivate** activates = nullptr;
	UINT32 activateCount = 0;

	THROW_ON_FAIL(MFTEnumEx(MFT_CATEGORY_VIDEO_ENCODER, MFT_ENUM_FLAG_ALL, nullptr, &outputTypeInfo, &activates, &activateCount));
	if(activateCount == 0) throw ComException(E_FAIL, "Can't find encoder");

	CComPtr<IMFActivate> activate = activates[0];

	for (UINT32 i = 0; i < activateCount; i++)
	{
		 activates[i]->Release();
	}
	CoTaskMemFree(activates);

	CComPtr<IMFTransform> transform;
	THROW_ON_FAIL(activate->ActivateObject(IID_PPV_ARGS(&transform)));
	return transform;
}

static void addPrivateData(IMFTransform* transform, IMFMediaType* mediaType)
{
	CComQIPtr<IWMCodecPrivateData> privateData = transform;
	if(!privateData) throw ComException(E_FAIL, "Can't get codec private data");
	DMO_MEDIA_TYPE dmoType = {};
	try
	{
		THROW_ON_FAIL(MFInitAMMediaTypeFromMFMediaType(mediaType, FORMAT_VideoInfo, (AM_MEDIA_TYPE*)&dmoType));
		THROW_ON_FAIL(privateData->SetPartialOutputType(&dmoType));

		ULONG dataSize = 0;
		THROW_ON_FAIL(privateData->GetPrivateData(nullptr, &dataSize));
		std::vector<BYTE> data;
		data.resize(dataSize);
		THROW_ON_FAIL(privateData->GetPrivateData(&data[0], &dataSize));  // FIXME: E_NOTIMPL error here
		THROW_ON_FAIL(mediaType->SetBlob(MF_MT_USER_DATA, &data[0], dataSize));
	}
	catch(...)
	{
		MoFreeMediaType(&dmoType);
		throw;
	}
	MoFreeMediaType(&dmoType);
}

CComPtr<IMFTransform> createEncoderTransform()
{
	CComPtr<IMFMediaType> inputType = createInputMediaType();
	CComPtr<IMFMediaType> outputType = createOutputMediaType();

	CComPtr<IMFTransform> transform = findRegisteredEncoderTransform(outputType);

	THROW_ON_FAIL(transform->SetInputType(0, inputType, 0));
	addPrivateData(transform, outputType);
	THROW_ON_FAIL(transform->SetOutputType(0, outputType, 0));
	return transform;
}

EXTERN_GUID(CLSID_CMSSCEncMediaObject2, 0xf7ffe0a0, 0xa4f5, 0x44b5, 0x94, 0x9e, 0x15, 0xed, 0x2b, 0xc6, 0x6f, 0x9d);

static void registerScreenEncoder()
{
	MFT_REGISTER_TYPE_INFO outputTypeInfo;
	outputTypeInfo.guidMajorType = MFMediaType_Video;
	outputTypeInfo.guidSubtype = MFVideoFormat_MSS2;
	THROW_ON_FAIL(MFTRegisterLocalByCLSID(CLSID_CMSSCEncMediaObject2, MFT_CATEGORY_VIDEO_ENCODER, L"Windows Media Video 9 Screen Encoder", MFT_ENUM_FLAG_ALL, 0, nullptr, 1, &outputTypeInfo));
}

int main()
{
	try
	{
		THROW_ON_FAIL(CoInitialize(nullptr));
		THROW_ON_FAIL(MFStartup(MF_VERSION));
		registerScreenEncoder();
		CComPtr<IMFTransform> transform = createEncoderTransform();
		std::cout << "Transform created " << transform << std::endl;
		THROW_ON_FAIL(MFShutdown());
	}
	catch(const ComException& e)
	{
		std::cout << e << std::endl;
	}
}