#include <errno.h>
#include <fcntl.h>
#include <luna/PathParser.h>
#include <luna/Sort.h>
#include <luna/String.h>
#include <luna/Utf8.h>
#include <luna/Vector.h>
#include <os/ArgumentParser.h>
#include <os/Directory.h>
#include <os/File.h>
#include <os/Process.h>
#include <os/Security.h>
#include <pwd.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/sysmacros.h>
#include <unistd.h>

static bool g_is_system = false;

FILE* g_init_log = nullptr;

struct Service
{
    String name;
    String command;
    bool restart { false };
    String environment;
    String standard_output;
    String standard_error;
    String standard_input;
    String working_directory;
    Option<uid_t> user {};
    Option<gid_t> group {};
    bool wait { false };
    Option<pid_t> pid {};
};

Vector<Service> g_services;

static void do_log(const char* format, ...)
{
    va_list ap;
    va_start(ap, format);

    if (!g_is_system) fprintf(g_init_log, "(user) ");

    if (g_init_log) vfprintf(g_init_log, format, ap);

    va_end(ap);
}

static void do_errlog(const char* format, ...)
{
    va_list ap;
    va_start(ap, format);

    fprintf(stderr, "(user) ");

    vfprintf(stderr, format, ap);

    va_end(ap);
}

// Request a successful exit from the system (for tests)
void sigterm_handler(int)
{
    do_log("[init] successful exit requested, complying\n");
    exit(0);
}

// Request a failure exit from the system (for tests)
void sigquit_handler(int)
{
    do_log("[init] failure exit requested, complying\n");
    exit(1);
}

static Result<void> service_child(const Service& service, SharedPtr<os::File> output, SharedPtr<os::File> error,
                                  SharedPtr<os::File> input)
{
    auto args = TRY(service.command.split(" \n"));

    if (output) dup2(output->fd(), STDOUT_FILENO);
    if (error) dup2(error->fd(), STDERR_FILENO);
    if (input) dup2(input->fd(), STDIN_FILENO);

    if (!service.working_directory.is_empty()) chdir(service.working_directory.chars());

    if (service.user.has_value())
    {
        setgid(service.group.value());
        setuid(service.user.value());
    }

    if (service.environment.is_empty()) { TRY(os::Process::exec(args[0].view(), args.slice(), false)); }
    else
    {
        auto env = TRY(service.environment.split(",\n"));
        TRY(os::Process::exec(args[0].view(), args.slice(), env.slice(), false));
    }

    return {};
}

static Result<void> try_start_service(Service& service)
{
    SharedPtr<os::File> new_stdout = {};
    SharedPtr<os::File> new_stderr = {};
    SharedPtr<os::File> new_stdin = {};

    if (!service.standard_output.is_empty())
    {
        new_stdout = TRY(os::File::open_or_create(service.standard_output.view(), os::File::Append, 0600));
        new_stdout->set_close_on_exec();
    }

    if (!service.standard_error.is_empty())
    {
        new_stderr = TRY(os::File::open_or_create(service.standard_error.view(), os::File::Append, 0600));
        new_stderr->set_close_on_exec();
    }

    if (!service.standard_input.is_empty())
    {
        new_stdin = TRY(os::File::open(service.standard_input.view(), os::File::ReadOnly));
        new_stdin->set_close_on_exec();
    }

    pid_t pid = TRY(os::Process::fork());
    if (pid == 0)
    {
        auto rc = service_child(service, new_stdout, new_stderr, new_stdin);
        if (rc.has_error())
        {
            do_errlog("[child %d] failed to start service %s due to error: %s\n", getpid(), service.name.chars(),
                      rc.error_string());
        }
        fclose(g_init_log);
        exit(127);
    }

    do_log("[init] created new child process %d for service %s\n", pid, service.name.chars());

    if (service.wait)
    {
        do_log("[init] waiting for child process %d to finish\n", pid);

        int status;
        TRY(os::Process::wait(pid, &status));

        do_log("[init] child process %d exited with code %d\n", pid, WEXITSTATUS(status));
    }
    else
        service.pid = pid;

    return {};
}

