// SPDX-License-Identifier: AGPL-3.0-only
// SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

module;
#include <fcntl.h>
#include <linux/close_range.h>
#include <poll.h>
#include <signal.h>
#include <sys/prctl.h>
#include <sys/syscall.h>
#include <sys/wait.h>
#include <unistd.h>
#include <cerrno>

module Crafter.Review:Process_impl;
import :Process;
import :Platform;
import :Text;
import std;

namespace fs = std::filesystem;

namespace Crafter::Review::Process {
    namespace {
        // ---- the reaper's registry and the spawner's queue ------------------------
        struct Entry {
            std::optional<std::int32_t> status;
        };
        struct Request;
        struct Shared {
            std::mutex registryMutex;
            std::condition_variable registryChanged;
            std::map<pid_t, Entry> registry;
            std::int32_t unreaped = 0;
            std::mutex spawnMutex;
            std::condition_variable spawnReady;
            std::deque<Request*> spawnQueue;
        };
        // Never destroyed. The reaper and the spawner are detached threads
        // that wait on these until the process ends, and a condition variable
        // destroyed under a waiter blocks exit(): glibc's pthread_cond_destroy
        // waits for its waiters, which here never come back.
        Shared& State() {
            static Shared* state = new Shared;
            return *state;
        }

        std::int32_t ReturnCode(std::int32_t status) {
            if (WIFEXITED(status)) return WEXITSTATUS(status);
            if (WIFSIGNALED(status)) return -WTERMSIG(status);
            return -1;
        }

        void ReaperLoop() {
            while (true) {
                std::int32_t status = 0;
                const pid_t pid = waitpid(-1, &status, 0);
                if (pid > 0) {
                    const std::lock_guard lock(State().registryMutex);
                    // A pid nobody registered is an orphan of one of our
                    // children, reparented to us as the subreaper: reaped,
                    // and that is all.
                    if (const auto found = State().registry.find(pid); found != State().registry.end() && !found->second.status) {
                        found->second.status = ReturnCode(status);
                        --State().unreaped;
                        State().registryChanged.notify_all();
                    }
                    continue;
                }
                if (errno == EINTR) continue;
                // ECHILD: no child at all. Wait for Start() to register one;
                // wake now and then for orphans reparented in the meantime.
                std::unique_lock lock(State().registryMutex);
                State().registryChanged.wait_for(lock, std::chrono::seconds(1), [] { return State().unreaped > 0; });
            }
        }

        // ---- the spawner -------------------------------------------------------------
        struct Request {
            const Spec* spec;
            std::promise<Child> result;
        };
        std::once_flag ThreadsStarted;

        std::optional<std::string> FindOnPath(const std::string& name, const std::string& path) {
            if (name.contains('/')) return name;
            for (const std::string& dir : Dyn::Split(path, ":")) {
                const std::string candidate = Dyn::JoinPath(dir.empty() ? std::string(".") : dir, name);
                if (access(candidate.c_str(), X_OK) == 0) return candidate;
            }
            return std::nullopt;
        }

