#pragma once

#include "GameCatalogExtensionVersion.h"

#include <AclAPI.h>
#include <Windows.h>
#include <combaseapi.h>
#include <rpc.h>
#include <sddl.h>
#include <wil\resource.h>
#include <wil\result_macros.h>

#include <sstream>
#include <string>

// Current API Version
// GCE_VERSION will be written at build time into GameCatalogExtensionVersion.h
// To modify this, you must also change the behavior of build-config/nuspecs/build-nupkg.ps1
inline const wchar_t* c_apiVersion{ GCE_VERSION };

inline const uint32_t c_maxApiVersionLength = 30;

// Registry Paths
inline const wchar_t* c_baseXboxExtensionsClsidPath{ L"SOFTWARE\\Microsoft\\Xbox\\GamingApp\\Extensions\\CLSID\\" };
inline const wchar_t* c_baseXboxExtensionsDataPath{ L"SOFTWARE\\Microsoft\\Xbox\\GamingApp\\Extensions\\Data\\" };
inline const wchar_t* c_baseClassesRootClsidPath{ L"CLSID\\" };
inline const wchar_t* c_baseClassesRootAppIdPath{ L"AppID\\" };

/*
* This security descriptor grants the following rights:
* 
* COM_RIGHTS_EXECUTE 1
* COM_RIGHTS_EXECUTE_LOCAL 2
* COM_RIGHTS_ACTIVATE_LOCAL 8
* (1 | 2 | 8 == 11)
*
* To the following SIDs:
* WD: Everyone. Required for access with packaged apps, because packaged apps require the current user to have access as well.
* AC: All application packages. To be removed in a future version of the API, only left in to aid with testing for third party publishers until the Xbox App is working reliably.
* 
* https://learn.microsoft.com/en-us/windows/win32/secauthz/security-descriptor-definition-language
* https://learn.microsoft.com/en-us/windows/win32/secauthz/ace-strings
* https://learn.microsoft.com/en-us/windows/win32/secauthz/sid-strings
*/
inline const wchar_t* c_comStringSecurityDescriptor{
    L"O:PSG:BUD:"
    L"(A;;11;;;WD)"
    L"(A;;11;;;AC)"
    L"(A;;11;;;S-1-15-2-1723189366-2159580849-2248400763-1481059666-1951766778-2756563051-3565589001)"
    L"S:(ML;;NX;;;LW)"
};

