/*
    在包含任何 Windows 头文件之前定义宏 WIN32_LEAN_AND_MEAN
*/
#define WIN32_LEAN_AND_MEAN
#include <algorithm>
#include <filesystem>
#include <iostream>
#include <shlobj.h>
#include <string>
#include <vector>
#include <windows.h>

#ifndef WHAT_FUCK
#include <shellapi.h>   // 需要在 windows.h 之后包含
#endif

namespace fs = std::filesystem;

class RegistryManager
{
private:
    std::string wtPath;

    bool checkWtExists()
    {
        if (fs::exists(wtPath) && fs::path(wtPath).filename() == "wt.exe") {
            std::cout << "找到 wt.exe: " << wtPath << std::endl;
            return true;
        } else {
            std::cout << "错误: wt.exe 不存在或路径无效: " << wtPath << std::endl;
            return false;
        }
    }

    bool isAdmin()
    {
        BOOL isAdmin = FALSE;
        PSID adminGroup = NULL;
        SID_IDENTIFIER_AUTHORITY ntAuthority = SECURITY_NT_AUTHORITY;

        if (AllocateAndInitializeSid(&ntAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0,
                                     &adminGroup)) {
            if (!CheckTokenMembership(NULL, adminGroup, &isAdmin)) {
                isAdmin = FALSE;
            }
            FreeSid(adminGroup);
        }

        return isAdmin != FALSE;
    }

    // 创建注册表键并设置值
    bool createRegistryKeyWithValue(const std::string& keyPath, const std::string& valueName, const std::string& valueData)
    {
        HKEY hKey;

        std::cout << "创建注册表键: " << keyPath << std::endl;
        if (!valueName.empty()) {
            std::cout << "设置值: " << valueName << " = " << valueData << std::endl;
        } else {
            std::cout << "设置默认值: " << valueData << std::endl;
        }
        
        LONG result =
            RegCreateKeyEx(HKEY_CLASSES_ROOT, keyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, NULL);

        if (result != ERROR_SUCCESS) {
            std::cout << "创建注册表键失败: " << keyPath << " 错误码: " << result << std::endl;
            return false;
        }

        if (!valueName.empty() || valueData.length() > 0) {
            result = RegSetValueEx(hKey, valueName.c_str(), 0, REG_SZ, (const BYTE*)valueData.c_str(),
                                   static_cast<DWORD>(valueData.length() + 1));

            if (result != ERROR_SUCCESS) {
                std::cout << "设置注册表值失败: " << valueName << " 错误码: " << result << std::endl;
                RegCloseKey(hKey);
                return false;
            }
        }

        RegCloseKey(hKey);
        return true;
    }

    // 创建注册表键（不设置值）
    bool createRegistryKey(const std::string& keyPath)
    {
        HKEY hKey;

        std::cout << "创建注册表键: " << keyPath << std::endl;

        LONG result =
            RegCreateKeyEx(HKEY_CLASSES_ROOT, keyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, NULL);

        if (result != ERROR_SUCCESS) {
            std::cout << "创建注册表键失败: " << keyPath << " 错误码: " << result << std::endl;
            return false;
        }

        RegCloseKey(hKey);
        return true;
    }

