#include "SystemParamsProvider_win.h"

struct IUnknown; // Workaround for "combaseapi.h(229): error C2187: syntax error: 'identifier' was unexpected here" when using /permissive-
#define _WIN32_DCOM
#include <comdef.h>
#include <Wbemidl.h>

#include <boost/algorithm/string/join.hpp>

#include <future>
#include <iostream>
#include <sstream>

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

namespace
{
    std::string convertWCSToMBS(const wchar_t* pstr, long wslen)
    {
        int len = ::WideCharToMultiByte(CP_ACP, 0, pstr, wslen, NULL, 0, NULL, NULL);

        std::string dblstr(len, '\0');
        len = ::WideCharToMultiByte(CP_ACP, 0 /* no flags */,
            pstr, wslen /* not necessary NULL-terminated */,
            &dblstr[0], len,
            NULL, NULL /* no default char */);

        return dblstr;
    }

    std::string convertBSTRToMBS(BSTR bstr)
    {
        int wslen = ::SysStringLen(bstr);
        return convertWCSToMBS((wchar_t*)bstr, wslen);
    }

    BSTR convertMBSToBSTR(const std::string& str)
    {
        int wslen = ::MultiByteToWideChar(CP_ACP, 0 /* no flags */,
            str.data(), static_cast<int>(str.length()),
            NULL, 0);

        BSTR wsdata = ::SysAllocStringLen(NULL, wslen);
        ::MultiByteToWideChar(CP_ACP, 0 /* no flags */,
            str.data(), static_cast<int>(str.length()),
            wsdata, wslen);
        return wsdata;
    }

    class QuerySink final : public IWbemObjectSink
    {
    public:
        typedef std::vector<std::string> Fields_type;
        typedef std::vector<std::vector<std::string>> Result_type;

    public:
        QuerySink(const std::string& wmiClass, const Fields_type& fields)
            : m_wmiClass(wmiClass)
            , m_fields(fields)
        {
            m_lRef = 0;
            m_bDone = false;
            InitializeCriticalSection(&m_threadLock);
        }
        ~QuerySink() {
            m_bDone = true;
            DeleteCriticalSection(&m_threadLock);
        }

    public:

        virtual ULONG STDMETHODCALLTYPE AddRef() override
        {
            return InterlockedIncrement(&m_lRef);
        }

        virtual ULONG STDMETHODCALLTYPE Release() override
        {
            LONG lRef = InterlockedDecrement(&m_lRef);
            if (lRef == 0)
                delete this;
            return lRef;
        }

        virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppv) override
        {
            if (riid == IID_IUnknown || riid == IID_IWbemObjectSink)
            {
                *ppv = (IWbemObjectSink *)this;
                AddRef();
                return WBEM_S_NO_ERROR;
            }
            else return E_NOINTERFACE;
        }

        virtual HRESULT STDMETHODCALLTYPE Indicate(
            LONG lObjectCount,
            IWbemClassObject __RPC_FAR *__RPC_FAR *apObjArray
        ) override
        {
            for (int i = 0; i < lObjectCount; i++)
            {
                std::vector<std::string> lineResult;
                lineResult.reserve(m_fields.size());
                for (const auto& field : m_fields)
                {
                    HRESULT hres = S_OK;
                    VARIANT varName;
                    hres = apObjArray[i]->Get(convertMBSToBSTR(field),
                        0, &varName, 0, 0);

                    if (FAILED(hres))
                    {
                        std::ostringstream ostr;
                        ostr << "Failed to get the data from the query. Error code = 0x"
                            << std::hex << hres;
                        std::cerr << ostr.str() << std::endl;
                        return WBEM_E_FAILED;       // Program has failed.
                    }

                    lineResult.push_back(convertBSTRToMBS(V_BSTR(&varName)));
                }
                
                m_result.push_back(lineResult);
            }

            return WBEM_S_NO_ERROR;
        }

