#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/Config.h>
#include <os/Directory.h>
#include <os/File.h>
#include <os/IPC.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 };
    bool wait_notify { 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();
    }

    os::IPC::Notifier notifier;
    if (service.wait_notify) { notifier = os::IPC::Notifier::create(); }

    pid_t pid = TRY(os::Process::fork());
    if (pid == 0)
    {
        if (service.wait_notify) { notifier.hook(); }
        auto rc = service_child(service, new_stdout, new_stderr, new_stdin);
        if (service.wait_notify) { notifier.unhook(); }
        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
    {
        if (service.wait_notify)
        {
            bool success = notifier.wait(1000);
            if (!success)
            {
                do_log("[init] service %s with pid %d failed to start successfully\n", service.name.chars(), pid);
            }
        }
        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::ConfigFile::open(path));

    Service service;

    auto name = file->read_string("Name");
    if (!name.has_value())
    {
        do_log("[init] service file is missing 'Name' entry, aborting!\n");
        return {};
    }
    service.name = TRY(String::from_string_view(name.value()));

    auto command = file->read_string("Command");
    if (!command.has_value())
    {
        do_log("[init] service file is missing 'Command' entry, aborting!\n");
        return {};
    }
    service.command = TRY(String::from_string_view(command.value()));

    service.restart = file->read_boolean_or("Restart", false);
    service.environment = TRY(String::from_string_view(file->read_string_or("Environment", {})));
    service.standard_output = TRY(String::from_string_view(file->read_string_or("StandardOutput", {})));
    service.standard_error = TRY(String::from_string_view(file->read_string_or("StandardError", {})));
    service.standard_input = TRY(String::from_string_view(file->read_string_or("StandardInput", {})));
    service.working_directory = TRY(String::from_string_view(file->read_string_or("WorkingDirectory", {})));

    if (g_is_system)
    {
        auto user = file->read_string("User");
        if (user.has_value())
        {
            auto* pw = getpwnam(user.value().chars());
            if (pw)
            {
                service.user = pw->pw_uid;
                service.group = pw->pw_gid;
            }
        }
    }

    service.wait = file->read_boolean_or("Wait", false);
    service.wait_notify = file->read_boolean_or("WaitUntilReady", false);

    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);
}