#include "wmiobject.h"

#include <Windows.h>
#include <slpublic.h>
#include <comdef.h>
#include <WbemIdl.h>

#include <QDebug>

#include "tool.h"
#include "variant.h"
#include "eventsink.h"
#include "disk.h"

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

namespace Win
{

class WmiObjectPrivate
{
public:
    WmiObjectPrivate(WmiObject *parent);
    ~WmiObjectPrivate();

    QList<QVariantMap> query(const QString &wql, const QStringList &props);
    QList<QVariantMap> monitor(const QString &wql);

    HRESULT hres;
    IWbemLocator *pLoc = NULL;
    IWbemServices *pSvc = NULL;

    WmiObject *q_ptr;
    Q_DECLARE_PUBLIC(WmiObject)
};

WmiObject::WmiObject(QObject *parent) :
    QObject(parent), dd_ptr(new WmiObjectPrivate(this))
{
}

WmiObject::~WmiObject()
{
    qDebug() << "delete wmiobject" << this;
}

QList<DiskInfo> WmiObject::QueryXPDiskInfo()
{
    QList<DiskInfo> diskinfos;
    auto objs = Query("SELECT * FROM Win32_DiskDrive", {"DeviceID", "PNPDeviceID"});
    int driveID = 0;
    for (auto &obj : objs) {
        int currentDriveID = driveID;
        QString DeviceID = obj.value("DeviceID").toString().trimmed();
        if (DeviceID.toUpper().contains("PHYSICALDRIVE")) {
            currentDriveID = DeviceID.toUpper().split("PHYSICALDRIVE").last().toInt();
        }
        DiskInfo info= getHardDriveSerialNumber(currentDriveID);
        info.name = DeviceID;
//        QString PNPDeviceID = obj.value("PNPDeviceID").toString().trimmed();
//        QString serial = PNPDeviceID.split("\\").last();
//        //42566538353834643432612D3338353565392037
//        qDebug() << info.name << PNPDeviceID << serial;
//        QString decode;
//        for (int i = 0; i < serial.length();) {
//            bool ok;
//            auto first = char(QString("%1%2").arg(serial.at(i)).arg(serial.at(i + 1)).toInt(&ok, 16));
//            auto second = char(QString("%1%2").arg(serial.at(i + 2)).arg(serial.at(i + 3)).toInt(&ok, 16));
//            decode += second;
//            decode += first;
//            i = i + 4;
//        }
//        decode = decode.trimmed();
//        info.serial_no = decode;
//        qDebug() << decode;
        diskinfos.push_back(info);
        ++driveID;
    }
    return diskinfos;
}

QList<DiskInfo> WmiObject::QueryDiskInfo()
{
    QList<DiskInfo> diskinfos;
    auto objs = Query("SELECT * FROM Win32_PhysicalMedia", {"SerialNumber", "Tag"});
    for (auto &obj : objs) {
        DiskInfo info;
        info.name = obj.value("Tag").toString().trimmed();
        info.serial_no = obj.value("SerialNumber").toString().trimmed();
        diskinfos.push_back(info);
    }
    return diskinfos;
}

QList<QVariantMap> WmiObject::Query(const QString &wql, const QStringList &props)
{
    Q_D(WmiObject);
    return d->query(wql, props);
}

QList<QVariantMap> WmiObject::Monitor(const QString &wql)
{
    Q_D(WmiObject);
    return d->monitor(wql);
}

WmiObjectPrivate::WmiObjectPrivate(WmiObject *parent) : q_ptr(parent)
{
    hres = CoCreateInstance(
               CLSID_WbemLocator,
               0,
               CLSCTX_INPROC_SERVER,
               IID_IWbemLocator, (LPVOID *) &pLoc);

    if (FAILED(hres)) {
        qDebug() << "Failed to create IWbemLocator object."
                 << " Err code = 0x"
                 << SUCCEEDED(hres) << Tool::qHexStr1(hres);
        return ;                 // Program has failed.
    }

    // Step 4: -----------------------------------------------------
    // Connect to WMI through the IWbemLocator::ConnectServer method


    // 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"), // Object path of WMI namespace
               NULL,                    // User name. NULL = current user
               NULL,                    // User password. NULL = current
               0,                       // Locale. NULL indicates current
               NULL,                    // Security flags.
               0,                       // Authority (for example, Kerberos)
               0,                       // Context object
               &pSvc                    // pointer to IWbemServices proxy
           );

    if (FAILED(hres)) {
        qDebug() << "Could not connect. Error code = 0x"
                 << SUCCEEDED(hres) << Tool::qHexStr1(hres);
        pLoc->Release();
        return;
    }

