/**
 * @file startui.cpp
 * @author apio (cloudapio.eu)
 * @brief Service to start and manage a UI session.
 *
 * @copyright Copyright (c) 2024, the Luna authors.
 *
 */

#include <errno.h>
#include <grp.h>
#include <os/ArgumentParser.h>
#include <os/File.h>
#include <os/FileSystem.h>
#include <os/IPC.h>
#include <os/Main.h>
#include <os/Process.h>
#include <pwd.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <unistd.h>

Result<Vector<gid_t>> read_supplementary_groups(const char* name)
{
    Vector<gid_t> extra_groups;

    setgrent();
    group* grp;
    while ((grp = getgrent()))
    {
        for (char** user = grp->gr_mem; *user; user++)
        {
            if (!strcmp(*user, name))
            {
                TRY(extra_groups.try_append(grp->gr_gid));
                break;
            }
        }
    }
    endgrent();

    return extra_groups;
}

Result<void> spawn_process_as_user(Slice<StringView> arguments, uid_t user, gid_t group,
                                   Slice<gid_t> supplementary_groups)
{
    pid_t child = TRY(os::Process::fork());
    if (child == 0)
    {
        if (setgroups(static_cast<int>(supplementary_groups.size()), supplementary_groups.data()) < 0)
            return err(errno);
        setgid(group);
        setuid(user);
        TRY(os::Process::exec(arguments[0], arguments, false));
    }
    return {};
}

Result<int> luna_main(int argc, char** argv)
{
    StringView username;

    os::ArgumentParser parser;
    parser.add_description("Service to start and manage a UI session.");
    parser.add_system_program_info("startui"_sv);
    parser.add_value_argument(username, 'u', "user", "the user to start the UI session as");
    parser.parse(argc, argv);

    if (geteuid() != 0)
    {
        os::eprintln("error: %s can only be started as root.", argv[0]);
        return 1;
    }

    if (username.is_empty())
    {
        os::eprintln("error: startui needs a --user parameter to run.");
        parser.short_usage(argv[0]);
        return 1;
    }

    struct passwd* pw = getpwnam(username.chars());
    if (!pw)
    {
        os::eprintln("error: user %s not found.", username.chars());
        return 1;
    }

    auto groups = TRY(read_supplementary_groups(username.chars()));

    auto system_groups = TRY(groups.shallow_copy());
    TRY(system_groups.try_append(3));

    setsid();

    // First of all, start the display server, in case we haven't been started by loginui.
    if (!os::FileSystem::exists("/tmp/wind.sock"))
    {
        // We need to wait for this one, since its sockets are required later.
        bool success = os::IPC::Notifier::run_and_wait(
            [&] {
                StringView wind_command[] = { "/usr/bin/wind" };
                os::Process::spawn(wind_command[0], Slice<StringView>(wind_command, 1));
            },
            1000);
        if (!success)
        {
            os::eprintln("startui: failed to start wind, timed out");
            return 1;
        }
    }

    clearenv();
    chdir(pw->pw_dir);
    setenv("PATH", "/usr/bin:/usr/local/bin", 1);
    setenv("USER", pw->pw_name, 1);
    setenv("HOME", pw->pw_dir, 1);
    setenv("SHELL", pw->pw_shell, 1);

    // We also need to wait for this one, since taskbar requires launch.sock.
    bool success = os::IPC::Notifier::run_and_wait(
        [&] {
            StringView launch_command[] = { "/usr/bin/launch" };
            spawn_process_as_user(Slice<StringView>(launch_command, 1), pw->pw_uid, pw->pw_gid, groups.slice());
        },
        1000);
    if (!success)
    {
        os::eprintln("startui: failed to start launch server, timed out");
        return 1;
    }

    StringView taskbar_command[] = { "/usr/bin/taskbar" };
    TRY(spawn_process_as_user(Slice<StringView>(taskbar_command, 1), pw->pw_uid, pw->pw_gid, system_groups.slice()));

    // Finally, start init in user mode to manage all configured optional services.
    StringView init_command[] = { "/usr/bin/init", "--user" };
    TRY(spawn_process_as_user(Slice<StringView>(init_command, 2), pw->pw_uid, pw->pw_gid, groups.slice()));

    // If any process exits, it's either the display server, one of the required UI components, or init --user. We care
    // about all of these, so let's just crash the entire UI session if they exit. In the future we might restart some
    // of them and continue, but for now we'll just exit on all of them.
    os::Process::wait(os::Process::ANY_CHILD, nullptr);

    return 0;
}