/*
 * Copyright 2023 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "processinfohelper.h"
#include <QDebug>
#include <QFile>
#include <fstream>
#include <iostream>
#include <libproc2/pids.h> // libproc2 的主头文件
#include <sstream>
#include <string>
#include <unordered_map>
#include <dirent.h>
#include <unordered_set>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <unistd.h>
#include <sys/types.h>



namespace {
void collectSocketInodes(int pid, std::unordered_set<unsigned long>& out) {
    std::string fdPath = "/proc/" + std::to_string(pid) + "/fd";
    DIR* dir = opendir(fdPath.c_str());
    if (!dir) return;

    dirent* entry;
    char linkBuf[4096] = {0};

    while ((entry = readdir(dir)) != nullptr) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) continue;

        std::string linkPath = fdPath + "/" + entry->d_name;
        ssize_t len = readlink(linkPath.c_str(), linkBuf, sizeof(linkBuf)-1);
        if (len <= 0) continue;

        linkBuf[len] = '\0';
        if (strstr(linkBuf, "socket:") == linkBuf) {
            unsigned long inode;
            if (sscanf(linkBuf + 7, "[%lu]", &inode) == 1) {
                out.insert(inode);
            }
        }
    }
    closedir(dir);
}

std::unordered_set<unsigned long> readSystemUdpInodes() {
    std::unordered_set<unsigned long> udpInodes;
    std::ifstream udpFile("/proc/net/udp");
    if (!udpFile.is_open()) return udpInodes;

    std::string line;
    std::getline(udpFile, line);

    while (std::getline(udpFile, line)) {
        std::istringstream iss(line);
        std::vector<std::string> fields;
        std::string field;

        while (iss >> field) fields.push_back(field);

        if (fields.size() < 10) continue;
        
        char* endPtr;
        unsigned long inode = strtoul(fields[9].c_str(), &endPtr, 10);
        if (endPtr == fields[9].c_str() || *endPtr != '\0') {
            continue;
        }
        udpInodes.insert(inode);
    }
    return udpInodes;
}

}   // namespace

namespace process_info_helper
{

std::string cmdline(int pid)
{
    if (pid <= 0) {
        return std::string();
    }
    std::string cmdlineFilePath = "/proc/" + std::to_string(pid) + "/cmdline";
    std::ifstream cmdlineFile(cmdlineFilePath);

    if (!cmdlineFile.is_open()) {
        qWarning() << "cmdlineFile open failed!" << pid;
        return std::string();
    }

    std::stringstream buffer;
    buffer << cmdlineFile.rdbuf();
    std::string cmdline = buffer.str();

    // Replace '\0' with ' '
    for (size_t i = 0; i < cmdline.length(); ++i) {
        if (cmdline[i] == '\0') {
            cmdline[i] = ' ';
        }
    }

    if (!cmdline.empty() && cmdline.at(cmdline.length() - 1) == ' ') {
        cmdline.pop_back();
    }

    return cmdline;
}

std::string commContent(int pid)
{
    if (pid <= 0) {
        return std::string();
    }
    std::string commPath = "/proc/" + std::to_string(pid) + "/comm";
    std::ifstream commFile(commPath);

    if (!commFile.is_open()) {
        qWarning() << "commFile open failed!" << pid;
        return std::string();
    }

    std::stringstream buffer;
    buffer << commFile.rdbuf();
    std::string commContent = buffer.str();

    // Replace '\0' with ' '
    for (size_t i = 0; i < commContent.length(); ++i) {
        if (commContent[i] == '\0') {
            commContent[i] = ' ';
        }
    }

    if (!commContent.empty() && commContent.at(commContent.length() - 1) == ' ') {
        commContent.pop_back();
    }

    return commContent;
}

bool processExists(int pid)
{
    std::ifstream file("/proc/" + std::to_string(pid));
    return file.good();
}

int parentPid(int pid)
{
    QString procStatus = QString("/proc/%1/status").arg(pid);
    QFile file(procStatus);
    if (!file.open(QIODevice::ReadOnly)) {
        qWarning() << "Get parent id failed, " << pid << file.errorString();
        return -1;
    }
    QTextStream ts(&file);
    QString line;
    while (ts.readLineInto(&line)) {
        if (line.contains("PPid:")) {
            QString strPpid = line.split(":").last();
            strPpid.remove(QChar::Space);
            return strPpid.toUInt();
        }
    }
    return -1;
}

std::vector<int> childPids(int parentPid)
{
    std::vector<int> pids;
    struct pids_info *info = NULL;

    // 设置我们需要获取的数据项
    enum pids_item items[] = {
        PIDS_ID_PID, // 进程ID
        PIDS_ID_PPID // 父进程ID
    };
    const int numitems = sizeof(items) / sizeof(enum pids_item);

    // 初始化进程信息结构
    if (procps_pids_new(&info, items, numitems) < 0) {
        qWarning() << "Failed to initialize pids_info";
        return pids;
    }

    // 获取进程信息
    struct pids_stack *stack;
    while ((stack = procps_pids_get(info, PIDS_FETCH_TASKS_ONLY))) {
        // 使用 PIDS_VAL 宏获取值
        int ppid = PIDS_VAL(1, s_int, stack, info); // PIDS_ID_PPID
        if (ppid == parentPid) {
            int pid = PIDS_VAL(0, s_int, stack, info); // PIDS_ID_PID
            pids.emplace_back(pid);
        }
    }

    procps_pids_unref(&info);

    return pids;
}

std::vector<int> childPids(const std::vector<int> &parentPids)
{
    std::vector<int> cPids;
    std::unordered_map<int, int> mp;
    for (int i = 0; i < parentPids.size(); ++i) {
        mp[parentPids[i]]++;
    }

    struct pids_info *info = NULL;

    // 设置我们需要获取的数据项
    enum pids_item items[] = {
        PIDS_ID_PID, // 进程ID
        PIDS_ID_PPID // 父进程ID
    };
    const int numitems = sizeof(items) / sizeof(enum pids_item);

    // 初始化进程信息结构
    if (procps_pids_new(&info, items, numitems) < 0) {
        qWarning() << "Failed to initialize pids_info";
        return cPids;
    }

    // 获取进程信息
    struct pids_stack *stack;
    while ((stack = procps_pids_get(info, PIDS_FETCH_TASKS_ONLY))) {
        // 使用 PIDS_VAL 宏获取值
        int ppid = PIDS_VAL(1, s_int, stack, info); // PIDS_ID_PPID
        if (mp.find(ppid) != mp.end()) {
            int pid = PIDS_VAL(0, s_int, stack, info); // PIDS_ID_PID
            cPids.emplace_back(pid);
        }
    }

    procps_pids_unref(&info);

    if (cPids.empty()) {
        return cPids;
    }

    // 函数其余部分保持不变
    std::vector<int> nextSubPids = childPids(cPids);
    if (nextSubPids.empty()) {
        return cPids;
    }

    for (const auto &pid : cPids) {
        if (std::find(nextSubPids.cbegin(), nextSubPids.cend(), pid) != nextSubPids.cend()) {
            nextSubPids.erase(std::remove(nextSubPids.begin(), nextSubPids.end(), pid), nextSubPids.end());
        }
    }

    cPids.insert(cPids.end(), nextSubPids.begin(), nextSubPids.end());

    return cPids;
}

std::string cgroup(int pid)
{
    if (pid <= 0) {
        return std::string();
    }
    std::string cgroupFilePath = "/proc/" + std::to_string(pid) + "/cgroup";
    std::ifstream cgroupFile(cgroupFilePath);

    if (!cgroupFile.is_open()) {
        qWarning() << "cgroupFile open failed!" << pid;
        return std::string();
    }

    std::string cgroup, line;
    while (std::getline(cgroupFile, line)) {
        auto lastPos = line.rfind("::");
        if (lastPos != std::string::npos) {
            cgroup = line.substr(lastPos + 2);
            break;
        }
    }
    cgroupFile.close();
    return cgroup;
}

std::vector<int> pidsOfCgroup(const std::string &cgroupPath)
{
    std::ifstream cgroupFile(cgroupPath + "/cgroup.procs");

    if (!cgroupFile.is_open()) {
        return {};
    }

    std::string line;
    std::vector<int> pids;
    while (std::getline(cgroupFile, line)) {
        pids.emplace_back(std::stoi(line));
    }
    cgroupFile.close();
    return pids;
}

bool isShellScript(const QString &filePath)
{
    // QProcess process;
    // process.start("file", QStringList() << "--mime-type" << filePath);
    // if (!process.waitForFinished()) {
    //     qDebug() << "执行file命令失败";
    //     return false;
    // }

    // QString output = process.readAllStandardOutput();
    // return (output.contains("text/x-shellscript") || output.contains("shell script"));
    std::ifstream file(filePath.toStdString());
    if (!file.is_open()) return false;

    std::string firstLine;
    std::getline(file, firstLine);
    file.close();

    return firstLine.find("#!") == 0;

}


bool isDigitString(const std::string& s) 
{
    if (s.empty()) return false;
    for (char c : s) {
        if (!isdigit(c)) return false;
    }
    return true;
}

std::vector<int> allPids() 
{
    std::vector<int> pids;
    DIR* procDir = opendir("/proc");
    if (!procDir) {
        perror("opendir(/proc) failed");
        return pids;
    }

    dirent* entry;
    while ((entry = readdir(procDir)) != nullptr) {
        if (!isDigitString(entry->d_name)) continue; 
        
        int pid = static_cast<int>(atoi(entry->d_name));
        if (pid > 0) pids.push_back(pid);
    }

    closedir(procDir);
    return pids;
}

std::string processName(int pid) 
{
    std::string commPath = "/proc/" + std::to_string(pid) + "/comm";
    std::ifstream commFile(commPath);
    std::string name((std::istreambuf_iterator<char>(commFile)),
                            std::istreambuf_iterator<char>());

    if (name.back() == '\n'){
        name.pop_back();
    }
    if (!name.empty()){
        std::replace(name.begin(), name.end() - 1, '\0', ' ');
    }
    return name;
}

std::string processCmdLine(int pid) 
{
    std::string cmdlinePath = "/proc/" + std::to_string(pid) + "/cmdline";

    std::ifstream cmdlineFile(cmdlinePath);
    std::string cmdline((std::istreambuf_iterator<char>(cmdlineFile)),
                            std::istreambuf_iterator<char>());

    if (!cmdline.empty())
    {
        std::replace(cmdline.begin(), cmdline.end() - 1, '\0', ' ');
    }
    return cmdline;
}


std::vector<int> pidsByName(const std::string& targetName, bool exactMatch) 
{
    std::vector<int> result;
    std::vector<int> pids = allPids();

    for (int pid : pids) {
        std::string procName = processName(pid);
        if (procName.empty()) continue;

        bool matched = false;
        if (exactMatch) {
            matched = (procName == targetName);
        } else {
            matched = (procName.find(targetName) != std::string::npos);
        }

        if (matched) {
            result.push_back(pid);
        }
    }

    return result;
}

std::unordered_map<std::string, std::vector<int>> pidsByNames(const std::vector<std::string>& targetNames, bool exactMatch) 
{
    std::unordered_map<std::string, std::vector<int>> result;
    std::unordered_set<std::string> targetSet(targetNames.begin(), targetNames.end());
    std::vector<int> pids = allPids();

    for (int pid : pids) {
        std::string procName = processName(pid);
        if (procName.empty()) continue;

        for (const auto& target : targetNames) {
            bool matched = false;
            if (exactMatch) {
                matched = (procName == target);
            } else {
                matched = (procName.find(target) != std::string::npos);
            }
            if (matched) {
                result[target].push_back(pid);
                break;
            }
        }
    }

    return result;
}

std::vector<int> pidsByCmdline(const std::string& targetName, bool exactMatch)
{
    std::vector<int> result;
    std::vector<int> pids = allPids();

    for (int pid : pids) {
        std::string procCmdline = processCmdLine(pid);
        if (procCmdline.empty()) continue;

        bool matched = false;
        if (exactMatch) {
            matched = (procCmdline == targetName);
        } else {
            matched = (procCmdline.find(targetName) != std::string::npos);
        }

        if (matched) {
            result.push_back(pid);
        }
    }

    return result;
}

std::unordered_map<std::string, std::vector<int>> pidsByCmdlines(const std::vector<std::string>& targetNames, bool exactMatch)
{
    std::unordered_map<std::string, std::vector<int>> result;
    std::unordered_set<std::string> targetSet(targetNames.begin(), targetNames.end());
    std::vector<int> pids = allPids();

    for (int pid : pids) {
        std::string procCmdline = processCmdLine(pid);
        if (procCmdline.empty()) continue;

        for (const auto& target : targetNames) {
            bool matched = false;
            if (exactMatch) {
                matched = (procCmdline == target);
            } else {
                matched = (procCmdline.find(target) != std::string::npos);
            }
            if (matched) {
                result[target].push_back(pid);
                break;
            }
        }
    }

    return result;
}

int udpSocketCount(int pid) {
    if (pid <= 0) return -1;

    std::unordered_set<unsigned long> procSockets;
    collectSocketInodes(pid, procSockets);
    if (procSockets.empty()) return 0;

    std::unordered_set<unsigned long> systemUdpInodes = readSystemUdpInodes();
    if (systemUdpInodes.empty()) return 0;

    return std::count_if(procSockets.begin(), procSockets.end(),
        [&systemUdpInodes](unsigned long inode) {
            return systemUdpInodes.count(inode);
        });
}

}
