/***************************************************************************
 *   Copyright (C) 2010 by Volker Lanz <vl@fidra.de>                       *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA            *
 ***************************************************************************/

#include "veneer/qt/win/wmi.h"

#include "veneer/device.h"

#include <QDebug>
#include <QVariant>

#ifdef _DEBUG
#	pragma comment(lib, "comsuppwd.lib")
#else
#	pragma comment(lib, "comsuppw.lib")
#endif

#pragma comment(lib, "wbemuuid.lib")

#include <Wbemidl.h>
#include <windows.h>
#include <rpc.h>
#include <comdef.h>

#pragma comment(lib, "wbemuuid.lib")

#include <iostream>

namespace Veneer
{
	class WmiItem
	{
		public:
			WmiItem(IWbemClassObject* p);
			~WmiItem();
			QVariant property(const QString& s) const;

		private:
			IWbemClassObject* m_ClassObject;
	};

	typedef QList<WmiItem*> WmiItems;

	WmiItem::WmiItem(IWbemClassObject* p) :
		m_ClassObject(p)
	{
	}

	WmiItem::~WmiItem()
	{
		m_ClassObject->Release();
		m_ClassObject = NULL;
	}

	QVariant WmiItem::property(const QString& s) const
	{
		// qDebug() << "get property " << s << " for item with class object at " << m_ClassObject;

		VARIANT variant;
		CIMTYPE cimType;
		HRESULT hres = m_ClassObject->Get(bstr_t(qPrintable(s)), 0, &variant, &cimType, 0);

		QVariant rval;
		if (SUCCEEDED(hres))
		{
			switch (cimType)
			{
				case CIM_EMPTY:
					break;

				case CIM_SINT8:
				case CIM_UINT8:
				case CIM_SINT16:
				case CIM_UINT16:
				case CIM_SINT32:
				case CIM_UINT32:
				case CIM_SINT64:
				case CIM_UINT64:
					rval = variant.intVal;
					break;

				case CIM_REAL32:
				case CIM_REAL64:
					rval = variant.fltVal;
						break;

				case CIM_BOOLEAN:
					rval = variant.boolVal;
					break;

				case CIM_STRING:
					rval = QString(reinterpret_cast<QChar*>(variant.bstrVal), wcslen(variant.bstrVal));
					break;

				case CIM_DATETIME:
					rval = variant.date;
					break;

				default:
					qWarning() << "cannot handle variant type " << cimType;
					break;
			}

			VariantClear(&variant);

			// qDebug() << "property: " << s << ", value: " << rval;
		}

		return rval;
	}

	class WmiQuery
	{
		public:
			WmiQuery();
			~WmiQuery();

		public:
			WmiItems exec(const QString& wql);

		private:
			bool m_InitSuccess;
			bool m_NeedUninitialize;
			IWbemLocator* m_Locator;
			IWbemServices* m_Server;
			IEnumWbemClassObject* m_Enumerator;
	};

	WmiQuery::WmiQuery() :
		m_InitSuccess(true),
		m_NeedUninitialize(false),
		m_Locator(NULL),
		m_Server(NULL),
		m_Enumerator(NULL)
	{
		HRESULT hres = CoInitialize(0);

		if (FAILED(hres) && hres != S_FALSE && hres != RPC_E_CHANGED_MODE)
		{
			qWarning() << "Failed to initialize COM library. Error code: 0x" << hex << quint32(hres);
			m_InitSuccess = false;
		}

		m_NeedUninitialize = (hres != S_FALSE && hres != RPC_E_CHANGED_MODE);

		if (m_InitSuccess)
		{
			hres = CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, NULL);

			if (FAILED(hres) && hres != RPC_E_TOO_LATE) // RPC_E_TOO_LATE --> security already initialized
			{
				if (m_NeedUninitialize)
					CoUninitialize();

				m_InitSuccess = false;

				qWarning() << "Failed to initialize security. Error code: " << hres;
			}
		}