        std::int32_t OpenStream(const Stream& stream, bool reading, std::int32_t& ours) {
            switch (stream.kind) {
                case Stream::Kind::Inherit:
                    return -1;
                case Stream::Kind::Null:
                    return open("/dev/null", (reading ? O_RDONLY : O_WRONLY) | O_CLOEXEC);
                case Stream::Kind::File:
                    return reading ? open(stream.file.c_str(), O_RDONLY | O_CLOEXEC)
                                   : open(stream.file.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
                case Stream::Kind::Pipe: {
                    std::array<std::int32_t, 2> ends {};
                    if (pipe2(ends.data(), O_CLOEXEC) != 0) return -2;
                    ours = reading ? ends[1] : ends[0];
                    return reading ? ends[0] : ends[1];
                }
            }
            return -1;
        }

        Child DoSpawn(const Spec& spec) {
            if (spec.argv.empty()) throw std::runtime_error("Start: no argv");
            // Everything the child needs is made here: between fork and exec
            // it may only make async-signal-safe calls.
            std::vector<std::string> envStrings;
            std::string pathVariable = "/usr/local/bin:/usr/bin:/bin";
            if (spec.environment) {
                for (const auto& [name, value] : *spec.environment) envStrings.push_back(std::format("{}={}", name, value));
                if (const auto found = spec.environment->find("PATH"); found != spec.environment->end()) pathVariable = found->second;
            } else {
                for (char** entry = environ; entry != nullptr && *entry != nullptr; ++entry) envStrings.emplace_back(*entry);
                if (const std::optional<std::string> path = Platform::GetEnv("PATH")) pathVariable = *path;
            }
            const std::optional<std::string> program = FindOnPath(spec.argv.front(), pathVariable);
            if (!program) {
                throw std::runtime_error(std::format("{}: {}", Dyn::ErrnoText(ENOENT), Dyn::Repr(spec.argv.front())));
            }
            std::vector<char*> argv;
            for (const std::string& arg : spec.argv) argv.push_back(const_cast<char*>(arg.c_str()));
            argv.push_back(nullptr);
            std::vector<char*> envp;
            for (const std::string& entry : envStrings) envp.push_back(const_cast<char*>(entry.c_str()));
            envp.push_back(nullptr);
            const std::string cwd = spec.cwd ? spec.cwd->string() : std::string();

            Child child;
            const std::int32_t in = OpenStream(spec.input, true, child.input);
            const std::int32_t out = OpenStream(spec.output, false, child.output);
            const std::int32_t err = OpenStream(spec.error, false, child.error);
            auto closeAll = [&] {
                for (const std::int32_t fd : { in, out, err, child.input, child.output, child.error }) {
                    if (fd >= 0) close(fd);
                }
            };
            if (in == -2 || out == -2 || err == -2 || (spec.input.kind != Stream::Kind::Inherit && in < 0)
                || (spec.output.kind != Stream::Kind::Inherit && out < 0) || (spec.error.kind != Stream::Kind::Inherit && err < 0)) {
                const std::int32_t failure = errno;
                closeAll();
                throw std::runtime_error(std::format("{}: {}", Dyn::ErrnoText(failure), Dyn::Repr(spec.argv.front())));
            }
            std::array<std::int32_t, 2> failurePipe {};
            if (pipe2(failurePipe.data(), O_CLOEXEC) != 0) {
                closeAll();
                throw std::runtime_error(Dyn::ErrnoText(errno));
            }
            sigset_t empty;
            sigemptyset(&empty);
            const pid_t parent = getpid();
            const std::int32_t deathSignal = spec.deathSignal;

            std::unique_lock registry(State().registryMutex);
            const pid_t pid = fork();
            if (pid == 0) {
                struct sigaction defaults {};
                defaults.sa_handler = SIG_DFL;
                sigaction(SIGPIPE, &defaults, nullptr);
                sigaction(SIGXFSZ, &defaults, nullptr);
                sigprocmask(SIG_SETMASK, &empty, nullptr);
                prctl(PR_SET_PDEATHSIG, deathSignal);
                if (getppid() != parent) _exit(127);
                if (in >= 0) dup2(in, 0);
                if (out >= 0) dup2(out, 1);
                if (err >= 0) dup2(err, 2);
                std::int32_t failure = 0;
                if (!cwd.empty() && chdir(cwd.c_str()) != 0) {
                    failure = errno;
                } else {
                    // Nothing but the three streams reaches the program: a lock
                    // or a socket of the driver's held open by a child would
                    // outlive the driver.
                    syscall(SYS_close_range, 3U, ~0U, CLOSE_RANGE_CLOEXEC);
                    execve(program->c_str(), argv.data(), envp.data());
                    failure = errno;
                }
                (void)!write(failurePipe[1], &failure, sizeof(failure));
                _exit(127);
            }
            if (pid < 0) {
                const std::int32_t failure = errno;
                registry.unlock();
                closeAll();
                close(failurePipe[0]);
                close(failurePipe[1]);
                throw std::runtime_error(std::format("fork: {}", Dyn::ErrnoText(failure)));
            }
            State().registry[pid] = {};
            ++State().unreaped;
            registry.unlock();
            State().registryChanged.notify_all();

            for (const std::int32_t fd : { in, out, err }) {
                if (fd >= 0) close(fd);
            }
            close(failurePipe[1]);
            std::int32_t failure = 0;
            ssize_t got = 0;
            do {
                got = read(failurePipe[0], &failure, sizeof(failure));
            } while (got < 0 && errno == EINTR);
            close(failurePipe[0]);
            if (got == static_cast<ssize_t>(sizeof(failure))) {
                Wait(pid);
                for (const std::int32_t fd : { child.input, child.output, child.error }) {
                    if (fd >= 0) close(fd);
                }
                throw std::runtime_error(std::format("{}: {}", Dyn::ErrnoText(failure), Dyn::Repr(spec.argv.front())));
            }
            child.pid = pid;
            return child;
        }

        void SpawnerLoop() {
            while (true) {
                Request* request = nullptr;
                {
                    std::unique_lock lock(State().spawnMutex);
                    State().spawnReady.wait(lock, [] { return !State().spawnQueue.empty(); });
                    request = State().spawnQueue.front();
                    State().spawnQueue.pop_front();
                }
                try {
                    request->result.set_value(DoSpawn(*request->spec));
                } catch (...) {
                    request->result.set_exception(std::current_exception());
                }
            }
        }

        void EnsureThreads() {
            std::call_once(ThreadsStarted, [] {
                std::thread(ReaperLoop).detach();
                std::thread(SpawnerLoop).detach();
            });
        }
    }

