#define _WIN32_DCOM
#include <iostream>
#include "drr.h"
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "socketManager.h"
#include <comdef.h>
#include <Wbemidl.h>
#pragma comment(lib,"rpcrt4.lib")
#pragma comment(lib, "wbemuuid.lib")
#include <chrono>
using namespace std;


bool cuda_check()
{
	bool isSupportCuda = false;
	HRESULT hres;

	// Initialize COM.
	hres = CoInitializeEx(0, COINIT_APARTMENTTHREADED);
	if (FAILED(hres))
	{
		cout << "Failed to initialize COM library. "
			<< "Error code = 0x"
			<< hex << hres << endl;
		return 1;              // Program has failed.
	}

	// Initialize 
	hres = CoInitializeSecurity(
		NULL,
		-1,      // COM negotiates service                  
		NULL,    // Authentication services
		NULL,    // Reserved
		RPC_C_AUTHN_LEVEL_DEFAULT,    // authentication
		RPC_C_IMP_LEVEL_IMPERSONATE,  // Impersonation
		NULL,             // Authentication info 
		EOAC_NONE,        // Additional capabilities
		NULL              // Reserved
	);


	if (FAILED(hres))
	{
		cout << "Failed to initialize security. "
			<< "Error code = 0x"
			<< hex << hres << endl;
		CoUninitialize();
		return 1;          // Program has failed.
	}

	// Obtain the initial locator to Windows Management
	// on a particular host computer.
	IWbemLocator *pLoc = 0;

	hres = CoCreateInstance(
		CLSID_WbemLocator,
		0,
		CLSCTX_INPROC_SERVER,
		IID_IWbemLocator, (LPVOID *)&pLoc);

	if (FAILED(hres))
	{
		cout << "Failed to create IWbemLocator object. "
			<< "Error code = 0x"
			<< hex << hres << endl;
		CoUninitialize();
		return 1;       // Program has failed.
	}

	IWbemServices *pSvc = 0;

	// Connect to the root\cimv2 namespace with the
	// current user and obtain pointer pSvc
	// to make IWbemServices calls.

	hres = pLoc->ConnectServer(

		_bstr_t(L"ROOT\\CIMV2"), // WMI namespace
		NULL,                    // User name
		NULL,                    // User password
		0,                       // Locale
		NULL,                    // Security flags                 
		0,                       // Authority       
		0,                       // Context object
		&pSvc                    // IWbemServices proxy
	);

	if (FAILED(hres))
	{
		cout << "Could not connect. Error code = 0x"
			<< hex << hres << endl;
		pLoc->Release();
		CoUninitialize();
		return 1;                // Program has failed.
	}

	//cout << "Connected to ROOT\\CIMV2 WMI namespace" << endl;

	// Set the IWbemServices proxy so that impersonation
	// of the user (client) occurs.
	hres = CoSetProxyBlanket(

		pSvc,                         // the proxy to set
		RPC_C_AUTHN_WINNT,            // authentication service
		RPC_C_AUTHZ_NONE,             // authorization service
		NULL,                         // Server principal name
		RPC_C_AUTHN_LEVEL_CALL,       // authentication level
		RPC_C_IMP_LEVEL_IMPERSONATE,  // impersonation level
		NULL,                         // client identity 
		EOAC_NONE                     // proxy capabilities     
	);

	if (FAILED(hres))
	{
		cout << "Could not set proxy blanket. Error code = 0x"
			<< hex << hres << endl;
		pSvc->Release();
		pLoc->Release();
		CoUninitialize();
		return 1;               // Program has failed.
	}


	// Use the IWbemServices pointer to make requests of WMI. 
	// Make requests here:

	// For example, query for all the running processes
	IEnumWbemClassObject* pEnumerator = NULL;
	hres = pSvc->ExecQuery(
		bstr_t("WQL"),
		bstr_t("SELECT * FROM Win32_VideoController"),
		WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
		NULL,
		&pEnumerator);

	if (FAILED(hres))
	{
		cout << "Query for processes failed. "
			<< "Error code = 0x"
			<< hex << hres << endl;
		pSvc->Release();
		pLoc->Release();
		CoUninitialize();
		return 1;               // Program has failed.
	}
	else
	{
		IWbemClassObject *pclsObj;
		ULONG uReturn = 0;

		while (pEnumerator)
		{
			hres = pEnumerator->Next(WBEM_INFINITE, 1,
				&pclsObj, &uReturn);
			if (0 == uReturn)
			{
				break;
			}

			VARIANT vtProp;
			VARIANT vtProp2;

			// Get the value of the Name property
			hres = pclsObj->Get(L"Description", 0, &vtProp, 0, 0);
			WCHAR* des;
			des = vtProp.bstrVal;
			if (des[0] == 'N' && des[1] == 'V' && des[2] == 'I' && des[3] == 'D') {
				isSupportCuda = true;
			}
			hres = pclsObj->Get(L"Availability", 0, &vtProp2, 0, 0);

			VariantClear(&vtProp);
			VariantClear(&vtProp2);
		}

	}

	// Cleanup
	// ========

	pSvc->Release();
	pLoc->Release();
	CoUninitialize();
	if (!isSupportCuda)
		std::cout << "Sorry, Your device doesn't have GPU!" << endl;

	return isSupportCuda;
}

int main() {
	cout << "*******************************************************************************" << endl;
	cout << " Please connect the socket." << endl;
	cout << " Help: " << endl;
	cout << " Send the message \"CT(__folder_path__)\" for loading a CT sequence." << endl;
	cout << "     Client will receive the message to inform if the operation is completed." << endl;
	cout << " Send the message \"PARA(rx,ry,rz,dx,dy,threshold)\" for generating a DRR image." << endl;
	cout << "     rx,ry,rz(float) are the rotation angle for the DRR image." << endl;
	cout << "     dx,dy(int) are the width and height of the DRR image." << endl;
	cout << "     threshold(float) can filter the required pixels for DRR images." << endl;
	cout << "     Client will receive a DRR array if the operation is completed." << endl;
	cout << "*******************************************************************************" << endl;
	if (!cuda_check()) return 0;
	socketManager* m = new socketManager;
	m->open_socket();
	return 0;
}