namespace Microsoft::Gaming::XboxApp::Extensions
{

namespace Impl
{
    inline HRESULT CreateXboxLocalMachineKey(
        std::wstring_view extensionId,
        std::wstring_view name,
        std::wstring_view clsId);
    inline HRESULT CreateXboxCurrentUserKey(std::wstring_view extensionId);
    inline HRESULT SetUserWriteAccessToRegistryObject(HKEY);
} // namespace Impl

enum class ExtensionInstallationState
{
    Uninstalled = 0,
    Installed = 1,
    Updating = 2,
    Uninitialized = 3,
};

enum class ExtensionType
{
    Server = 0,
};

enum class Architecture
{
    X64,
    X86
};

inline HRESULT RegisterExtensionWithXbox(std::wstring_view extensionId, std::wstring_view name, std::wstring_view clsId)
{
    RETURN_IF_FAILED(Impl::CreateXboxLocalMachineKey(extensionId, name, clsId));
    RETURN_IF_FAILED(Impl::CreateXboxCurrentUserKey(extensionId));

    return S_OK;
}

inline HRESULT RegisterExtensionWithCOM(
    std::wstring_view clsId,
    std::wstring_view serverExecutablePath,
    Architecture serverArchitecture)
{
    std::wstring formattedClsId;
    formattedClsId.append(L"{").append(clsId).append(L"}");

    auto keyWow64{ serverArchitecture == Architecture::X64 ? KEY_WOW64_64KEY : KEY_WOW64_32KEY };

    wil::unique_hkey classesRootClsidKey;
    auto clsidRegistryPath{ c_baseClassesRootClsidPath + formattedClsId };
    RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
        HKEY_CLASSES_ROOT,
        clsidRegistryPath.data(),
        0,
        nullptr,
        0,
        keyWow64 | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
        nullptr,
        classesRootClsidKey.addressof(),
        nullptr));

    RETURN_IF_WIN32_ERROR(
        RegSetValueExW(classesRootClsidKey.get(), L"", 0, REG_SZ, (BYTE*)L"ExtServer", 10 * sizeof(wchar_t)));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        classesRootClsidKey.get(),
        L"AppId",
        0,
        REG_SZ,
        (BYTE*)formattedClsId.data(),
        static_cast<DWORD>(formattedClsId.size() * sizeof(wchar_t))));

    wil::unique_hkey localServerKey;
    auto localServerSubkeyPath{ clsidRegistryPath + L"\\LocalServer32" };
    RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
        HKEY_CLASSES_ROOT,
        localServerSubkeyPath.data(),
        0,
        nullptr,
        0,
        keyWow64 | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
        nullptr,
        localServerKey.addressof(),
        nullptr));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        localServerKey.get(),
        L"",
        0,
        REG_SZ,
        (BYTE*)serverExecutablePath.data(),
        static_cast<DWORD>(serverExecutablePath.size() * sizeof(wchar_t))));

    wil::unique_hkey classesRootAppIdKey;
    auto appIdRegistryPath{ c_baseClassesRootAppIdPath + formattedClsId };
    RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
        HKEY_CLASSES_ROOT,
        appIdRegistryPath.data(),
        0,
        nullptr,
        0,
        keyWow64 | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
        nullptr,
        classesRootAppIdKey.addressof(),
        nullptr));

    RETURN_IF_WIN32_ERROR(
        RegSetValueExW(classesRootAppIdKey.get(), L"", 0, REG_SZ, (BYTE*)L"ExtServer", 10 * sizeof(wchar_t)));

    // We set this to 0 to allow for elevated instances of the COM server to be able to run and have the Xbox app connect to it
    DWORD appIdFlags = 0;
    RETURN_IF_WIN32_ERROR(
        RegSetValueExW(classesRootAppIdKey.get(), L"AppIdFlags", 0, REG_DWORD, (BYTE*)&appIdFlags, sizeof(appIdFlags)));

    const wchar_t interactiveUserStr[] = L"Interactive User";
    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        classesRootAppIdKey.get(),
        L"RunAs",
        0,
        REG_SZ,
        (BYTE*)interactiveUserStr,
        ARRAYSIZE(interactiveUserStr) * sizeof(wchar_t)));

    wil::unique_hlocal_security_descriptor securityDescriptor{};
    DWORD securityDescriptorSize;
    RETURN_LAST_ERROR_IF(!ConvertStringSecurityDescriptorToSecurityDescriptorW(
        c_comStringSecurityDescriptor,
        SDDL_REVISION_1,
        &securityDescriptor,
        &securityDescriptorSize));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        classesRootAppIdKey.get(),
        L"LaunchPermission",
        0,
        REG_BINARY,
        reinterpret_cast<byte*>(securityDescriptor.get()),
        securityDescriptorSize));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        classesRootAppIdKey.get(),
        L"AccessPermission",
        0,
        REG_BINARY,
        reinterpret_cast<byte*>(securityDescriptor.get()),
        securityDescriptorSize));

    return S_OK;
}

inline void UnregisterExtensionWithCOM(std::wstring_view clsId)
{
    std::wstring formattedClsId;
    formattedClsId.append(L"{").append(clsId).append(L"}");

    // Attempt to delete both x64 and x86 registry view, to be certain registry state isn't left behind.
    std::wstring clsidSubkeyRegistryPath{ c_baseClassesRootClsidPath + formattedClsId + L"\\LocalServer32" };
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, clsidSubkeyRegistryPath.data(), KEY_WOW64_64KEY, 0));
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, clsidSubkeyRegistryPath.data(), KEY_WOW64_32KEY, 0));

    // Attempt to delete both x64 and x86 registry view, to be certain registry state isn't left behind.
    auto clsidRegistryPath{ c_baseClassesRootClsidPath + formattedClsId };
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, clsidRegistryPath.data(), KEY_WOW64_64KEY, 0));
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, clsidRegistryPath.data(), KEY_WOW64_32KEY, 0));

    // Attempt to delete both x64 and x86 registry view, to be certain registry state isn't left behind.
    auto appIdRegistryPath{ c_baseClassesRootAppIdPath + formattedClsId };
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, appIdRegistryPath.data(), KEY_WOW64_64KEY, 0));
    LOG_IF_WIN32_ERROR(RegDeleteKeyExW(HKEY_CLASSES_ROOT, appIdRegistryPath.data(), KEY_WOW64_32KEY, 0));
}

