/*
    SPDX-License-Identifier: GPL-2.0-only OR GPL-3.0-only OR LicenseRef-KDE-Accepted-GPL
    SPDX-FileCopyrightText: 2021 Harald Sitter <sitter@kde.org>
*/

#include <filesystem>
#include <fstream>
#include <iostream>
#include <libintl.h>
#include <regex>
#include <string>
#include <string_view>

constexpr std::uintmax_t byteFactor = 1024;
// Grub and theme is about 20MiB but we add some extra space to account for future growth.
constexpr std::uintmax_t themeSize = (20 * byteFactor * byteFactor /* MiB */);
// initrd + vmlinuz is about 110MiB per kernel. Somewhat subject to change. We also add some additional wiggle room in the code.
constexpr std::uintmax_t singleSize = (110 * byteFactor * byteFactor /* MiB */);

std::uintmax_t toMiB(std::uintmax_t bytes)
{
    return bytes / byteFactor / byteFactor;
}

template<typename... Args>
std::string format(const std::string_view &format, Args... args)
{
    constexpr int nullchar = 1; // \0
    const int size = std::snprintf(nullptr, 0, format.data(), args...) + nullchar;
    if (size <= nullchar) {
        return "Unable to determine size of description string.";
    }

    auto data = std::make_unique<char[]>( size );
    std::snprintf(data.get(), size, format.data(), args...);

    return std::string(data.get(), size);
}

static bool endsWith(const std::string_view &string, const std::string_view &suffix)
{
    if (string.size() < suffix.size()) {
        return false;
    }
    return string.compare(string.size() - suffix.size(), suffix.size(), suffix) == 0;
}

uint countKernels(const std::string_view &path)
{
    uint probableKernels = 0;

    if (std::ifstream file(path.data()); file.is_open()) {
        std::string line;
        while (std::getline(file, line)) {
            // Expression should at least have one digit after the name. Otherwise it may be a meta package.
            // We then match everything up to the next entry deliminator ';' so as to collect the entire package name.
            // This way we can then filter dbgsyms.
            static std::regex expression("linux-image-\\d+[^;]+");
            auto begin = std::sregex_iterator(line.cbegin(), line.cend(), expression);
            auto end = std::sregex_iterator();
            for (auto it = begin; it != end; ++it) {
                const std::string match = it->str();
                if (endsWith(match, "-dbgsym")) { // filter dbgsym out
                    continue;
                }
                std::cout << "Found kernel: " << match << '\n';
                ++probableKernels;
            }
        }
    } else {
        std::cerr << "Failed to open path " << path << '\n';
    }

    return probableKernels;
}

int main(int /*argc*/, char ** /*argv*/)
{
    std::setlocale(LC_ALL, ""); // apply env locale
    textdomain("neon-boot-space");

    const uint probableKernels = countKernels("/var/lib/PackageKit/prepared-update");
    if (probableKernels == 0) {
        std::cout << "Update doesn't include kernels it seems\n";
        return 0;
    }
    std::cout << "Found kernels: " << probableKernels << '\n';

    // require 2 kernel leeway space. 1 is wiggle room and 1 is required because initrds aren't rebuilt in-place
    // but rather next to the existing initrd. i.e. when rebuilding an initrd it needs twice its usual space.
    const std::uintmax_t estimate = ((2 + probableKernels)  * singleSize) + themeSize;
    const std::filesystem::space_info info = std::filesystem::space("/boot/");
    if (info.available > estimate) {
        std::cout << "Enough space. Wanted " << estimate << " has " << info.available << '\n';
        return 0;
    }

    // NB: it's indeed actually spelled competed without l
    if (std::ofstream file("/var/lib/PackageKit/offline-update-competed"); file.is_open()) {
        file << "[PackageKit Offline Update Results]\n"
             << "Success=false\n"
             << "ErrorCode=boot-out-of-space\n"
             << "ErrorDetails="
             << format(gettext("/boot does not have enough space to safely apply the update (requires %d MiB; has %d MiB)"),
                   toMiB(estimate),
                   toMiB(info.available))
            << '\n';
    } else {
        std::cerr << "Failed to open offline-update-completed file";
    }

    // Unstage update. Remove all the marker files, then reload the daemon and isolate into the actual
    // default.target to reach the desktop and report on the problem.

    for (const auto path : {"/system-update",
                            "/var/lib/PackageKit/offline-update-action",
                            "/run/systemd/generator.early/default.target"}) {
        try {
            std::filesystem::remove(path);
        } catch (std::filesystem::filesystem_error const &e) {
            // Don't abort if any of these fail. Continue cleaning up! Worst case the update continues
            // but we can't do anything about that anyway if one of the marker files fails to remove.
            std::cerr << e.what() << '\n';
        }
    }

    std::cout.flush();
    if (std::system("systemctl daemon-reload") != 0) {
        std::cerr << "daemon reload failed\n";
    }
    if (std::system("systemctl isolate --no-block default.target") != 0) {
        std::cerr << "isolate failed\n";
    }

    return 1;
}
