/*
 * MIT License
 *
 * Copyright (c) 2025 Vanessa T. <nessa@neko-tools.de>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

#include <stdexcept>
#include <format>

#include <winternl.h>
#include <ntstatus.h>
#include <psapi.h>

#include "Process.h"
#include "Utility.h"

Process::Process(DWORD processId) : m_processId(processId) {

}

Process::~Process() {
    close();
}

DWORD Process::processId() const {
    return m_processId;
}

void Process::open() {
    if (m_hProcess) return;

    m_limited = false;
    m_hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, m_processId);

    if (!m_hProcess) {
        m_limited = true;
        m_hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, m_processId);
    }

    if (!m_hProcess) {
        checkAndThrowLastWindowsError("Failed to open process " + std::to_string(m_processId));
    }
}

void Process::close() {
    if (m_hProcess) {
        CloseHandle(m_hProcess);
    }
}

std::wstring Process::imagePath() const {
    if (!isOpen()) throw std::runtime_error("Process is not open");

    ULONG length = 0;
    NTSTATUS status = NtQueryInformationProcess(m_hProcess, ProcessImageFileName, nullptr, 0, &length);
    if (status != STATUS_SUCCESS && status != STATUS_INFO_LENGTH_MISMATCH) {
        throw std::runtime_error("Unable to get process image path length: NTSTATUS = " +
                                 std::format("{:0x}", static_cast<ULONG>(status)));
    }

    std::vector<wchar_t> buffer(length + 1);
    if (!QueryFullProcessImageNameW(m_hProcess, 0, buffer.data(), &length)) {
        checkAndThrowLastWindowsError("Unable to get process image path");
    }

    return {buffer.data()};
}

Process::MemoryUsage Process::ramUsage() const {
    if (!isOpen()) throw std::runtime_error("Process is not open");

    MemoryUsage usage{};

    PROCESS_MEMORY_COUNTERS_EX counters = {
            .cb = sizeof counters
    };

    if (!GetProcessMemoryInfo(m_hProcess, reinterpret_cast<PROCESS_MEMORY_COUNTERS *>(&counters), sizeof(counters))) {
        checkAndThrowLastWindowsError("Unable to get process memory usage");
    }

    usage.pageFaults = counters.PageFaultCount;
    usage.peakWorkingSetSize = counters.PeakWorkingSetSize;
    usage.workingSetSize = counters.WorkingSetSize;
    usage.quotaPeakPagedPoolUsage = counters.QuotaPeakPagedPoolUsage;
    usage.quotaPagedPoolUsage = counters.QuotaPagedPoolUsage;
    usage.quotaPeakNonPagedPoolUsage = counters.QuotaPeakNonPagedPoolUsage;
    usage.quotaNonPagedPoolUsage = counters.QuotaNonPagedPoolUsage;
    usage.pageFileUsage = counters.PagefileUsage;
    usage.peakPageFileUsage = counters.PeakPagefileUsage;
    usage.privateUsage = counters.PrivateUsage;

    return usage;
}

bool Process::isLimited() const {
    return m_limited;
}

std::vector<Process> Process::enumerate() {
    std::vector<Process> processes{};

    bool foundAllProcesses = false;
    std::vector<DWORD> processIds(1024);
    while (!foundAllProcesses) {
        auto bufferSize = processIds.size() * sizeof(DWORD);
        DWORD bytesReturned;

        if (!EnumProcesses(processIds.data(), bufferSize, &bytesReturned)) {
            checkAndThrowLastWindowsError("Unable to enumerate processes");
        }

        if (bytesReturned <= bufferSize) {
            processIds.resize(bytesReturned / sizeof(DWORD));
            foundAllProcesses = true;
        } else {
            // Resize to 1.5 times the current size
            processIds.resize(processes.size() + processes.size() / 2);
        }
    }

    processes.reserve(processIds.size() - 1);
    for (auto processId: processIds) {
        // Skip system idle process
        if (processId == 0) continue;
        processes.emplace_back(processId);
    }

    return processes;
}

bool Process::isOpen() const {
    return m_hProcess != nullptr;
}