    Child Start(const Spec& spec) {
        EnsureThreads();
        Request request { .spec = &spec };
        std::future<Child> result = request.result.get_future();
        {
            const std::lock_guard lock(State().spawnMutex);
            State().spawnQueue.push_back(&request);
        }
        State().spawnReady.notify_one();
        return result.get();
    }

    std::optional<std::int32_t> Wait(std::int64_t pid, std::optional<std::chrono::milliseconds> timeout) {
        std::unique_lock lock(State().registryMutex);
        auto exited = [pid] {
            const auto found = State().registry.find(static_cast<pid_t>(pid));
            return found == State().registry.end() || found->second.status.has_value();
        };
        if (timeout) {
            if (!State().registryChanged.wait_for(lock, *timeout, exited)) return std::nullopt;
        } else {
            State().registryChanged.wait(lock, exited);
        }
        const auto found = State().registry.find(static_cast<pid_t>(pid));
        if (found == State().registry.end()) return -1;
        const std::int32_t code = *found->second.status;
        State().registry.erase(found);
        return code;
    }

    void Signal(std::int64_t pid, std::int32_t signal) {
        if (pid > 0) kill(static_cast<pid_t>(pid), signal);
    }

    void KillTree(std::int64_t root, bool includeRoot) {
        std::set<std::int64_t> stopped;
        for (std::int32_t round = 0; round < 50; ++round) {
            const auto kids = Platform::PpidMap();
            bool added = false;
            for (const std::int64_t pid : Platform::Tree(root, kids)) {
                if (pid == root && !includeRoot) continue;
                if (pid <= 1) continue;
                if (stopped.insert(pid).second) {
                    kill(static_cast<pid_t>(pid), SIGSTOP);
                    added = true;
                }
            }
            if (!added) break;
        }
        for (const std::int64_t pid : stopped) kill(static_cast<pid_t>(pid), SIGKILL);
    }

    void KillDescendants() {
        KillTree(getpid(), false);
    }

    Captured Run(const std::vector<std::string>& argv, const std::optional<fs::path>& cwd,
                 std::optional<std::chrono::milliseconds> timeout, std::string_view input,
                 const std::optional<std::map<std::string, std::string>>& environment) {
        Spec spec { .argv = argv, .environment = environment, .cwd = cwd };
        spec.input.kind = input.empty() ? Stream::Kind::Null : Stream::Kind::Pipe;
        spec.output.kind = Stream::Kind::Pipe;
        spec.error.kind = Stream::Kind::Pipe;
        Child child = Start(spec);
        Captured captured;
        const auto deadline = timeout ? std::optional(std::chrono::steady_clock::now() + *timeout) : std::nullopt;
        std::size_t sent = 0;
        if (child.input >= 0) fcntl(child.input, F_SETFL, fcntl(child.input, F_GETFL) | O_NONBLOCK);
        std::array<char, 65536> buffer {};
        while (child.output >= 0 || child.error >= 0 || child.input >= 0) {
            std::vector<pollfd> watched;
            if (child.input >= 0) watched.push_back({ .fd = child.input, .events = POLLOUT, .revents = 0 });
            if (child.output >= 0) watched.push_back({ .fd = child.output, .events = POLLIN, .revents = 0 });
            if (child.error >= 0) watched.push_back({ .fd = child.error, .events = POLLIN, .revents = 0 });
            std::int32_t waitMs = -1;
            if (deadline) {
                const auto left = std::chrono::duration_cast<std::chrono::milliseconds>(*deadline - std::chrono::steady_clock::now()).count();
                if (left <= 0) {
                    Signal(child.pid, SIGKILL);
                    captured.timedOut = true;
                    break;
                }
                waitMs = static_cast<std::int32_t>(std::min<std::int64_t>(left, 1000));
            }
            const std::int32_t ready = poll(watched.data(), watched.size(), waitMs);
            if (ready < 0 && errno != EINTR) break;
            for (const pollfd& entry : watched) {
                if (entry.revents == 0) continue;
                if (entry.fd == child.input) {
                    if (entry.revents & (POLLERR | POLLHUP)) {
                        close(child.input);
                        child.input = -1;
                        continue;
                    }
                    const ssize_t wrote = write(child.input, input.data() + sent, input.size() - sent);
                    if (wrote > 0) sent += static_cast<std::size_t>(wrote);
                    if (sent >= input.size() || (wrote < 0 && errno != EAGAIN && errno != EINTR)) {
                        close(child.input);
                        child.input = -1;
                    }
                    continue;
                }
                const ssize_t got = read(entry.fd, buffer.data(), buffer.size());
                if (got > 0) {
                    (entry.fd == child.output ? captured.out : captured.err).append(buffer.data(), static_cast<std::size_t>(got));
                } else if (got == 0 || (errno != EINTR && errno != EAGAIN)) {
                    close(entry.fd);
                    (entry.fd == child.output ? child.output : child.error) = -1;
                }
            }
        }
        for (const std::int32_t fd : { child.input, child.output, child.error }) {
            if (fd >= 0) close(fd);
        }
        captured.returnCode = Wait(child.pid).value_or(-1);
        return captured;
    }

