#include "processmonitor.h"
BOOL EnablePrivilege(LPCWSTR privilege)
{
    HANDLE hToken;
    TOKEN_PRIVILEGES tp;
    LUID luid;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        return FALSE;
    }

    if (!LookupPrivilegeValue(NULL, privilege, &luid)) {
        CloseHandle(hToken);
        return FALSE;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) {
        CloseHandle(hToken);
        return FALSE;
    }

    CloseHandle(hToken);
    return GetLastError() == ERROR_SUCCESS;
}
ProcessMonitor::ProcessMonitor()
{
    filter.insert({"System", true});
    filter.insert({"svchost.exe", true});
    filter.insert({"wininit.exe", true});
    filter.insert({"NVDisplay.Container.exe", true});
}

std::vector<std::shared_ptr<ProcessMonitor::ProcessInfo>> ProcessMonitor::checkProcesses()
{
    data.clear();
    mapData.clear();
    // 创建进程快照
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnapshot != INVALID_HANDLE_VALUE) {
        PROCESSENTRY32 pe32;
        pe32.dwSize = sizeof(PROCESSENTRY32);

        // 获取第一个进程信息
        if (Process32First(hSnapshot, &pe32)) {
            do {
                // 获取当前进程的PID
                DWORD pid = pe32.th32ProcessID;
                if (pid != 0) {
                    // 获取进程信息
                    std::shared_ptr<ProcessInfo> info = getProcessInfo(pid);
                    info->parentPid = pe32.th32ParentProcessID; // 直接从pe32获取父进程PID
                    std::string processName = QString::fromWCharArray(pe32.szExeFile).toStdString();
                    info->processName = processName;
                    data.push_back(info);
                    mapData[pid] = info;
                }
            } while (Process32Next(hSnapshot, &pe32)); // 循环获取下一个进程信息
        }
        CloseHandle(hSnapshot);
    }
    std::vector<std::shared_ptr<ProcessInfo>> timeProcess;
    // 获取全部根节点信息
    auto rootNodes = root();
    std::shared_ptr<ProcessInfo> explorerNode = nullptr;
    for (auto rootNode : rootNodes) {
        //在过滤里面
        if (rootNode) {
            if (filter.find(rootNode->processName) != filter.end()) {
            } else {
                // qDebug() << "[" << QString::fromStdString(rootNode->processName) << "]"
                //          << rootNode->pid << rootNode->parentPid;
                timeProcess.push_back(rootNode);
            }
            if (rootNode->processName == "explorer.exe" || rootNode->processName == "Explorer.exe") {
                explorerNode = rootNode;
            }
        }
    }
    // 获取全部根节点的过滤
    if (explorerNode) {
        for (const auto& item : data) {
            if (item->parentPid == explorerNode->pid) {
                // qDebug() << "-- explorerNode [" << QString::fromStdString(item->processName) << "]"
                //          << item->pid << item->parentPid;
                timeProcess.push_back(item);
            }
        }
    }
    return timeProcess;
}
DWORD ProcessMonitor::getRootPid(DWORD pid)
{
    if (mapData.find(pid) != mapData.end()) {
        const auto value = mapData.at(pid);
        if (pid == value->pid) {
            return pid;
        }
        if (pid == value->parentPid) {
            return pid;
        }
        DWORD id = getRootPid(value->parentPid);
        if (id == -1) {
            return pid;
        }
    }
    return -1;
}

std::set<std::shared_ptr<ProcessMonitor::ProcessInfo>> ProcessMonitor::root()
{
    std::set<std::shared_ptr<ProcessInfo>> roots;
    for (const auto& item : data) {
        DWORD rootPid = getRootPid(item->parentPid);
        if (rootPid != -1) {
            roots.insert(mapData[rootPid]);
        }
    }
    return roots;
}
std::shared_ptr<ProcessMonitor::ProcessInfo> ProcessMonitor::getProcessInfo(DWORD pid)
{
    std::shared_ptr<ProcessInfo> info = std::make_shared<ProcessInfo>();
    // ProcessInfo info;
    info->pid = pid;

    HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
    if (hProcess) {
        char processName[MAX_PATH] = {0};
        if (GetModuleFileNameExA(hProcess, NULL, processName, sizeof(processName) / sizeof(char))) {
            info->name = std::string(processName);
        }

        // 获取进程时间
        if (GetProcessTimes(hProcess,
                            &info->creationTime,
                            &info->exitTime,
                            &info->kernelTime,
                            &info->userTime)) {
        }

        CloseHandle(hProcess);
    }
    return info;
}