#include "Config.h"

#include <windows.h>
#include <vector>

namespace
{
    DWORD ReadDword(HKEY key, const wchar_t* name, DWORD defaultValue)
    {
        DWORD value = defaultValue;
        DWORD size = sizeof(value);
        DWORD type = 0;
        if (RegQueryValueExW(key, name, nullptr, &type, reinterpret_cast<BYTE*>(&value), &size) != ERROR_SUCCESS ||
            type != REG_DWORD)
        {
            return defaultValue;
        }
        return value;
    }

    std::wstring ReadStringRaw(HKEY key, const wchar_t* name)
    {
        DWORD type = 0;
        DWORD size = 0;
        if (RegQueryValueExW(key, name, nullptr, &type, nullptr, &size) != ERROR_SUCCESS ||
            (type != REG_SZ && type != REG_EXPAND_SZ) || size == 0)
        {
            return {};
        }

        std::vector<wchar_t> buffer((size / sizeof(wchar_t)) + 1, L'\0');
        if (RegQueryValueExW(key, name, nullptr, &type, reinterpret_cast<BYTE*>(buffer.data()), &size) != ERROR_SUCCESS)
        {
            return {};
        }
        return std::wstring(buffer.data());
    }

    bool WriteDword(HKEY key, const wchar_t* name, bool value, std::wstring& error)
    {
        DWORD data = value ? 1u : 0u;
        LONG result = RegSetValueExW(key, name, 0, REG_DWORD, reinterpret_cast<const BYTE*>(&data), sizeof(data));
        if (result != ERROR_SUCCESS)
        {
            error = L"Registry-Wert konnte nicht geschrieben werden: " + std::wstring(name) + L" (Fehler " + std::to_wstring(result) + L")";
            return false;
        }
        return true;
    }

    bool WriteString(HKEY key, const wchar_t* name, const std::wstring& value, DWORD type, std::wstring& error)
    {
        const DWORD bytes = static_cast<DWORD>((value.size() + 1) * sizeof(wchar_t));
        LONG result = RegSetValueExW(
            key,
            name,
            0,
            type,
            reinterpret_cast<const BYTE*>(value.c_str()),
            bytes);
        if (result != ERROR_SUCCESS)
        {
            error = L"Registry-Wert konnte nicht geschrieben werden: " + std::wstring(name) + L" (Fehler " + std::to_wstring(result) + L")";
            return false;
        }
        return true;
    }
}

namespace config
{
    std::wstring ExpandEnvironment(const std::wstring& value)
    {
        if (value.empty())
        {
            return value;
        }

        DWORD needed = ExpandEnvironmentStringsW(value.c_str(), nullptr, 0);
        if (needed == 0)
        {
            return value;
        }

        std::vector<wchar_t> buffer(needed, L'\0');
        if (ExpandEnvironmentStringsW(value.c_str(), buffer.data(), needed) == 0)
        {
            return value;
        }
        return std::wstring(buffer.data());
    }

    HookConfig Load()
    {
        HookConfig value;
        HKEY key = nullptr;
        if (RegOpenKeyExW(HKEY_CURRENT_USER, kConfigKey, 0, KEY_QUERY_VALUE, &key) != ERROR_SUCCESS)
        {
            return value;
        }

        value.enabled = ReadDword(key, L"Enabled", 1) != 0;
        value.enableLogging = ReadDword(key, L"EnableLogging", 1) != 0;

        value.scriptPath = ReadStringRaw(key, L"ScriptPath");
        value.arguments = ReadStringRaw(key, L"Arguments");
        value.workingDirectory = ReadStringRaw(key, L"WorkingDirectory");
        value.startOnConnect = ReadDword(key, L"StartOnConnect", 1) != 0;
        value.stopOnDisconnect = ReadDword(key, L"StopOnDisconnect", 0) != 0;
        value.hidden = ReadDword(key, L"Hidden", 1) != 0;
        value.preventDuplicates = ReadDword(key, L"PreventDuplicates", 1) != 0;

        value.triggerEnabled = ReadDword(key, L"TriggerEnabled", 0) != 0;
        value.triggerProgramPath = ReadStringRaw(key, L"TriggerProgramPath");
        value.triggerArguments = ReadStringRaw(key, L"TriggerArguments");
        value.triggerWorkingDirectory = ReadStringRaw(key, L"TriggerWorkingDirectory");
        value.triggerHidden = ReadDword(key, L"TriggerHidden", 1) != 0;
        value.triggerPreventDuplicates = ReadDword(key, L"TriggerPreventDuplicates", 1) != 0;

        RegCloseKey(key);
        return value;
    }