    void InstallSignalThread(std::function<void(std::int32_t)> onSignal) {
        sigset_t set;
        sigemptyset(&set);
        sigaddset(&set, SIGINT);
        sigaddset(&set, SIGTERM);
        sigaddset(&set, SIGHUP);
        pthread_sigmask(SIG_BLOCK, &set, nullptr);
        struct sigaction ignore {};
        ignore.sa_handler = SIG_IGN;
        sigaction(SIGPIPE, &ignore, nullptr);
        std::thread([set, onSignal = std::move(onSignal)] {
            while (true) {
                std::int32_t signal = 0;
                if (sigwait(&set, &signal) == 0) onSignal(signal);
            }
        }).detach();
    }

    void BecomeSubreaper() {
        prctl(PR_SET_CHILD_SUBREAPER, 1);
    }

    std::vector<std::string> Supervised(std::span<const std::string> argv) {
        std::vector<std::string> out { Platform::SelfExe().string(), std::string(SuperviseVerb) };
        out.insert(out.end(), argv.begin(), argv.end());
        return out;
    }

    std::int32_t SupervisedReturnCode(std::int32_t supervisorCode) {
        if (supervisorCode > 128 && supervisorCode <= 128 + 64) return -(supervisorCode - 128);
        return supervisorCode;
    }

    std::int32_t SuperviseMain(std::span<const std::string_view> args) {
        if (args.empty()) {
            std::println(std::cerr, "{}: nothing to run", SuperviseVerb);
            return 64;
        }
        prctl(PR_SET_CHILD_SUBREAPER, 1);
        sigset_t watched;
        sigemptyset(&watched);
        sigaddset(&watched, SIGTERM);
        sigaddset(&watched, SIGINT);
        sigaddset(&watched, SIGHUP);
        sigaddset(&watched, SIGCHLD);
        sigprocmask(SIG_BLOCK, &watched, nullptr);
        sigset_t empty;
        sigemptyset(&empty);

        std::vector<std::string> strings(args.begin(), args.end());
        std::vector<char*> argv;
        for (const std::string& arg : strings) argv.push_back(const_cast<char*>(arg.c_str()));
        argv.push_back(nullptr);
        const std::optional<std::string> program = FindOnPath(strings.front(), Platform::GetEnv("PATH").value_or("/usr/bin:/bin"));
        if (!program) {
            std::println(std::cerr, "{}: {}", Dyn::ErrnoText(ENOENT), Dyn::Repr(strings.front()));
            return 127;
        }
        const pid_t self = getpid();
        const pid_t target = fork();
        if (target == 0) {
            struct sigaction defaults {};
            defaults.sa_handler = SIG_DFL;
            sigaction(SIGPIPE, &defaults, nullptr);
            sigaction(SIGXFSZ, &defaults, nullptr);
            sigprocmask(SIG_SETMASK, &empty, nullptr);
            prctl(PR_SET_PDEATHSIG, SIGKILL);
            if (getppid() != self) _exit(127);
            execv(program->c_str(), argv.data());
            _exit(127);
        }
        if (target < 0) return 127;

        std::optional<std::int32_t> status;
        while (!status) {
            siginfo_t info {};
            const std::int32_t signal = sigwaitinfo(&watched, &info);
            if (signal < 0) continue;
            if (signal == SIGCHLD) {
                while (true) {
                    std::int32_t raw = 0;
                    const pid_t pid = waitpid(-1, &raw, WNOHANG);
                    if (pid <= 0) break;
                    if (pid == target) status = raw;
                }
                continue;
            }
            // A timeout, an abort, or the driver gone: the whole tree goes,
            // the target with it, and its status is what that left.
            KillTree(self, false);
            std::int32_t raw = 0;
            while (waitpid(target, &raw, 0) < 0 && errno == EINTR) {}
            status = raw;
        }
        // The target is gone; whatever it started goes too, and every orphan
        // that was reparented to this subreaper is reaped.
        KillTree(self, false);
        while (true) {
            const pid_t pid = waitpid(-1, nullptr, 0);
            if (pid < 0 && errno == EINTR) continue;
            if (pid < 0) break;
        }
        if (WIFEXITED(*status)) return WEXITSTATUS(*status);
        if (WIFSIGNALED(*status)) return 128 + WTERMSIG(*status);
        return 1;
    }
}