        virtual HRESULT STDMETHODCALLTYPE SetStatus(
            /* [in] */ LONG lFlags,
            /* [in] */ HRESULT hResult,
            /* [in] */ BSTR strParam,
            /* [in] */ IWbemClassObject __RPC_FAR *pObjParam
        ) override
        {
            if (lFlags == WBEM_STATUS_COMPLETE)
            {
                // Call complete
                EnterCriticalSection(&m_threadLock);
                m_bDone = true;
                LeaveCriticalSection(&m_threadLock);
                m_resultPromise.set_value(std::move(m_result));
            }
            else if (lFlags == WBEM_STATUS_PROGRESS)
            {
                // Call in progress...
            }

            return WBEM_S_NO_ERROR;
        }

        std::future<Result_type> getFuture()
        {
            return m_resultPromise.get_future();
        }

        bool IsDone() const
        {
            bool done = true;

            EnterCriticalSection(&m_threadLock);
            done = m_bDone;
            LeaveCriticalSection(&m_threadLock);

            return done;
        }

    private:
        LONG m_lRef;
        bool m_bDone;
        mutable CRITICAL_SECTION m_threadLock; // for thread safety
        const std::string m_wmiClass;
        const Fields_type m_fields;
        std::vector<std::vector<std::string>> m_result;
        std::promise<Result_type> m_resultPromise;
    };

    void initializeCOM()
    {
        // We need to initialize COM for SystemParamsProvider   
        HRESULT hres = 0;

        // Initialize COM. ------------------------------------------
        hres = CoInitializeEx(0, COINIT_MULTITHREADED);
        if (FAILED(hres))
        {
            std::ostringstream ostr;
            ostr << "Failed to initialize COM library. Error code = 0x" << std::hex << hres;
            throw std::runtime_error(ostr.str());
        }

        // Set general COM security levels --------------------------
        hres = CoInitializeSecurity(NULL,
            -1,                          // COM authentication
            NULL,                        // Authentication services
            NULL,                        // Reserved
            RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication 
            RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation  
            NULL,                        // Authentication info
            EOAC_NONE,                   // Additional capabilities 
            NULL);                       // Reserved

        if (FAILED(hres) && hres != RPC_E_TOO_LATE) // Ignore error if the security has already been initialized
        {
            CoUninitialize();
            std::ostringstream ostr;
            ostr << "Failed to initialize security. Error code = 0x" << std::hex << hres;
            throw std::runtime_error(ostr.str());
        }
    }
}

namespace detail
{
    class SystemParamsProvider_winImpl final
    {
    public:
        SystemParamsProvider_winImpl()
        {
            HRESULT hres = 0;

            initializeCOM();

            // Obtain the initial locator to WMI -------------------------
            hres = CoCreateInstance(
                CLSID_WbemLocator,
                0,
                CLSCTX_INPROC_SERVER,
                IID_IWbemLocator, (LPVOID *)&m_pLoc);

            if (FAILED(hres))
            {
                CoUninitialize();
                std::ostringstream ostr;
                ostr << "Failed to create IWbemLocator object. Err code = 0x" << std::hex << hres;
                throw std::runtime_error(ostr.str());
            }

            // Connect to WMI through the IWbemLocator::ConnectServer method
            // Connect to the local root\cimv2 namespace
            // and obtain pointer pSvc to make IWbemServices calls.
            hres = m_pLoc->ConnectServer(_bstr_t(L"ROOT\\CIMV2"),
                NULL,
                NULL,
                0,
                NULL,
                0,
                0,
                &m_pSvc);

            if (FAILED(hres))
            {
                m_pLoc->Release();
                CoUninitialize();
                std::ostringstream ostr;
                ostr << "Could not connect. Error code = 0x" << std::hex << hres;
                throw std::runtime_error(ostr.str());
            }

            // Set security levels on the proxy -------------------------
            hres = CoSetProxyBlanket(m_pSvc,   // Indicates the proxy to set
                RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
                RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
                NULL,                        // Server principal name 
                RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx 
                RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
                NULL,                        // client identity
                EOAC_NONE);                  // proxy capabilities 

            if (FAILED(hres))
            {
                m_pSvc->Release();
                m_pLoc->Release();
                CoUninitialize();
                std::ostringstream ostr;
                ostr << "Could not set proxy blanket. Error code = 0x" << std::hex << hres;
                throw std::runtime_error(ostr.str());
            }

        }

        ~SystemParamsProvider_winImpl()
        {
            m_pSvc->Release();
            m_pLoc->Release();
            CoUninitialize();
        }