    bool Save(const HookConfig& value, std::wstring& error)
    {
        HKEY key = nullptr;
        DWORD disposition = 0;
        LONG result = RegCreateKeyExW(
            HKEY_CURRENT_USER,
            kConfigKey,
            0,
            nullptr,
            REG_OPTION_NON_VOLATILE,
            KEY_SET_VALUE,
            nullptr,
            &key,
            &disposition);
        if (result != ERROR_SUCCESS)
        {
            error = L"Konfigurationsschlüssel konnte nicht erstellt werden (Fehler " + std::to_wstring(result) + L").";
            return false;
        }

        bool ok =
            WriteDword(key, L"Enabled", value.enabled, error) &&
            WriteDword(key, L"EnableLogging", value.enableLogging, error) &&
            WriteString(key, L"ScriptPath", value.scriptPath, REG_EXPAND_SZ, error) &&
            WriteString(key, L"Arguments", value.arguments, REG_SZ, error) &&
            WriteString(key, L"WorkingDirectory", value.workingDirectory, REG_EXPAND_SZ, error) &&
            WriteDword(key, L"StartOnConnect", value.startOnConnect, error) &&
            WriteDword(key, L"StopOnDisconnect", value.stopOnDisconnect, error) &&
            WriteDword(key, L"Hidden", value.hidden, error) &&
            WriteDword(key, L"PreventDuplicates", value.preventDuplicates, error) &&
            WriteDword(key, L"TriggerEnabled", value.triggerEnabled, error) &&
            WriteString(key, L"TriggerProgramPath", value.triggerProgramPath, REG_EXPAND_SZ, error) &&
            WriteString(key, L"TriggerArguments", value.triggerArguments, REG_SZ, error) &&
            WriteString(key, L"TriggerWorkingDirectory", value.triggerWorkingDirectory, REG_EXPAND_SZ, error) &&
            WriteDword(key, L"TriggerHidden", value.triggerHidden, error) &&
            WriteDword(key, L"TriggerPreventDuplicates", value.triggerPreventDuplicates, error);

        RegCloseKey(key);
        return ok;
    }

    std::wstring GetRegisteredPluginPath()
    {
        HKEY key = nullptr;
        if (RegOpenKeyExW(HKEY_CURRENT_USER, kAddInKey, 0, KEY_QUERY_VALUE, &key) != ERROR_SUCCESS)
        {
            return {};
        }
        std::wstring value = ReadStringRaw(key, L"Name");
        RegCloseKey(key);
        return value;
    }

    bool RegisterPlugin(const std::wstring& dllPath, std::wstring& error)
    {
        HKEY key = nullptr;
        DWORD disposition = 0;
        LONG result = RegCreateKeyExW(
            HKEY_CURRENT_USER,
            kAddInKey,
            0,
            nullptr,
            REG_OPTION_NON_VOLATILE,
            KEY_SET_VALUE,
            nullptr,
            &key,
            &disposition);
        if (result != ERROR_SUCCESS)
        {
            error = L"MSTSC AddIn-Schlüssel konnte nicht erstellt werden (Fehler " + std::to_wstring(result) + L").";
            return false;
        }

        bool ok = WriteString(key, L"Name", dllPath, REG_SZ, error);
        RegCloseKey(key);
        return ok;
    }

    bool UnregisterPlugin(std::wstring& error)
    {
        LONG result = RegDeleteTreeW(HKEY_CURRENT_USER, kAddInKey);
        if (result == ERROR_SUCCESS || result == ERROR_FILE_NOT_FOUND || result == ERROR_PATH_NOT_FOUND)
        {
            return true;
        }
        error = L"MSTSC-Plugin-Registrierung konnte nicht entfernt werden (Fehler " + std::to_wstring(result) + L").";
        return false;
    }
}