		if (m_InitSuccess)
		{
			hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*) &m_Locator);

			if (FAILED(hres))
			{
				qWarning() << "Failed to create IWbemLocator object. Error code: " << hres;

				if (m_NeedUninitialize)
					CoUninitialize();

				m_InitSuccess = false;
			}
		}

		if (m_InitSuccess)
		{
			hres = m_Locator->ConnectServer(_bstr_t("ROOT\\CIMV2"), NULL, NULL, 0, NULL, 0, 0, &m_Server);

			if (FAILED(hres))
			{
				qWarning() << "Could not connect. Error code = " << hres << endl;

				m_Locator->Release();

				if (m_NeedUninitialize)
					CoUninitialize();

				m_InitSuccess = false;
			}
		}

		if (m_InitSuccess)
		{
			hres = CoSetProxyBlanket(m_Server, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);

			if (FAILED(hres))
			{
				qWarning() << "Could not set proxy blanket. Error code: " << hres;

				m_Server->Release();
				m_Locator->Release();

				if (m_NeedUninitialize)
					CoUninitialize();

				m_InitSuccess = false;
			}
		}
	}

	WmiQuery::~WmiQuery()
	{
		if (m_InitSuccess)
		{
			if (m_Server)
				m_Server->Release();

			if (m_Locator)
				m_Locator->Release();

			if (m_NeedUninitialize)
				CoUninitialize();
		}
	}

	WmiItems WmiQuery::exec(const QString& wql)
	{
		// qDebug() << "query: " << wql;

		WmiItems rval;

		if (m_Server == NULL)
			return rval;

		HRESULT hres;

		hres = m_Server->ExecQuery(bstr_t("WQL"), bstr_t(qPrintable(wql)), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &m_Enumerator);

		if (FAILED(hres))
			qWarning() << "Query with string '" << wql << "' has failed. Error code: " << hres;
		else
		{
			ULONG uReturn = 0;

			while (m_Enumerator)
			{
				IWbemClassObject* pclsObj;
				hres = m_Enumerator->Next(WBEM_INFINITE, 1, &pclsObj, &uReturn);

				if (uReturn == 0)
					break;

				rval.append(new WmiItem(pclsObj));
			}
		}

		// qDebug() << "num rows: " << rval.size();

		return rval;
	}

	QList<Device> wmiListHotPluggables()
	{
		QList<Device> rval;

		WmiQuery wmi;

		const WmiItems drives = wmi.exec("SELECT * FROM Win32_DiskDrive");

		foreach(const WmiItem* drive, drives)
		{
			// qDebug() << "drive: " << drive->property("Caption").toString();

			// NOTE: WQL cannot handle "... WIN32_DiskDrive.DeviceID = x", i.e. spaces around the
			// equal sign in the ASSOCIATORS OF {} subclause.
			const QString partitionQuery = QString("ASSOCIATORS OF {Win32_DiskDrive.DeviceID='%1'} WHERE AssocClass = Win32_DiskDriveToDiskPartition").arg(drive->property("DeviceID").toString());

			const WmiItems partitions = wmi.exec(partitionQuery);

			foreach(const WmiItem* partition, partitions)
			{
				// qDebug() << "partition: " << partition->property("Caption").toString();

				const QString logicalQuery = QString("ASSOCIATORS OF {Win32_DiskPartition.DeviceID='%1'} WHERE AssocClass = Win32_LogicalDiskToPartition").arg(partition->property("DeviceID").toString());

				const WmiItems logicals = wmi.exec(logicalQuery);

				foreach(const WmiItem* logical, logicals)
				{
					// qDebug() << "logical: " << logical->property("DeviceID") << " with drive type " << logical->property("DriveType");

					if (logical->property("DriveType").toInt() != 2)
						continue;

					// TODO: we leave the vendor empty and only fill in the product for the device; is
					// there a foolproof way to get both vendor and product for all kinds of devices?
					rval.append(Veneer::Device(logical->property("DeviceID").toString(), "drive-removable-media-usb-pendrive", QString(), drive->property("Caption").toString()));
				}

				qDeleteAll(logicals);
			}

			qDeleteAll(partitions);
		}

		qDeleteAll(drives);

		return rval;
	}
}

