#include "fastfetch.h"
#include "common/processing.h"
#include "common/io.h"
#include "common/stringUtils.h"
#include "common/mallocHelper.h"

#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <poll.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/wait.h>

#if !(__ANDROID__ || __OpenBSD__)
    #include <spawn.h>
#endif

#if defined(__FreeBSD__) || defined(__APPLE__)
    #include <sys/types.h>
    #include <sys/user.h>
    #include <sys/sysctl.h>
#endif
#if defined(__APPLE__)
    #include <libproc.h>
#elif defined(__sun)
    #include <procfs.h>
#elif defined(__OpenBSD__)
    #include <sys/param.h>
    #include <sys/sysctl.h>
    #include <kvm.h>
#elif defined(__NetBSD__)
    #include <sys/types.h>
    #include <sys/sysctl.h>
#elif defined(__HAIKU__)
    #include <OS.h>
    #include <image.h>
#endif

#ifndef environ
extern char** environ;
#endif

enum { FF_PIPE_BUFSIZ = 8192 };

static inline int ffPipe2(int* fds, int flags)
{
    #ifndef FF_HAVE_PIPE2
        if(pipe(fds) == -1)
            return -1;
        fcntl(fds[0], F_SETFL, fcntl(fds[0], F_GETFL) | flags);
        fcntl(fds[1], F_SETFL, fcntl(fds[1], F_GETFL) | flags);
        return 0;
    #else
        return pipe2(fds, flags);
    #endif
}


// Not thread-safe
const char* ffProcessSpawn(char* const argv[], bool useStdErr, FFProcessHandle* outHandle)
{
    int pipes[2];
    if(ffPipe2(pipes, O_CLOEXEC) == -1)
        return "pipe() failed";

    pid_t childPid = -1;
    int nullFile = ffGetNullFD();

    #if !(__ANDROID__ || __OpenBSD__)

    // NetBSD / Darwin: native syscall
    // Linux (glibc): clone3-execve
    // FreeBSD: vfork-execve
    // illumos: vforkx-execve
    // OpenBSD / Android (bionic): fork-execve

    posix_spawn_file_actions_t file_actions;
    posix_spawn_file_actions_init(&file_actions);
    posix_spawn_file_actions_adddup2(&file_actions, pipes[1], useStdErr ? STDERR_FILENO : STDOUT_FILENO);
    posix_spawn_file_actions_adddup2(&file_actions, nullFile, useStdErr ? STDOUT_FILENO : STDERR_FILENO);

    static char* oldLang = NULL;
    static int langIndex = -1;

    if (langIndex >= 0)
    {
        // Found before
        if (oldLang) // oldLang was set only if it needed to be changed
        {
            if (environ[langIndex] != oldLang)
            {
                // environ is changed outside of this function
                langIndex = -1;
            }
            else
                environ[langIndex] = (char*) "LANG=C.UTF-8";
        }
    }
    if (langIndex < 0)
    {
        for (int i = 0; environ[i] != NULL; i++)
        {
            if (ffStrStartsWith(environ[i], "LANG="))
            {
                langIndex = i;
                const char* langValue = environ[i] + 5; // Skip "LANG="
                if (ffStrEqualsIgnCase(langValue, "C") ||
                    ffStrStartsWithIgnCase(environ[i], "C.") ||
                    ffStrEqualsIgnCase(langValue, "en_US") ||
                    ffStrStartsWithIgnCase(langValue, "en_US."))
                    break; // No need to change LANG
                oldLang = environ[i];
                environ[i] = (char*) "LANG=C.UTF-8"; // Set LANG to C.UTF-8 for consistent output
                break;
            }
        }
    }

    int ret = posix_spawnp(&childPid, argv[0], &file_actions, NULL, argv, environ);

    if (oldLang)
        environ[langIndex] = oldLang;

    posix_spawn_file_actions_destroy(&file_actions);

    if (ret != 0)
    {
        close(pipes[0]);
        close(pipes[1]);
        return "posix_spawnp() failed";
    }

    #else

    // https://github.com/termux/termux-packages/issues/25369
    childPid = fork();
    if(childPid == -1)
    {
        close(pipes[0]);
        close(pipes[1]);
        return "fork() failed";
    }

    if(childPid == 0)
    {
        // Child process
        dup2(pipes[1], useStdErr ? STDERR_FILENO : STDOUT_FILENO);
        dup2(nullFile, useStdErr ? STDOUT_FILENO : STDERR_FILENO);
        putenv("LANG=C.UTF-8");
        execvp(argv[0], argv);
        _exit(127);
    }

    #endif

    close(pipes[1]);
    outHandle->pid = childPid;
    outHandle->pipeRead = pipes[0];
    return NULL;
}