    public:
        SystemParams retrieveSystemParams()
        {   
            auto biosFuture = queryWmiInfo("Win32_BIOS", { "SerialNumber" });
            auto computerSystemProductFuture = queryWmiInfo("Win32_ComputerSystemProduct", { "UUID" });
            auto diskDriveFuture = queryWmiInfo("Win32_DiskDrive", { "SerialNumber" }, "DeviceId = \"\\\\\\\\.\\\\PHYSICALDRIVE0\"");
            auto osFuture = queryWmiInfo("Win32_OperatingSystem", { "SerialNumber" });
            auto nicFuture = queryWmiInfo("Win32_NetworkAdapter", { "MACAddress" }, "PhysicalAdapter = TRUE");
            auto mainboardFuture = queryWmiInfo("Win32_BaseBoard", { "SerialNumber" });

            SystemParams result;
            result[SystemParamTypes::biosSerialNum] = processBiosData(biosFuture.get());
            result[SystemParamTypes::computerUUID] = processComputerSystemProductData(computerSystemProductFuture.get());
            result[SystemParamTypes::diskSerialNum] = processDiskDrive(diskDriveFuture.get());
            result[SystemParamTypes::osId] = processOs(osFuture.get());
            result[SystemParamTypes::nicMac] = processNic(nicFuture.get());
            result[SystemParamTypes::mainboardSerialNum] = processMainBoard(mainboardFuture.get());

            return result;
        }
    
    private:
        typedef std::vector<std::vector<std::string>> QueryResult_type;

    private:
        std::future<QueryResult_type> queryWmiInfo(const std::string& wmiClass, const std::vector<std::string>& fieldNames, const std::string& condition = "")
        {
            HRESULT hres = 0;

            std::ostringstream queryStr;
            queryStr << "SELECT " << boost::algorithm::join(fieldNames, ",") << " FROM " << wmiClass;

            if (!condition.empty())
            {
                queryStr << " WHERE " << condition;
            }

            // Use the IWbemServices pointer to make requests of WMI ----
            QuerySink* pResponseSink = new QuerySink(wmiClass, fieldNames);
            pResponseSink->AddRef();
            auto resultFuture = pResponseSink->getFuture();
            hres = m_pSvc->ExecQueryAsync(bstr_t("WQL"),
                convertMBSToBSTR(queryStr.str()),
                WBEM_FLAG_BIDIRECTIONAL,
                NULL,
                pResponseSink);

            if (FAILED(hres))
            {
                pResponseSink->Release();
                std::ostringstream ostr;
                ostr << "Query for " << boost::algorithm::join(fieldNames, ", ") << " failed. Error code = 0x" << std::hex << hres;
                throw std::runtime_error(ostr.str());
            }

            return resultFuture;
        }

        static std::string processBiosData(const QueryResult_type& biosInfo)
        {
            if (biosInfo.size() > 0)
            {
                return biosInfo[0][0];
            }
            return {};
        }

        static std::string processComputerSystemProductData(const QueryResult_type& csInfo)
        {
            if (csInfo.size() > 0)
            {
                return csInfo[0][0];
            }
            return {};
        }

        static std::string processDiskDrive(const QueryResult_type& csInfo)
        {
            if (csInfo.size() > 0)
            {
                return csInfo[0][0];
            }
            return {};
        }

        static std::string processOs(const QueryResult_type& csInfo)
        {
            if (csInfo.size() > 0)
            {
                return csInfo[0][0];
            }
            return {};
        }

        static std::string processNic(const QueryResult_type& csInfo)
        {
            if (csInfo.size() > 0)
            {
                return csInfo[0][0];
            }
            return {};
        }

        static std::string processMainBoard(const QueryResult_type& csInfo)
        {
            if (csInfo.size() > 0)
            {
                return csInfo[0][0];
            }
            return {};
        }

    private:
        IWbemLocator* m_pLoc = nullptr;
        IWbemServices *m_pSvc = nullptr;
    };
}

SystemParamsProvider_win::SystemParamsProvider_win()
    : m_impl(std::make_unique<detail::SystemParamsProvider_winImpl>())
{
}

SystemParamsProvider_win::~SystemParamsProvider_win()
{
}

SystemParams SystemParamsProvider_win::retrieveSystemParams()
{
    return m_impl->retrieveSystemParams();
}