    // 递归删除注册表键及其所有子键
    bool deleteRegistryKeyRecursive(HKEY hRoot, const std::string& subKey)
    {
        HKEY hKey;

        // 尝试打开键
        LONG result = RegOpenKeyEx(hRoot, subKey.c_str(), 0, KEY_READ | KEY_WRITE, &hKey);

        if (result != ERROR_SUCCESS) {
            if (result == ERROR_FILE_NOT_FOUND) {
                // 键不存在，视为成功删除
                return true;
            }
            std::cout << "打开注册表键失败: " << subKey << " 错误码: " << result << std::endl;
            return false;
        }

        // 递归删除所有子键
        DWORD subKeyCount = 0;
        DWORD maxSubKeyLen = 0;

        result = RegQueryInfoKey(hKey, NULL, NULL, NULL, &subKeyCount, &maxSubKeyLen, NULL, NULL, NULL, NULL, NULL, NULL);

        if (result == ERROR_SUCCESS && subKeyCount > 0) {
            std::vector<char> subKeyName(maxSubKeyLen + 2, 0);   // +2 保证有足够空间

            for (DWORD i = 0; i < subKeyCount; i++) {
                DWORD nameLen = maxSubKeyLen + 1;
                result = RegEnumKeyEx(hKey,
                                      0,   // 总是删除第一个子键
                                      &subKeyName[0], &nameLen, NULL, NULL, NULL, NULL);

                if (result == ERROR_SUCCESS) {
                    std::string childPath = subKey + "\\" + std::string(&subKeyName[0]);
                    if (!deleteRegistryKeyRecursive(hRoot, childPath)) {
                        RegCloseKey(hKey);
                        return false;
                    }
                } else {
                    RegCloseKey(hKey);
                    std::cout << "枚举子键失败: " << subKey << " 错误码: " << result << std::endl;
                    return false;
                }
            }
        }

        RegCloseKey(hKey);

        // 删除当前键
        result = RegDeleteKey(hRoot, subKey.c_str());
        if (result != ERROR_SUCCESS) {
            std::cout << "删除注册表键失败: " << subKey << " 错误码: " << result << std::endl;
            return false;
        }

        return true;
    }

    bool deleteRegistryKey(const std::string& keyPath)
    {
        return deleteRegistryKeyRecursive(HKEY_CLASSES_ROOT, keyPath);
    }

    bool installContextMenu()
    {
        std::cout << "安装 Windows Terminal 右键菜单..." << std::endl;

        bool success = true;

        // 1. 创建背景右键菜单项
        std::cout << "\n=== 安装背景右键菜单 ===" << std::endl;

        // 1.1 创建 wt 键
        if (!createRegistryKey("Directory\\Background\\shell\\wt")) {
            std::cout << "创建背景菜单键失败" << std::endl;
            success = false;
        }

        // 1.2 设置默认值
        if (!createRegistryKeyWithValue("Directory\\Background\\shell\\wt", "", "Terminal")) {
            std::cout << "设置背景菜单名称失败" << std::endl;
            success = false;
        }

        // 1.3 设置图标
        if (!createRegistryKeyWithValue("Directory\\Background\\shell\\wt", "Icon", "\"" + wtPath + "\"")) {
            std::cout << "设置背景菜单图标失败" << std::endl;
            success = false;
        }

        // 1.4 创建 command 子键
        if (!createRegistryKey("Directory\\Background\\shell\\wt\\command")) {
            std::cout << "创建背景命令键失败" << std::endl;
            success = false;
        }

        // 1.5 设置命令
        if (!createRegistryKeyWithValue("Directory\\Background\\shell\\wt\\command", "", "\"" + wtPath + "\"")) {
            std::cout << "设置背景命令失败" << std::endl;
            success = false;
        }

        // 2. 创建目录右键菜单项
        std::cout << "\n=== 安装目录右键菜单 ===" << std::endl;

        // 2.1 创建 wt 键
        if (!createRegistryKey("Directory\\shell\\wt")) {
            std::cout << "创建目录菜单键失败" << std::endl;
            success = false;
        }

        // 2.2 设置默认值
        if (!createRegistryKeyWithValue("Directory\\shell\\wt", "", "Terminal")) {
            std::cout << "设置目录菜单名称失败" << std::endl;
            success = false;
        }

        // 2.3 设置图标
        if (!createRegistryKeyWithValue("Directory\\shell\\wt", "Icon", "\"" + wtPath + "\"")) {
            std::cout << "设置目录菜单图标失败" << std::endl;
            success = false;
        }

        // 2.4 创建 command 子键
        if (!createRegistryKey("Directory\\shell\\wt\\command")) {
            std::cout << "创建目录命令键失败" << std::endl;
            success = false;
        }

        // 2.5 设置命令 - 注意：这里需要转义双引号
        std::string command = "\"" + wtPath + "\" -d \"%V\"";
        if (!createRegistryKeyWithValue("Directory\\shell\\wt\\command", "", command)) {
            std::cout << "设置目录命令失败" << std::endl;
            success = false;
        }

        return success;
    }