const char* ffProcessReadOutput(FFProcessHandle* handle, FFstrbuf* buffer)
{
    assert(handle->pipeRead != -1);
    assert(handle->pid != -1);

    const int32_t timeout = instance.config.general.processingTimeout;
    FF_AUTO_CLOSE_FD int childPipeFd = handle->pipeRead;
    pid_t childPid = handle->pid;
    handle->pipeRead = -1;
    handle->pid = -1;
    char str[FF_PIPE_BUFSIZ];

    for (;;)
    {
        if (timeout >= 0)
        {
            struct pollfd pollfd = { childPipeFd, POLLIN, 0 };
            int pollret = poll(&pollfd, 1, timeout);
            if (pollret == 0)
            {
                kill(childPid, SIGTERM);
                waitpid(childPid, NULL, 0);
                return "poll(&pollfd, 1, timeout) timeout (try increasing --processing-timeout)";
            }
            else if (pollret < 0 || (pollfd.revents & POLLERR))
            {
                kill(childPid, SIGTERM);
                waitpid(childPid, NULL, 0);
                return pollret < 0
                    ? "poll(&pollfd, 1, timeout) error: pollret < 0"
                    : "poll(&pollfd, 1, timeout) error: pollfd.revents & POLLERR";
            }
        }

        ssize_t nRead = read(childPipeFd, str, FF_PIPE_BUFSIZ);
        if (nRead > 0)
            ffStrbufAppendNS(buffer, (uint32_t) nRead, str);
        else if (nRead == 0)
        {
            int stat_loc = 0;
            if (childPid > 0 && waitpid(childPid, &stat_loc, 0) == childPid)
            {
                if (!WIFEXITED(stat_loc))
                    return "child process exited abnormally";
                if (WEXITSTATUS(stat_loc) == 127)
                    return "command not found";
                // We only handle 127 as an error. See `getTerminalVersionUrxvt` in `terminalshell.c`
                return NULL;
            }
            return NULL;
        }
        else if (nRead < 0)
            break;
    }

    return "read(childPipeFd, str, FF_PIPE_BUFSIZ) failed";
}