inline ExtensionInstallationState GetExtensionInstallationState(std::wstring_view extensionId)
{
    wil::unique_hkey localMachineKey;

    std::wstring registryPath;
    registryPath.append(c_baseXboxExtensionsDataPath).append(extensionId);

    DWORD extensionInstallationState;
    DWORD size{ sizeof(extensionInstallationState) };

    if (FAILED_WIN32(RegGetValueW(
            HKEY_LOCAL_MACHINE,
            registryPath.data(),
            L"State",
            RRF_SUBKEY_WOW6464KEY | RRF_RT_REG_DWORD,
            nullptr,
            &extensionInstallationState,
            &size)))
    {
        return ExtensionInstallationState::Uninstalled;
    }

    // Convert the registry value to the appropriate ExtensionInstallationState
    // Any unknown values will default to Uninstalled
    switch (extensionInstallationState)
    {
        case static_cast<DWORD>(ExtensionInstallationState::Uninstalled):
            return ExtensionInstallationState::Uninstalled;
        case static_cast<DWORD>(ExtensionInstallationState::Installed):
            return ExtensionInstallationState::Installed;
        case static_cast<DWORD>(ExtensionInstallationState::Updating):
            return ExtensionInstallationState::Updating;
        case static_cast<DWORD>(ExtensionInstallationState::Uninitialized):
            return ExtensionInstallationState::Uninitialized;
        default:
            return ExtensionInstallationState::Uninstalled;
    }
}

inline HRESULT UpdateExtensionInstallationState(
    std::wstring_view extensionId,
    ExtensionInstallationState extensionInstallationState)
{
    wil::unique_hkey localMachineKey;
    wil::unique_hkey currentUserKey;

    std::wstring registryPath;
    registryPath.append(c_baseXboxExtensionsDataPath).append(extensionId);

    RETURN_IF_WIN32_ERROR(RegOpenKeyExW(
        HKEY_LOCAL_MACHINE,
        registryPath.data(),
        0,
        KEY_WOW64_64KEY | KEY_SET_VALUE,
        localMachineKey.addressof()));

    RETURN_IF_WIN32_ERROR(RegOpenKeyExW(
        HKEY_CURRENT_USER,
        registryPath.data(),
        0,
        KEY_WOW64_64KEY | KEY_SET_VALUE,
        currentUserKey.addressof()));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        localMachineKey.get(),
        L"State",
        0,
        REG_DWORD,
        (BYTE*)&extensionInstallationState,
        sizeof(extensionInstallationState)));

    RETURN_IF_WIN32_ERROR(RegSetValueExW(
        currentUserKey.get(),
        L"State",
        0,
        REG_DWORD,
        (BYTE*)&extensionInstallationState,
        sizeof(extensionInstallationState)));

    return S_OK;
}

inline HRESULT SetExtensionFolderPermissions(std::wstring_view serverExecutablePath)
{
    std::wstring executableDirectory{ serverExecutablePath.substr(0, serverExecutablePath.find_last_of('\\')) };

    PACL pOldAcl{ nullptr };
    wil::unique_hlocal_security_descriptor securityDescriptor{};
    RETURN_IF_WIN32_ERROR(GetNamedSecurityInfoW(
        executableDirectory.data(),
        SE_OBJECT_TYPE::SE_FILE_OBJECT,
        DACL_SECURITY_INFORMATION,
        NULL,
        NULL,
        &pOldAcl,
        NULL,
        wil::out_param(securityDescriptor)));

    BYTE anyPackageSid[SECURITY_MAX_SID_SIZE]{};
    DWORD size = SECURITY_MAX_SID_SIZE;
    RETURN_LAST_ERROR_IF(!CreateWellKnownSid(
        WELL_KNOWN_SID_TYPE::WinBuiltinAnyPackageSid,
        nullptr,
        static_cast<PSID>(anyPackageSid),
        &size));

    EXPLICIT_ACCESS explicitAccess{};
    explicitAccess.grfAccessPermissions = GENERIC_READ | GENERIC_EXECUTE;
    explicitAccess.grfAccessMode = ACCESS_MODE::SET_ACCESS;
    explicitAccess.grfInheritance = OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;

    BuildTrusteeWithSidW(&explicitAccess.Trustee, anyPackageSid);

    wil::unique_hlocal_ptr<ACL> newAcl{};
    RETURN_IF_WIN32_ERROR(SetEntriesInAclW(
        /* cCountOfExplicitEntries */ 1,
        &explicitAccess,
        pOldAcl,
        wil::out_param(newAcl)));

    RETURN_IF_WIN32_ERROR(SetNamedSecurityInfoW(
        executableDirectory.data(),
        SE_OBJECT_TYPE::SE_FILE_OBJECT,
        DACL_SECURITY_INFORMATION,
        NULL,
        NULL,
        newAcl.get(),
        NULL));

    return S_OK;
}