static void start_service(Service& service)
{
    auto rc = try_start_service(service);
    if (rc.has_error())
    {
        do_errlog("[init] failed to start service %s due to error: %s\n", service.name.chars(), rc.error_string());
    }
}

static Result<void> load_service(const os::Path& path)
{
    do_log("[init] reading service file: %s\n", path.name().chars());

    auto file = TRY(os::File::open(path, os::File::ReadOnly));

    Service service;

    while (true)
    {
        auto line = TRY(file->read_line());
        if (line.is_empty()) break;

        line.trim("\n");
        if (line.is_empty()) continue;

        auto parts = TRY(line.split_once('='));
        if (parts.size() < 2 || parts[0].is_empty() || parts[1].is_empty())
        {
            do_log("[init] file contains invalid line, aborting: '%s'\n", line.chars());
            return {};
        }

        if (parts[0].view() == "Name")
        {
            service.name = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "Description")
        {
            // We let users specify this in the config file, but init doesn't actually use it.
            continue;
        }

        if (parts[0].view() == "Command")
        {
            service.command = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "Restart")
        {
            if (parts[1].view() == "true" || parts[1].view().to_uint().value_or(0) == 1)
            {
                service.restart = true;
                continue;
            }
            service.restart = false;
            continue;
        }

        if (parts[0].view() == "Environment")
        {
            service.environment = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "StandardOutput")
        {
            service.standard_output = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "StandardError")
        {
            service.standard_error = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "StandardInput")
        {
            service.standard_input = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "WorkingDirectory")
        {
            service.working_directory = move(parts[1]);
            continue;
        }

        if (g_is_system && parts[0].view() == "User")
        {
            auto* pw = getpwnam(parts[1].chars());
            if (!pw) continue;
            service.user = pw->pw_uid;
            service.group = pw->pw_gid;
            continue;
        }

        if (parts[0].view() == "Wait")
        {
            if (parts[1].view() == "true" || parts[1].view().to_uint().value_or(0) == 1)
            {
                service.wait = true;
                continue;
            }
            service.wait = false;
            continue;
        }

        do_log("[init] skipping unknown entry name %s\n", parts[0].chars());
    }

    if (service.name.is_empty())
    {
        do_log("[init] service file is missing 'Name' entry, aborting!\n");
        return {};
    }

    if (service.command.is_empty())
    {
        do_log("[init] service file is missing 'Command' entry, aborting!\n");
        return {};
    }

    do_log("[init] loaded service %s into memory\n", service.name.chars());

    TRY(g_services.try_append(move(service)));

    return {};
}

static Result<void> load_services(StringView path)
{
    auto dir = TRY(os::Directory::open(path));

    auto services = TRY(dir->list_names(os::Directory::Filter::ParentAndBase));
    sort(services.begin(), services.end(), String::compare);

    for (const auto& entry : services) TRY(load_service({ dir->fd(), entry.view() }));

    return {};
}

static Result<void> start_services(StringView path)
{
    TRY(load_services(path));
    for (auto& service : g_services)
    {
        do_log("[init] starting service %s\n", service.name.chars());
        start_service(service);
    }

    return {};
}

static Result<void> set_hostname()
{
    auto file = TRY(os::File::open("/etc/hostname", os::File::ReadOnly));

    auto hostname = TRY(file->read_line());
    hostname.trim("\n");

    if (hostname.is_empty())
    {
        do_log("[init] /etc/hostname is empty or invalid, keeping the default hostname\n");
        return {};
    }

    Utf8StringDecoder decoder(hostname.chars());
    if (decoder.code_points().has_error())
    {
        do_log("[init] /etc/hostname is not valid UTF-8, keeping the default hostname\n");
        return {};
    }

    if (sethostname(hostname.chars(), hostname.length()) < 0) return {};

    do_log("[init] successfully set system hostname to '%s'\n", hostname.chars());

    return {};
}

