#include "cgroup.hpp"

#include <unistd.h>

#include <chrono>
#include <sstream>
#include <thread>

#include "spec.hpp"
#include "util.hpp"

namespace fs = std::filesystem;

namespace doos {

namespace {
bool hasController(const fs::path& dir, const std::string& name) {
    try {
        std::istringstream in(fsutil::readFile(dir / "cgroup.controllers"));
        for (std::string c; in >> c;) {
            if (c == name) return true;
        }
    } catch (const std::exception&) {
    }
    return false;
}

bool tryWrite(const fs::path& file, const std::string& text) {
    try {
        fsutil::writeFile(file, text);
        return true;
    } catch (const std::exception&) {
        return false;
    }
}
}  // namespace

fs::path Cgroup::base(std::string& whyNot) {
    std::string self;
    try {
        self = fsutil::readFile("/proc/self/cgroup");
    } catch (const std::exception& e) {
        whyNot = e.what();
        return {};
    }
    std::string marker = "user@" + std::to_string(::geteuid()) + ".service";
    auto at = self.find(marker);
    if (self.rfind("0::", 0) != 0 || at == std::string::npos) {
        whyNot = "no cgroup v2 delegated to the user manager (" + marker + " not found)";
        return {};
    }
    return fs::path("/sys/fs/cgroup") / self.substr(3, at + marker.size() - 3).substr(1);
}

std::unique_ptr<Cgroup> Cgroup::create(const std::string& name, const Limits& limits, std::string& whyNot,
                                       std::string& notEnforced) {
    try {
        fs::path b = base(whyNot);
        if (b.empty()) return nullptr;
        fs::path top = b / std::string(spec::kCgroupDir);
        std::error_code ec;
        fs::create_directory(top, ec);
        if (ec && ec != std::errc::file_exists) {
            whyNot = "cannot create " + top.string() + ": " + ec.message();
            return nullptr;
        }
        // Controllers must be enabled for children of `doos/` for the limit files to exist in `doos/<name>`.
        std::string enable;
        for (const char* c : {"memory", "pids", "cpu"}) {
            if (hasController(b, c)) enable += std::string(" +") + c;
        }
        if (!enable.empty()) tryWrite(top / "cgroup.subtree_control", enable);

        fs::path dir = top / name;
        fs::create_directory(dir, ec);
        if (ec && ec != std::errc::file_exists) {
            whyNot = "cannot create " + dir.string() + ": " + ec.message();
            return nullptr;
        }
        for (const char* leaf : {"main", "shell"}) fs::create_directory(dir / leaf, ec);
        if (::access((dir / "main" / "cgroup.freeze").c_str(), W_OK) != 0) {
            whyNot = "cgroup.freeze not writable in " + dir.string();
            return nullptr;
        }
        auto note = [&](const std::string& what) { notEnforced += (notEnforced.empty() ? "" : ", ") + what; };
        if (limits.memoryBytes) {
            if (!tryWrite(dir / "memory.max", std::to_string(limits.memoryBytes))) note("memory");
            else {
                tryWrite(dir / "memory.swap.max", "0");  // otherwise the cap can be dodged through swap
                tryWrite(dir / "memory.oom.group", "1"); // an OOM kill takes the whole container, not one process
            }
        }
        if (limits.pids && !tryWrite(dir / "pids.max", std::to_string(limits.pids))) note("pids");
        if (limits.milliCpus) {
            long quota = static_cast<long>(limits.milliCpus) * spec::kCpuPeriodUs / 1000;
            if (!tryWrite(dir / "cpu.max", std::to_string(quota) + " " + std::to_string(spec::kCpuPeriodUs))) note("cpus");
        }
        return std::unique_ptr<Cgroup>(new Cgroup(dir, true));
    } catch (const std::exception& e) {
        whyNot = e.what();
        return nullptr;
    }
}

std::unique_ptr<Cgroup> Cgroup::open(const std::string& name) {
    std::string why;
    fs::path b = base(why);
    if (b.empty()) return nullptr;
    fs::path dir = b / std::string(spec::kCgroupDir) / name;
    std::error_code ec;
    if (!fs::is_directory(dir / "shell", ec)) return nullptr;
    return std::unique_ptr<Cgroup>(new Cgroup(dir, false));
}

Cgroup::~Cgroup() {
    if (!owner_) return;
    // Everything dies with the container's init; kill any straggler, then remove the leaves and the group.
    tryWrite(dir_ / "cgroup.kill", "1");
    for (int i = 0; i < 40; ++i) {
        std::error_code ec;
        fs::remove(dir_ / "main", ec);
        fs::remove(dir_ / "shell", ec);
        if (fs::remove(dir_, ec)) return;
        std::this_thread::sleep_for(std::chrono::milliseconds(25));
    }
}

void Cgroup::addMain(pid_t pid) const { fsutil::writeFile(dir_ / "main" / "cgroup.procs", std::to_string(pid)); }
void Cgroup::addShell(pid_t pid) const { fsutil::writeFile(dir_ / "shell" / "cgroup.procs", std::to_string(pid)); }

void Cgroup::freeze(bool on) const {
    fsutil::writeFile(dir_ / "main" / "cgroup.freeze", on ? "1" : "0");
    frozen_ = on;
}

std::uint64_t Cgroup::counter(const char* file, const char* key) const {
    try {
        std::istringstream in(fsutil::readFile(dir_ / file));
        std::string k;
        std::uint64_t v;
        while (in >> k >> v) {
            if (k == key) return v;
        }
    } catch (const std::exception&) {
    }
    return 0;
}

std::uint64_t Cgroup::oomKills() const { return counter("memory.events", "oom_kill"); }
std::uint64_t Cgroup::pidFailures() const { return counter("pids.events", "max"); }

}  // namespace doos