inline uint64_t DuplicateCurrentProcessHandleForXboxApp()
{
    RPC_CALL_ATTRIBUTES callAttributes = {};
    callAttributes.Flags = RPC_QUERY_CLIENT_PID;
    callAttributes.Version = RPC_CALL_ATTRIBUTES_VERSION;

    auto status = RpcServerInqCallAttributes(0, &callAttributes);
    if (status != RPC_S_OK)
    {
        std::wstringstream stream;
        stream << L"RpcServerInqCallAttributes failed with status code: " << status << std::endl;
        OutputDebugStringW(stream.str().data());

        return 0;
    }

    uint64_t callingPid = reinterpret_cast<uint64_t>(callAttributes.ClientPID);
    wil::unique_handle callingProcess{ OpenProcess(
        PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE | PROCESS_DUP_HANDLE,
        FALSE,
        static_cast<DWORD>(callingPid)) };

    if (!callingProcess)
    {
        LOG_LAST_ERROR();
        return 0;
    }

    wil::unique_handle thisProcess{
        OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, FALSE, GetCurrentProcessId())
    };

    if (!thisProcess)
    {
        LOG_LAST_ERROR();
        return 0;
    }

    HANDLE duplicatedHandle{};
    if (!DuplicateHandle(
            GetCurrentProcess(),
            thisProcess.get(),
            callingProcess.get(),
            &duplicatedHandle,
            0,
            FALSE,
            DUPLICATE_SAME_ACCESS))
    {
        LOG_LAST_ERROR();
        return 0;
    }

    return reinterpret_cast<uint64_t>(duplicatedHandle);
}

