#include "wmi.hpp"
#include "util/windows/com.hpp"
#include "util/windows/unicode.hpp"

#include <synchapi.h>
#include <wchar.h>
#include <math.h>

namespace
{
    // Provide our bstr_t to avoid libstdc++ dependency
    struct bstr_t
    {
        explicit bstr_t(const wchar_t* str) noexcept: _bstr(SysAllocString(str)) {}
        ~bstr_t(void) noexcept { SysFreeString(_bstr); }
        explicit operator const wchar_t*(void) const noexcept { return _bstr; }
        operator BSTR(void) const noexcept { return _bstr; }

        private:
            BSTR _bstr;
    };
}

static BOOL CALLBACK InitHandleFunction(PINIT_ONCE, PVOID lpParameter, PVOID* lpContext)
{
    HRESULT hres;

    // Obtain the initial locator to WMI
    IWbemLocator* pLoc = nullptr;
    hres = CoCreateInstance(
        CLSID_WbemLocator,
        nullptr,
        CLSCTX_INPROC_SERVER,
        IID_IWbemLocator,
        (LPVOID*) &pLoc);

    if (FAILED(hres))
    {
        *((const char**)lpContext) = "Failed to create IWbemLocator object";
        return FALSE;
    }

    // Connect to WMI through the IWbemLocator::ConnectServer method
    IWbemServices* pSvc = nullptr;

    // Connect to the root\cimv2 namespace with
    // the current user and obtain pointer pSvc
    // to make IWbemServices calls.
    hres = pLoc->ConnectServer(
        bstr_t((const wchar_t*) lpParameter), // Object path of WMI namespace
        nullptr,                              // User name. nullptr = current user
        nullptr,                              // User password. nullptr = current
        0,                                    // Locale. nullptr indicates current
        0,                                    // Security flags.
        0,                                    // Authority (for example, Kerberos)
        0,                                    // Context object
        &pSvc                                 // pointer to IWbemServices proxy
    );
    pLoc->Release();
    pLoc = nullptr;

    if (FAILED(hres))
    {
        *((const char**)lpContext) = "Could not connect WMI server";
        return FALSE;
    }

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

    if (FAILED(hres))
    {
        pSvc->Release();
        *((const char**)lpContext) = "Could not set proxy blanket";
        return FALSE;
    }

    *((IWbemServices**)lpContext) = pSvc;
    return TRUE;
}

FFWmiQuery::FFWmiQuery(const wchar_t* queryStr, FFstrbuf* error, FFWmiNamespace wmiNs)
    : pEnumerator(nullptr)
{
    const char* context = ffInitCom();
    if (context)
    {
        if (error)
            ffStrbufAppendS(error, context);
        return;
    }

    static INIT_ONCE s_InitOnce[(int) FFWmiNamespace::LAST] = {};
    if (InitOnceExecuteOnce(
        &s_InitOnce[(int)wmiNs],
        &InitHandleFunction,
        (PVOID) (wmiNs == FFWmiNamespace::CIMV2 ? L"ROOT\\CIMV2" : L"ROOT\\WMI"),
        (void**)&context) == FALSE
    ) {
        if(error)
            ffStrbufAppendS(error, context);
        return;
    }

    // Use the IWbemServices pointer to make requests of WMI
    HRESULT hres = ((IWbemServices*)context)->ExecQuery(
        bstr_t(L"WQL"),
        bstr_t(queryStr),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        nullptr,
        &pEnumerator);

    if (FAILED(hres))
    {
        if(error)
            ffStrbufAppendF(error, "Query for '%ls' failed. Error code = 0x%lX", queryStr, hres);
    }
}

bool FFWmiRecord::getString(const wchar_t* key, FFstrbuf* strbuf)
{
    bool result = true;

    FFWmiVariant vtProp;

    CIMTYPE type;
    if(FAILED(obj->Get(key, 0, &vtProp, &type, nullptr)) || vtProp.vt != VT_BSTR)
    {
        result = false;
    }
    else
    {
        switch(vtProp.vt)
        {
            case VT_BSTR:
                if(type == CIM_DATETIME)
                {
                    ISWbemDateTime *pDateTime;
                    BSTR dateStr;
                    if(FAILED(CoCreateInstance(__uuidof(SWbemDateTime), 0, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pDateTime))))
                        result = false;
                    else if(FAILED(pDateTime->put_Value(vtProp.bstrVal)))
                        result = false;
                    else if(FAILED(pDateTime->GetFileTime(VARIANT_TRUE, &dateStr)))
                        result = false;
                    else
                        ffStrbufSetNWS(strbuf, SysStringLen(dateStr), dateStr);
                }
                else
                {
                    ffStrbufSetNWS(strbuf, SysStringLen(vtProp.bstrVal), vtProp.bstrVal);
                }
                break;

            case VT_LPSTR:
                ffStrbufAppendS(strbuf, vtProp.pcVal);
                break;

            case VT_LPWSTR:
            default:
                ffStrbufSetWS(strbuf, vtProp.bstrVal);
                break;
        }
    }
    return result;
}

