#include <errno.h>
#include <fcntl.h>
#include <luna/PathParser.h>
#include <luna/Sort.h>
#include <luna/String.h>
#include <luna/Vector.h>
#include <os/Directory.h>
#include <os/File.h>
#include <os/Process.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>

FILE* g_init_log = nullptr;

// Request a successful exit from the system (for tests)
void sigterm_handler(int)
{
    _exit(0);
}

// Request a failure exit from the system (for tests)
void sigquit_handler(int)
{
    _exit(1);
}

struct Service
{
    String name;
    String command;
    bool restart { false };
    String environment;
    String standard_output;
    String standard_error;
    String standard_input;
    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_init_log) vfprintf(g_init_log, format, ap);

    va_end(ap);
}

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.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_log("[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_log("[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")
        {
            if (!service.command.is_empty())
            {
                do_log("[init] 'Command' cannot be specified after 'Script' has already been set! (%s)\n",
                       line.chars());
                return {};
            }
            service.command = move(parts[1]);
            continue;
        }

        if (parts[0].view() == "Script")
        {
            if (!service.command.is_empty())
            {
                do_log("[init] 'Script' cannot be specified after 'Command' has already been set! (%s)\n",
                       line.chars());
                return {};
            }
            service.command = TRY(String::format("/bin/sh -- %s"_sv, parts[1].chars()));
            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() == "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' or 'Script' 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()
{
    auto dir = TRY(os::Directory::open("/etc/init"));

    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()
{
    TRY(load_services());
    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 (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);
}

int main()
{
    if (getpid() != 1)
    {
        fprintf(stderr, "error: init not running as PID 1.\n");
        return 1;
    }

    // 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/console", "r");
    stdout = fopen("/dev/console", "w");
    stderr = fopen("/dev/console", "w");

    mount_tmpfs();

    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_log("[init] failed to register handler for SIGTERM\n");
    if (signal(SIGQUIT, sigquit_handler) == SIG_ERR) do_log("[init] failed to register handler for SIGQUIT\n");

    start_services();

    while (1)
    {
        int status;
        auto rc = os::Process::wait(os::Process::ANY_CHILD, &status);
        if (rc.has_error()) continue;

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