void ffProcessGetInfoLinux(pid_t pid, FFstrbuf* processName, FFstrbuf* exe, const char** exeName, FFstrbuf* exePath)
{
    assert(processName->length > 0);
    ffStrbufClear(exe);
    if (exePath) ffStrbufClear(exePath);

    #if defined(__linux__) || defined(__GNU__)

    char filePath[64];
    snprintf(filePath, sizeof(filePath), "/proc/%d/cmdline", (int)pid);

    if(ffReadFileBuffer(filePath, exe))
    {
        const char* p = exe->chars;
        uint32_t len = (uint32_t) strlen(p);

        if (len + 1 < exe->length)
        {
            const char* name = memrchr(p, '/', len);
            if (name) name++; else name = p;

            // For interpreters, try to find the real script path in the arguments
            if (ffStrStartsWith(name, "python")
                #ifndef __ANDROID__
                || ffStrEquals(name, "guile") // for shepherd
                #endif
            )
            {
                // `cmdline` always ends with a trailing '\0', and ffReadFileBuffer appends another \0
                // So `exe->chars` is always double '\0' terminated
                for (p = p + len + 1; *p && *p == '-'; p += strlen(p) + 1) // Skip arguments
                    assert(p - exe->chars < exe->allocated);
                if (*p)
                {
                    len = (uint32_t) strlen(p);
                    memmove(exe->chars, p, len + 1);
                }
            }
        }

        assert(len < exe->allocated);
        exe->length = len;
        ffStrbufTrimLeft(exe, '-'); //Login shells start with a dash
    }

    if (exePath)
    {
        snprintf(filePath, sizeof(filePath), "/proc/%d/exe", (int)pid);
        char buf[PATH_MAX];
        ssize_t length = readlink(filePath, buf, PATH_MAX - 1);
        if (length > 0) // doesn't contain trailing NUL
        {
            buf[length] = '\0';
            // When the process is a deleted executable, the resolved path is like `/usr/bin/app (deleted)`
            // But we can still access the binary via `/proc/pid/exe`. See #2136
            if (ffPathExists(buf, FF_PATHTYPE_ANY))
                ffStrbufSetNS(exePath, (uint32_t)length, buf);
        }

        if (exePath->length == 0)
            ffStrbufSetS(exePath, filePath);
    }

    #elif defined(__APPLE__)

    size_t len = 0;
    int mibs[] = { CTL_KERN, KERN_PROCARGS2, pid };
    if (sysctl(mibs, ARRAY_SIZE(mibs), NULL, &len, NULL, 0) == 0)
    {// try get arg0
        //don't know why if don't let len longer, proArgs2 and len will change during the following sysctl() in old MacOS version.
        len++;
        FF_AUTO_FREE char* const procArgs2 = malloc(len);
        if (sysctl(mibs, ARRAY_SIZE(mibs), procArgs2, &len, NULL, 0) == 0)
        {
            // https://gist.github.com/nonowarn/770696#file-getargv-c-L46
            uint32_t argc = *(uint32_t*) procArgs2;
            const char* realExePath = procArgs2 + sizeof(argc);

            const char* arg0 = memchr(realExePath, '\0', len - (size_t) (realExePath - procArgs2));
            if (exePath)
                ffStrbufSetNS(exePath, (uint32_t) (arg0 - realExePath), realExePath);

            do arg0++; while (*arg0 == '\0');
            assert(arg0 < procArgs2 + len);

            if (argc > 1)
            {
                // #977
                const char* p = strrchr(arg0, '/');
                if (p)
                    p++;
                else
                    p = arg0;
                if (ffStrStartsWithIgnCase(p, "python")) // /opt/homebrew/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/Resources/Python.app/Contents/MacOS/Python /Users/carter/.local/bin/xonsh
                    arg0 = p + strlen(p) + 1;
            }

            if (*arg0 == '-') arg0++; // Login shells

            ffStrbufSetS(exe, arg0);
        }
    }

    if (exePath || exe->length == 0)
    {
        char buf[PROC_PIDPATHINFO_MAXSIZE];
        int length = proc_pidpath(pid, buf, ARRAY_SIZE(buf));
        if (length > 0)
        {
            if (exe->length == 0)
                ffStrbufSetNS(exe, (uint32_t) length, buf);
            if (exePath)
            {
                // We don't use exec_path above as exePath because it's a relative path and can be different
                // from the actual executable being run (for example, when the original file is moved)
                ffStrbufSetNS(exePath, (uint32_t) length, buf);
            }
        }
    }

    #elif defined(__FreeBSD__) || defined(__NetBSD__)

    size_t size = ARG_MAX;
    FF_AUTO_FREE char* args = malloc(size);

    static_assert(ARG_MAX > PATH_MAX, "");

    if(exePath && sysctl(
        (int[]){CTL_KERN,
        #if __FreeBSD__
            KERN_PROC, KERN_PROC_PATHNAME, pid
        #else
            KERN_PROC_ARGS, pid, KERN_PROC_PATHNAME
        #endif
        }, 4,
        args, &size,
        NULL, 0
    ) == 0)
        ffStrbufSetNS(exePath, (uint32_t) (size - 1), args);

    size = ARG_MAX;
    if(sysctl(
        (int[]){CTL_KERN,
            #if __FreeBSD__
                KERN_PROC, KERN_PROC_ARGS, pid
            #else
                KERN_PROC_ARGS, pid, KERN_PROC_ARGV,
            #endif
        }, 4,
        args, &size,
        NULL, 0
    ) == 0)
    {
        char* arg0 = args;
        size_t arg0Len = strlen(args);
        if (size > arg0Len + 1)
        {
            char* p = (char*) memrchr(args, '/', arg0Len);
            if (p)
                p++;
            else
                p = arg0;
            if (ffStrStartsWith(p, "python")) // /usr/local/bin/python3.9 /home/carter/.local/bin/xonsh
            {
                arg0 += arg0Len + 1;
            }
        }
        if (arg0[0] == '-') arg0++;
        ffStrbufSetS(exe, arg0);
    }

    #elif defined(__sun)

    char filePath[128];
    snprintf(filePath, sizeof(filePath), "/proc/%d/psinfo", (int) pid);
    psinfo_t proc;
    if (ffReadFileData(filePath, sizeof(proc), &proc) == sizeof(proc))
    {
        const char* args = proc.pr_psargs;
        if (args[0] == '-') ++args;
        const char* end = strchr(args, ' ');
        ffStrbufSetNS(exe, end ? (uint32_t) (end - args) : (uint32_t) strlen(args), args);
    }

    if (exePath)
    {
        snprintf(filePath, sizeof(filePath), "/proc/%d/path/a.out", (int) pid);
        char buf[PATH_MAX];
        ssize_t length = readlink(filePath, buf, PATH_MAX - 1);
        if (length > 0) // doesn't contain trailing NUL
        {
            buf[length] = '\0';
            ffStrbufSetNS(exePath, (uint32_t)length, buf);
        }
    }

    #elif defined(__OpenBSD__)

    kvm_t* kd = kvm_open(NULL, NULL, NULL, KVM_NO_FILES, NULL);
    int count = 0;
    const struct kinfo_proc* proc = kvm_getprocs(kd, KERN_PROC_PID, pid, sizeof(struct kinfo_proc), &count);
    if (proc)
    {
        char** argv = kvm_getargv(kd, proc, 0);
        if (argv)
        {
            const char* arg0 = argv[0];
            if (arg0[0] == '-') arg0++;
            ffStrbufSetS(exe, arg0);
        }
    }
    kvm_close(kd);

    #elif defined(__HAIKU__)

    image_info info;
    int32 cookie = 0;

    while (get_next_image_info(pid, &cookie, &info) == B_OK)
    {
        if (info.type != B_APP_IMAGE) continue;
        ffStrbufSetS(exe, info.name);

        if (exePath)
            ffStrbufSet(exePath, exe);
        break;
    }

    #endif

    if(exe->length == 0)
        ffStrbufSet(exe, processName);

    assert(exe->length > 0);
    uint32_t lastSlashIndex = ffStrbufLastIndexC(exe, '/');
    if(lastSlashIndex < exe->length)
        *exeName = exe->chars + lastSlashIndex + 1;
}