    qDebug() << "Connected to ROOT\\CIMV2 WMI namespace";

    // Step 5: --------------------------------------------------
    // 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
               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)) {
        qDebug() << "Could not set proxy blanket. Error code = 0x"
                 << SUCCEEDED(hres) << Tool::qHexStr1(hres);
        pSvc->Release();
        pLoc->Release();
        return;
    }
}

WmiObjectPrivate::~WmiObjectPrivate()
{
    // Cleanup
    // ========
    //  TODO: clean
    //  hres = pSvc->CancelAsyncCall(pStubSink);
    qDebug() << "WMI: Monitor End -------------";

    qDebug() << "delete WmiObjectPrivate" << this;
    pSvc->Release();
    pLoc->Release();
    CoUninitialize();
}

QList<QVariantMap> WmiObjectPrivate::query(const QString &wql, const QStringList &props)
{
    qDebug() << "WMI: Query Begin -------------";
    qDebug() << "WMI: Query " << wql << props;
    QList<QVariantMap> result;
    // Step 6: --------------------------------------------------
    // Use the IWbemServices pointer to make requests of WMI ----

    // bstr_t("SELECT * FROM Win32_OperatingSystem"),
    // For example, get the name of the operating system
    IEnumWbemClassObject *pEnumerator = NULL;
    hres = pSvc->ExecQuery(
               bstr_t("WQL"),
               bstr_t(wql.toStdWString().c_str()),
               WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
               NULL,
               &pEnumerator);

    if (FAILED(hres)) {
        qCritical() << "Query for operating system name failed."
                    << " Error code = 0x"
                    << SUCCEEDED(hres) << Tool::qHexStr1(hres);
        return result;
    }

    // Step 7: -------------------------------------------------
    // Get the data from the query in step 6 -------------------

    IWbemClassObject    *pclsObj    = NULL;
    ULONG               uReturn     = 0;

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

//        qDebug() << "WMI: Dump object Begin -------------";
        QVariantMap values;
        VARIANT vtProp;
        for (QString prop : props) {
            hr = pclsObj->Get(prop.toStdWString().c_str(), 0, &vtProp, 0, 0);
//            Win::dumpProp(vtProp, prop.toStdWString().c_str());
            values.insert(prop, Win::qVariant(vtProp));
            VariantClear(&vtProp);
        }
        if (props.isEmpty()) {
            qDebug() << "dump" << pclsObj;
            Win::dumpWmiObj(pclsObj);
        }
        result.append(values);
        pclsObj->Release();
//        qDebug() << "WMI: Dump object End -------------";
    }
    pEnumerator->Release();

    qDebug() << "WMI: Query End -------------";
    return result;
}

QList<QVariantMap> WmiObjectPrivate::monitor(const QString &wql)
{
    Q_Q(WmiObject);

    QList<QVariantMap> result;
    // Step 6: -------------------------------------------------
    // Receive event notifications -----------------------------

    // Use an unsecured apartment for security
    IUnsecuredApartment *pUnsecApp = NULL;

    hres = CoCreateInstance(CLSID_UnsecuredApartment, NULL,
                            CLSCTX_LOCAL_SERVER, IID_IUnsecuredApartment,
                            (void **)&pUnsecApp);

    EventSink *pSink = new EventSink(q);
    pSink->AddRef();

    IUnknown *pStubUnk = NULL;
    pUnsecApp->CreateObjectStub(pSink, &pStubUnk);

    IWbemObjectSink *pStubSink = NULL;
    pStubUnk->QueryInterface(IID_IWbemObjectSink,
                             (void **) &pStubSink);

    // The ExecNotificationQueryAsync method will call
    // The EventQuery::Indicate method when an event occurs
    hres = pSvc->ExecNotificationQueryAsync(
               _bstr_t("WQL"),
               _bstr_t("SELECT * "
                       "FROM __InstanceOperationEvent WITHIN 1 "
                       "WHERE TargetInstance ISA 'Win32_Process'"),
               WBEM_FLAG_SEND_STATUS,
               NULL,
               pStubSink);

    // Check for errors.
    if (FAILED(hres)) {
        printf("ExecNotificationQueryAsync failed "
               "with = 0x%X\n", hres);
        pUnsecApp->Release();
        pStubUnk->Release();
        pSink->Release();
        pStubSink->Release();
        return result;
    }

    qDebug() << "WMI: Monitor Start -------------";


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

//    pUnsecApp->Release();
//    pStubUnk->Release();
//    pSink->Release();
//    pStubSink->Release();

    return result;
}

WmiObject *GetWmiObject()
{
    static WmiObject obj;
    return &obj;
}

}
