#include "registry.h"

#include <Windows.h>

#include <QDebug>

#define MAX_KEY_LENGTH 255
#define MAX_VALUE_NAME 16383

extern BOOL IsWow64();

class RegistryPrivate
{
public:
    RegistryPrivate(Registry *parent) : q_ptr(parent) {}

    void open()
    {
        DWORD mask = isWow64 ? KEY_WOW64_64KEY : KEY_WOW64_32KEY;
//        qDebug() << "open" << isWow64 << mask << keyPath << hKey;
        LONG retCode = RegOpenKeyEx(hKey,
                                    keyPath.toStdWString().c_str(),
                                    REG_OPTION_NON_VOLATILE,
                                    KEY_READ | mask,
                                    &hRegistry);

        if (ERROR_SUCCESS != retCode) {
            qDebug() << "RegOpenKeyEx Failed" << retCode << hKey << keyPath;
        }
    }

    QStringList childGroups() const
    {
        QStringList groups;
        TCHAR    achKey[MAX_KEY_LENGTH];   // buffer for subkey name
        DWORD    cbName;                   // size of name string
        TCHAR    achClass[MAX_PATH] = TEXT("");  // buffer for class name
        DWORD    cchClassName = MAX_PATH;  // size of class string
        DWORD    cSubKeys = 0;             // number of subkeys
        DWORD    cbMaxSubKey;              // longest subkey size
        DWORD    cchMaxClass;              // longest class string
        DWORD    cValues;              // number of values for key
        DWORD    cchMaxValue;          // longest value name
        DWORD    cbMaxValueData;       // longest value data
        DWORD    cbSecurityDescriptor; // size of security descriptor
        FILETIME ftLastWriteTime;      // last write time

        DWORD i, retCode;
        // Get the class name and the value count.
        retCode = RegQueryInfoKey(
                      hRegistry,                    // key handle
                      achClass,                // buffer for class name
                      &cchClassName,           // size of class string
                      NULL,                    // reserved
                      &cSubKeys,               // number of subkeys
                      &cbMaxSubKey,            // longest subkey size
                      &cchMaxClass,            // longest class string
                      &cValues,                // number of values for this key
                      &cchMaxValue,            // longest value name
                      &cbMaxValueData,         // longest value data
                      &cbSecurityDescriptor,   // security descriptor
                      &ftLastWriteTime);       // last write time

        if (cSubKeys) {
            for (i = 0; i < cSubKeys; i++) {
                cbName = MAX_KEY_LENGTH;
                retCode = RegEnumKeyEx(hRegistry, i,
                                       achKey,
                                       &cbName,
                                       NULL,
                                       NULL,
                                       NULL,
                                       &ftLastWriteTime);
                if (retCode == ERROR_SUCCESS) {
//                    qDebug() << i + 1 << QString::fromStdWString(achKey);
                    groups << QString::fromStdWString(achKey);
                }
            }
        }
        return groups;
    }

    QStringList keys()const
    {
        QStringList keyList;
        TCHAR    achClass[MAX_PATH] = TEXT("");  // buffer for class name
        DWORD    cchClassName = MAX_PATH;  // size of class string
        DWORD    cSubKeys = 0;             // number of subkeys
        DWORD    cbMaxSubKey;              // longest subkey size
        DWORD    cchMaxClass;              // longest class string
        DWORD    cValues;              // number of values for key
        DWORD    cchMaxValue;          // longest value name
        DWORD    cbMaxValueData;       // longest value data
        DWORD    cbSecurityDescriptor; // size of security descriptor
        FILETIME ftLastWriteTime;      // last write time

        DWORD i, retCode;

        TCHAR  achValue[MAX_VALUE_NAME];
        DWORD cchValue = MAX_VALUE_NAME;

        // Get the class name and the value count.
        retCode = RegQueryInfoKey(
                      hRegistry,                    // key handle
                      achClass,                // buffer for class name
                      &cchClassName,           // size of class string
                      NULL,                    // reserved
                      &cSubKeys,               // number of subkeys
                      &cbMaxSubKey,            // longest subkey size
                      &cchMaxClass,            // longest class string
                      &cValues,                // number of values for this key
                      &cchMaxValue,            // longest value name
                      &cbMaxValueData,         // longest value data
                      &cbSecurityDescriptor,   // security descriptor
                      &ftLastWriteTime);       // last write time

        // Enumerate the key values.
        if (cValues) {
            for (i = 0, retCode = ERROR_SUCCESS; i < cValues; i++) {
                cchValue = MAX_VALUE_NAME;
                achValue[0] = '\0';
                retCode = RegEnumValue(hRegistry, i,
                                       achValue,
                                       &cchValue,
                                       NULL,
                                       NULL,
                                       NULL,
                                       NULL);

                if (retCode == ERROR_SUCCESS) {
//                    qDebug() << i + 1 << QString::fromStdWString(achValue);
                    keyList << QString::fromStdWString(achValue);
                }
            }
        }
        return keyList;
    }