bool FFWmiRecord::getSigned(const wchar_t* key, int64_t* integer)
{
    bool result = true;

    FFWmiVariant vtProp;

    CIMTYPE type;
    if(FAILED(obj->Get(key, 0, &vtProp, &type, nullptr)))
    {
        result = false;
    }
    else
    {
        switch(vtProp.vt)
        {
            case VT_BSTR: *integer = wcstoll(vtProp.bstrVal, nullptr, 10); break;
            case VT_I1: *integer = vtProp.cVal; break;
            case VT_I2: *integer = vtProp.iVal; break;
            case VT_INT:
            case VT_I4: *integer = vtProp.intVal; break;
            case VT_I8: *integer = vtProp.llVal; break;
            case VT_UI1: *integer = (int64_t)vtProp.bVal; break;
            case VT_UI2: *integer = (int64_t)vtProp.uiVal; break;
            case VT_UINT:
            case VT_UI4: *integer = (int64_t)vtProp.uintVal; break;
            case VT_UI8: *integer = (int64_t)vtProp.ullVal; break;
            case VT_BOOL: *integer = vtProp.boolVal != VARIANT_FALSE; break;
            default: *integer = 0; result = false;
        }
    }
    return result;
}

bool FFWmiRecord::getUnsigned(const wchar_t* key, uint64_t* integer)
{
    bool result = true;

    FFWmiVariant vtProp;

    if(FAILED(obj->Get(key, 0, &vtProp, nullptr, nullptr)))
    {
        result = false;
    }
    else
    {
        switch(vtProp.vt)
        {
            case VT_BSTR: *integer = wcstoull(vtProp.bstrVal, nullptr, 10); break;
            case VT_I1: *integer = (uint64_t)vtProp.cVal; break;
            case VT_I2: *integer = (uint64_t)vtProp.iVal; break;
            case VT_INT:
            case VT_I4: *integer = (uint64_t)vtProp.intVal; break;
            case VT_I8: *integer = (uint64_t)vtProp.llVal; break;
            case VT_UI1: *integer = vtProp.bVal; break;
            case VT_UI2: *integer = vtProp.uiVal; break;
            case VT_UINT:
            case VT_UI4: *integer = vtProp.uintVal; break;
            case VT_UI8: *integer = vtProp.ullVal; break;
            case VT_BOOL: *integer = vtProp.boolVal != VARIANT_FALSE; break;
            default: *integer = 0; result = false;
        }
    }
    return result;
}

bool FFWmiRecord::getReal(const wchar_t* key, double* real)
{
    bool result = true;

    FFWmiVariant vtProp;

    if(FAILED(obj->Get(key, 0, &vtProp, nullptr, nullptr)))
    {
        result = false;
    }
    else
    {
        switch(vtProp.vt)
        {
            case VT_BSTR: *real = wcstod(vtProp.bstrVal, nullptr); break;
            case VT_I1: *real = vtProp.cVal; break;
            case VT_I2: *real = vtProp.iVal; break;
            case VT_INT:
            case VT_I4: *real = vtProp.intVal; break;
            case VT_I8: *real = (double)vtProp.llVal; break;
            case VT_UI1: *real = vtProp.bVal; break;
            case VT_UI2: *real = vtProp.uiVal; break;
            case VT_UINT:
            case VT_UI4: *real = vtProp.uintVal; break;
            case VT_UI8: *real = (double)vtProp.ullVal; break;
            case VT_R4: *real = vtProp.fltVal; break;
            case VT_R8: *real = vtProp.dblVal; break;
            case VT_BOOL: *real = vtProp.boolVal != VARIANT_FALSE; break;
            default: *real = NAN; result = false;
        }
    }
    return result;
}