    bool uninstallContextMenu()
    {
        std::cout << "卸载 Windows Terminal 右键菜单..." << std::endl;

        bool success = true;

        // 删除背景右键菜单项
        if (!deleteRegistryKey("Directory\\Background\\shell\\wt")) {
            std::cout << "警告: Directory\\Background\\shell\\wt 删除失败，可能不存在或已删除" << std::endl;
        }

        // 删除目录右键菜单项
        if (!deleteRegistryKey("Directory\\shell\\wt")) {
            std::cout << "警告: Directory\\shell\\wt 删除失败，可能不存在或已删除" << std::endl;
        }

        return success;
    }

    void refreshExplorer()
    {
        // 发送消息刷新资源管理器
        SendMessageTimeout(HWND_BROADCAST, WM_SETTINGCHANGE, 0,
                           (LPARAM) "Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\Advanced", SMTO_ABORTIFHUNG, 1000,
                           NULL);

        // 通知资源管理器更新
        SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);
    }

public:
    RegistryManager(const std::string& path) : wtPath(path)
    {
        // 规范化路径（统一使用反斜杠）
        std::replace(wtPath.begin(), wtPath.end(), '/', '\\');
    }

    bool install()
    {
        if (!checkWtExists()) {
            return false;
        }

        if (!isAdmin()) {
            std::cout << "需要管理员权限来修改注册表" << std::endl;
            std::cout << "请以管理员身份重新运行此程序" << std::endl;
            return false;
        }

        bool success = installContextMenu();
        if (success) {
            refreshExplorer();
            std::cout << "\n安装成功!" << std::endl;
        } else {
            std::cout << "\n安装失败!" << std::endl;
        }

        return success;
    }

    bool uninstall()
    {
        if (!isAdmin()) {
            std::cout << "需要管理员权限来修改注册表" << std::endl;
            std::cout << "请以管理员身份重新运行此程序" << std::endl;
            return false;
        }

        bool success = uninstallContextMenu();
        if (success) {
            refreshExplorer();
            std::cout << "卸载成功!" << std::endl;
        } else {
            std::cout << "卸载失败!" << std::endl;
        }

        return success;
    }
};

void printUsage(const char* programName)
{
    std::cout << "Windows Terminal 右键菜单管理工具" << std::endl;
    std::cout << "用法:" << std::endl;
    std::cout << "  " << programName << " install <wt.exe路径>" << std::endl;
    std::cout << "  " << programName << " uninstall" << std::endl;
    std::cout << std::endl;
    std::cout << "示例:" << std::endl;
    std::cout << "  " << programName << " install \"D:\\Program\\terminal\\wt.exe\"" << std::endl;
    std::cout << "  " << programName << " uninstall" << std::endl;
}

int main(int argc, char* argv[])
{
    // 设置控制台输出为UTF-8
    SetConsoleOutputCP(CP_UTF8);

    if (argc < 2) {
        printUsage(argv[0]);
        return 1;
    }

    std::string action = argv[1];

    if (action == "install") {
        if (argc < 3) {
            std::cout << "错误: 需要指定 wt.exe 的路径" << std::endl;
            printUsage(argv[0]);
            return 1;
        }

        std::string wtPath = argv[2];
        RegistryManager manager(wtPath);

        return manager.install() ? 0 : 1;
    } else if (action == "uninstall") {
        std::string dummyPath = "C:\\Windows\\System32\\cmd.exe";
        RegistryManager manager(dummyPath);

        return manager.uninstall() ? 0 : 1;
    } else {
        std::cout << "错误: 未知的操作: " << action << std::endl;
        printUsage(argv[0]);
        return 1;
    }
}