    QVariant value(const QString &key) const
    {
        QVariant value;
        DWORD dwLength = 0;
        DWORD dwType = 0;
        auto retCode = RegQueryValueEx(hRegistry, key.toStdWString().c_str(),
                                       NULL, NULL, NULL, &dwLength);

        if (dwLength <= 0) {
//            qDebug() << "RegQueryValueEx " << key << "Failed" << retCode << "DataLength" << dwLength;
            return value;
        }

        //Allocate Memory & Initialization
        auto data = new byte[dwLength];
        memset(data, 0, dwLength);

        retCode = RegQueryValueEx(hRegistry, key.toStdWString().c_str(),
                                  NULL, &dwType, data, &dwLength);

        if (retCode != ERROR_SUCCESS) {
            qDebug() << "RegQueryValueEx " << key << "Failed" << retCode << "DataLength" << dwLength;
        }

        switch (dwType) {
        case 1: // string
            value = QString::fromWCharArray((wchar_t *)data, dwLength / 2 - 1);
            break;
        case 3: // bytearray
            value = QByteArray::fromRawData((char *)data, dwLength);
            break;
        case 4: { // DWORD
            DWORD dwData = *(DWORD *)data;
            value = QVariant::fromValue(dwData);
            break;
        }
        default:
            break;
        }
//        qDebug() << key << dwLength << dwType << value;

        //Release Memory
        if (data) {
            delete []data;
        }
        return value;
    }

    bool    isWow64 = true;
    HKEY    hKey;
    HKEY    hRegistry;
    QString keyPath;

    Registry *q_ptr;
    Q_DECLARE_PUBLIC(Registry)
};


Registry::Registry(QObject *parent):
    QObject(parent), dd_ptr(new RegistryPrivate(this))
{
}


Registry::Registry(HKEY key,
                   const QString &regKey,
                   bool autoArch,
                   bool useWOW64,
                   QObject *parent):
    QObject(parent), dd_ptr(new RegistryPrivate(this))
{
    open(key, regKey, autoArch, useWOW64);
}

Registry::~Registry()
{
    //Close Registry
    Q_D(Registry);
    RegCloseKey(d->hRegistry);
}

void Registry::open(HKEY key, const QString &keyPath, bool autoArch, bool useWOW64)
{
    Q_D(Registry);

    static bool isSysWOW64 = IsWow64();
    d->keyPath = keyPath;
    d->hKey = key;
    if (autoArch) {
        d->isWow64 = isSysWOW64;
    } else {
        d->isWow64 = useWOW64;
    }

    d->open();
}

QStringList Registry::childGroups() const
{
    Q_D(const Registry);
    return d->childGroups();
}

QStringList Registry::keys() const
{
    Q_D(const Registry);
    return d->keys();
}

QVariant Registry::value(const QString &key) const
{
    Q_D(const Registry);
    return d->value(key);
}

Registry *Registry::subKey(const QString &keyPath)
{
    Q_D(const Registry);
    auto subKeyPath = d->keyPath.isEmpty() ?
                      keyPath : QString("%1\\%2").arg(d->keyPath).arg(keyPath);
    return new Registry(d->hKey,
                        subKeyPath,
                        false,
                        d->isWow64,
                        this);
}