const char* ffProcessGetBasicInfoLinux(pid_t pid, FFstrbuf* name, pid_t* ppid, int32_t* tty)
{
    if (pid <= 0)
        return "Invalid pid";

    #if defined(__linux__) || defined(__GNU__)

    char procFilePath[64];
    #if __linux__
    if (ppid || tty)
    #endif
    {
        snprintf(procFilePath, sizeof(procFilePath), "/proc/%d/stat", (int)pid);
        char buf[PROC_FILE_BUFFSIZ];
        ssize_t nRead = ffReadFileData(procFilePath, sizeof(buf) - 1, buf);
        if(nRead <= 8)
            return "ffReadFileData(/proc/pid/stat, PROC_FILE_BUFFSIZ-1, buf) failed";
        buf[nRead] = '\0'; // pid (comm) state ppid pgrp session tty

        const char* pState = NULL;

        {
            // comm in `/proc/pid/stat` is not encoded, and may contain ' ', ')' or even `\n`
            const char* start = memchr(buf, '(', (size_t) nRead);
            if (!start)
                return "memchr(stat, '(') failed";
            start++;
            const char* end = memrchr(start, ')', (size_t) nRead - (size_t) (start - buf));
            if (!end)
                return "memrchr(stat, ')') failed";
            ffStrbufSetNS(name, (uint32_t) (end - start), start);
            ffStrbufTrimRightSpace(name);
            if (name->chars[0] == '\0')
                return "process name is empty";
            pState = end + 2; // skip ") "
        }

        #if !__linux__
        if (ppid || tty)
        #endif
        {
            int ppid_, tty_;
            if(sscanf(pState + 2, "%d %*d %*d %d", &ppid_, &tty_) < 2)
                return "sscanf(stat) failed";

            if (ppid)
                *ppid = (pid_t) ppid_;
            if (tty)
                *tty = tty_ & 0xFF;
        }
    }
    #if __linux__
    else
    {
        snprintf(procFilePath, sizeof(procFilePath), "/proc/%d/comm", (int)pid);
        ssize_t nRead = ffReadFileBuffer(procFilePath, name);
        if(nRead <= 0)
            return "ffReadFileBuffer(/proc/pid/comm, name) failed";
        ffStrbufTrimRightSpace(name);
    }
    #endif

    #elif defined(__APPLE__)

    struct kinfo_proc proc;
    size_t size = sizeof(proc);
    if(sysctl(
        (int[]){CTL_KERN, KERN_PROC, KERN_PROC_PID, pid}, 4,
        &proc, &size,
        NULL, 0
    ))
        return "sysctl(KERN_PROC_PID) failed";

    ffStrbufSetS(name, proc.kp_proc.p_comm); //trancated to 16 chars
    if (ppid)
        *ppid = (pid_t)proc.kp_eproc.e_ppid;
    if (tty)
    {
        *tty = ((proc.kp_eproc.e_tdev >> 24) & 0xFF) == 0x10
            ? proc.kp_eproc.e_tdev & 0xFFFFFF
            : -1;
    }

    #elif defined(__FreeBSD__)

    #ifdef __DragonFly__
        #define ki_comm kp_comm
        #define ki_ppid kp_ppid
        #define ki_tdev kp_tdev
        #define ki_flag kp_flags
    #endif

    struct kinfo_proc proc;
    size_t size = sizeof(proc);
    if(sysctl(
        (int[]){CTL_KERN, KERN_PROC, KERN_PROC_PID, pid}, 4,
        &proc, &size,
        NULL, 0
    ))
        return "sysctl(KERN_PROC_PID) failed";

    ffStrbufSetS(name, proc.ki_comm);
    if (ppid)
        *ppid = (pid_t)proc.ki_ppid;
    if (tty)
    {
        if (proc.ki_tdev != NODEV && proc.ki_flag & P_CONTROLT)
        {
            const char* ttyName = devname(proc.ki_tdev, S_IFCHR);
            if (ffStrStartsWith(ttyName, "pts/"))
                *tty = (int32_t) strtol(ttyName + strlen("pts/"), NULL, 10);
            else
                *tty = -1;
        }
        else
            *tty = -1;
    }

    #elif defined(__NetBSD__)

    struct kinfo_proc2 proc;
    size_t size = sizeof(proc);
    if(sysctl(
        (int[]){CTL_KERN, KERN_PROC2, KERN_PROC_PID, pid, sizeof(proc), 1}, 6,
        &proc, &size,
        NULL, 0
    ) != 0)
        return "sysctl(KERN_PROC_PID) failed";

    ffStrbufSetS(name, proc.p_comm);
    if (ppid)
        *ppid = (pid_t)proc.p_ppid;
    if (tty)
    {
        if (proc.p_flag & P_CONTROLT)
        {
            const char* ttyName = devname(proc.p_tdev, S_IFCHR);
            if (ffStrStartsWith(ttyName, "pts/"))
                *tty = (int32_t) strtol(ttyName + strlen("pts/"), NULL, 10);
            else
                *tty = -1;
        }
        else
            *tty = -1;
    }

    #elif defined(__sun)
    char path[128];
    snprintf(path, sizeof(path), "/proc/%d/psinfo", (int) pid);
    psinfo_t proc;
    if (ffReadFileData(path, sizeof(proc), &proc) != sizeof(proc))
        return "ffReadFileData(psinfo) failed";

    ffStrbufSetS(name, proc.pr_fname);
    if (ppid)
        *ppid = proc.pr_ppid;
    if (tty)
        *tty = (int) proc.pr_ttydev;

    #elif defined(__OpenBSD__)

    kvm_t* kd = kvm_open(NULL, NULL, NULL, KVM_NO_FILES, NULL);
    int count = 0;
    const struct kinfo_proc* proc = kvm_getprocs(kd, KERN_PROC_PID, pid, sizeof(struct kinfo_proc), &count);
    if (proc)
    {
        ffStrbufSetS(name, proc->p_comm);
        if (ppid)
            *ppid = proc->p_ppid;
        if (tty)
            *tty = (int) proc->p_tdev;
    }
    kvm_close(kd);
    if (!proc)
        return "kvm_getprocs() failed";

    #elif defined(__HAIKU__)

    team_info info;
    if (get_team_info(pid, &info) == B_OK)
    {
        ffStrbufSetS(name, info.name);
        if (ppid)
            *ppid = info.parent;
    }

    FF_UNUSED(tty);

    #else

    return "Unsupported platform";

    #endif

    return NULL;
}