static void mount_tmpfs()
{
    if (mount("/tmp", "tmpfs", "tmpfs") < 0) exit(255);

    if (chmod("/tmp", 01777) < 0) exit(255);
}

static void mount_shmfs()
{
    if (mkdir("/dev/shm", 0755) < 0) exit(255);

    if (mount("/dev/shm", "tmpfs", "tmpfs") < 0) exit(255);

    if (chmod("/dev/shm", 01777) < 0) exit(255);
}

static void mount_devpts()
{
    if (mkdir("/dev/pts", 0755) < 0) exit(255);

    if (mount("/dev/pts", "devpts", "devpts") < 0) exit(255);
}

static void wait_for_child(int)
{
    int status;
    auto rc = os::Process::wait(os::Process::ANY_CHILD, &status);
    if (rc.has_error())
    {
        do_log("[init] waitpid error %s", rc.error_string());
        return;
    }

    pid_t child = rc.release_value();

    for (auto& service : g_services)
    {
        if (service.pid.has_value() && service.pid.value() == child)
        {
            if (WIFEXITED(status))
            {
                do_log("[init] service %s exited with status %d\n", service.name.chars(), WEXITSTATUS(status));
            }
            else { do_log("[init] service %s was terminated by signal %d\n", service.name.chars(), WTERMSIG(status)); }

            if (service.restart)
            {
                do_log("[init] restarting service %s\n", service.name.chars());

                start_service(service);
            }

            break;
        }
    }
}

Result<int> sysinit(StringView path)
{
    if (getpid() != 1)
    {
        fprintf(stderr, "error: init not running as PID 1.\n");
        return 1;
    }

    g_is_system = true;

    // Before this point, we don't even have an stdin, stdout and stderr. Set it up now so that child processes (and us)
    // can print stuff.
    stdin = fopen("/dev/null", "r");
    stdout = fopen("/dev/kmsg", "w");
    stderr = fopen("/dev/kmsg", "w");

    TRY(os::Security::pledge("stdio rpath wpath cpath fattr host mount proc exec id", nullptr));

    mount_tmpfs();
    mount_shmfs();
    mount_devpts();

    umask(022);

    g_init_log = fopen("/dev/uart0", "w");
    check(g_init_log);
    setlinebuf(g_init_log);
    fcntl(fileno(g_init_log), F_SETFD, FD_CLOEXEC);

    set_hostname();

    if (signal(SIGTERM, sigterm_handler) == SIG_ERR) do_errlog("[init] failed to register handler for SIGTERM\n");
    if (signal(SIGQUIT, sigquit_handler) == SIG_ERR) do_errlog("[init] failed to register handler for SIGQUIT\n");

    TRY(os::Security::pledge("stdio rpath wpath cpath proc exec id", nullptr));

    if (path.is_empty()) path = "/etc/init";
    start_services(path);

    while (1) { wait_for_child(0); }
}

Result<int> user_init(StringView path)
{
    setpgid(0, 0);

    g_init_log = fopen("/dev/uart0", "w");
    check(g_init_log);
    setlinebuf(g_init_log);
    fcntl(fileno(g_init_log), F_SETFD, FD_CLOEXEC);

    TRY(os::Security::pledge("stdio rpath wpath cpath proc exec", nullptr));

    if (path.is_empty()) path = "/etc/user";
    start_services(path);

    TRY(os::Security::pledge("stdio rpath wpath proc exec", nullptr));

    while (1) { wait_for_child(0); }
}

Result<int> luna_main(int argc, char** argv)
{
    bool user;
    StringView service_path;

    os::ArgumentParser parser;
    parser.add_description("The init system for Luna.");
    parser.add_system_program_info("init"_sv);
    parser.add_switch_argument(user, 'u', "user"_sv, "initialize a user session instead of the system");
    parser.add_value_argument(service_path, 's', "service-path"_sv,
                              "change the default service path (/etc/init or /etc/user)");
    parser.parse(argc, argv);

    signal(SIGCHLD, wait_for_child);

    if (user) return user_init(service_path);
    return sysinit(service_path);
}