namespace Impl
{
    inline HRESULT CreateXboxLocalMachineKey(
        std::wstring_view extensionId,
        std::wstring_view name,
        std::wstring_view clsId)
    {

        // Set the CLSID key - requires Administrator privileges to modify
        wil::unique_hkey xboxExtensionsClsidKey;

        std::wstring xboxExtensionsClsidRegistryPath(c_baseXboxExtensionsClsidPath);
        xboxExtensionsClsidRegistryPath.append(extensionId);

        RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
            HKEY_LOCAL_MACHINE,
            xboxExtensionsClsidRegistryPath.data(),
            0,
            nullptr,
            0,
            KEY_WOW64_64KEY | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
            nullptr,
            xboxExtensionsClsidKey.addressof(),
            nullptr));

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            xboxExtensionsClsidKey.get(),
            nullptr,
            0,
            REG_SZ,
            (BYTE*)clsId.data(),
            static_cast<DWORD>(clsId.size() * sizeof(wchar_t))));

        // Set the data - modifiable by anyone
        wil::unique_hkey xboxExtensionsDataKey;

        std::wstring xboxExtensionsDataRegistryPath(c_baseXboxExtensionsDataPath);
        xboxExtensionsDataRegistryPath.append(extensionId);

        RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
            HKEY_LOCAL_MACHINE,
            xboxExtensionsDataRegistryPath.data(),
            0,
            nullptr,
            0,
            KEY_WOW64_64KEY | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
            nullptr,
            xboxExtensionsDataKey.addressof(),
            nullptr));

        SetUserWriteAccessToRegistryObject(xboxExtensionsDataKey.get());

        ExtensionType extensionType{ ExtensionType::Server };
        ExtensionInstallationState extensionState{ ExtensionInstallationState::Uninitialized };

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            xboxExtensionsDataKey.get(),
            L"ExtensionType",
            0,
            REG_DWORD,
            (BYTE*)&extensionType,
            sizeof(extensionType)));

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            xboxExtensionsDataKey.get(),
            L"Name",
            0,
            REG_SZ,
            (BYTE*)name.data(),
            static_cast<DWORD>(name.size() * sizeof(wchar_t))));

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            xboxExtensionsDataKey.get(),
            L"State",
            0,
            REG_DWORD,
            (BYTE*)&extensionState,
            sizeof(extensionState)));

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            xboxExtensionsDataKey.get(),
            L"Version",
            0,
            REG_SZ,
            (BYTE*)c_apiVersion,
            static_cast<DWORD>(wcsnlen_s(c_apiVersion, c_maxApiVersionLength) * sizeof(wchar_t))));

        return S_OK;
    }

    inline HRESULT CreateXboxCurrentUserKey(std::wstring_view extensionId)
    {
        wil::unique_hkey currentUserKey;

        std::wstring registryPath;
        registryPath.append(c_baseXboxExtensionsDataPath).append(extensionId);

        RETURN_IF_WIN32_ERROR(RegCreateKeyExW(
            HKEY_CURRENT_USER,
            registryPath.data(),
            0,
            nullptr,
            0,
            KEY_WOW64_64KEY | KEY_CREATE_SUB_KEY | KEY_SET_VALUE,
            nullptr,
            &currentUserKey,
            nullptr));

        ExtensionInstallationState extensionState{ ExtensionInstallationState::Uninitialized };

        RETURN_IF_WIN32_ERROR(RegSetValueExW(
            currentUserKey.get(),
            L"State",
            0,
            REG_DWORD,
            (BYTE*)&extensionState,
            sizeof(extensionState)));

        return S_OK;
    }

    inline void SetExplicitAccessForSid(PSID sid, DWORD accessPermissions, EXPLICIT_ACCESS& explicitAccess)
    {
        explicitAccess.grfAccessPermissions = accessPermissions;
        explicitAccess.grfAccessMode = ACCESS_MODE::SET_ACCESS;
        explicitAccess.grfInheritance = OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;
        BuildTrusteeWithSidW(&explicitAccess.Trustee, sid);
    }

    inline HRESULT SetUserWriteAccessToRegistryObject(HKEY key)
    {
        EXPLICIT_ACCESS explicitAccess[5]{};

        // Set ReadWrite access for 'SYSTEM'
        DWORD systemSidSize{ SECURITY_MAX_SID_SIZE };
        BYTE systemSid[SECURITY_MAX_SID_SIZE]{};
        RETURN_IF_WIN32_BOOL_FALSE(CreateWellKnownSid(
            WELL_KNOWN_SID_TYPE::WinLocalSystemSid,
            nullptr,
            static_cast<PSID>(systemSid),
            &systemSidSize));
        SetExplicitAccessForSid(systemSid, KEY_ALL_ACCESS, explicitAccess[0]);

        // Set ReadWrite access for 'CREATOR OWNER'
        DWORD creatorOwnerSidSize{ SECURITY_MAX_SID_SIZE };
        BYTE creatorOwnerSid[SECURITY_MAX_SID_SIZE]{};
        RETURN_IF_WIN32_BOOL_FALSE(CreateWellKnownSid(
            WELL_KNOWN_SID_TYPE::WinCreatorOwnerSid,
            nullptr,
            static_cast<PSID>(creatorOwnerSid),
            &creatorOwnerSidSize));
        SetExplicitAccessForSid(creatorOwnerSid, KEY_ALL_ACCESS, explicitAccess[1]);

        // Set ReadWrite access for 'Administrators'
        DWORD adminSidSize{ SECURITY_MAX_SID_SIZE };
        BYTE adminSid[SECURITY_MAX_SID_SIZE]{};
        RETURN_IF_WIN32_BOOL_FALSE(CreateWellKnownSid(
            WELL_KNOWN_SID_TYPE::WinBuiltinAdministratorsSid,
            nullptr,
            static_cast<PSID>(adminSid),
            &adminSidSize));
        SetExplicitAccessForSid(adminSid, KEY_ALL_ACCESS, explicitAccess[2]);

        // Set Read access for 'ALL_APPLICATION_PACKAGES'
        DWORD packagesSidSize{ SECURITY_MAX_SID_SIZE };
        BYTE packagesSid[SECURITY_MAX_SID_SIZE]{};
        RETURN_IF_WIN32_BOOL_FALSE(CreateWellKnownSid(
            WELL_KNOWN_SID_TYPE::WinBuiltinAnyPackageSid,
            nullptr,
            static_cast<PSID>(packagesSid),
            &packagesSidSize));
        SetExplicitAccessForSid(packagesSid, KEY_READ, explicitAccess[3]);

        // Set ReadWrite access 'Users'
        DWORD usersSidSize{ SECURITY_MAX_SID_SIZE };
        BYTE usersSid[SECURITY_MAX_SID_SIZE]{};
        RETURN_IF_WIN32_BOOL_FALSE(CreateWellKnownSid(
            WELL_KNOWN_SID_TYPE::WinBuiltinUsersSid,
            nullptr,
            static_cast<PSID>(usersSid),
            &usersSidSize));
        SetExplicitAccessForSid(usersSid, KEY_ALL_ACCESS, explicitAccess[4]);

        wil::unique_hlocal_ptr<ACL> newAcl{};
        RETURN_IF_WIN32_ERROR(SetEntriesInAclW(
            /* cCountOfExplicitEntries */ 5,
            explicitAccess,
            NULL,
            wil::out_param(newAcl)));

        RETURN_IF_WIN32_ERROR(SetSecurityInfo(
            key,
            SE_OBJECT_TYPE::SE_REGISTRY_KEY,
            DACL_SECURITY_INFORMATION,
            NULL,
            NULL,
            newAcl.get(),
            NULL));

        return S_OK;
    }
} // namespace Impl

} // namespace Microsoft::Gaming::XboxApp::Extensions