#include <os/ArgumentParser.h>
#include <pwd.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <termios.h>
#include <unistd.h>

static struct termios orig;

void restore_terminal()
{
    tcsetattr(STDIN_FILENO, TCSANOW, &orig);
}

void signal_handler(int signo)
{
    restore_terminal();
    raise(signo);
}

char* getpass()
{
    if (!isatty(STDIN_FILENO))
    {
        // FIXME: Just read from /dev/tty (the controlling terminal). Problem: that doesn't exist yet.
        fprintf(stderr, "error: password must be read from a terminal!");
        return nullptr;
    }

    tcsetpgrp(STDIN_FILENO, getpgid(0));

    fputs("Password: ", stdout);

    if (tcgetattr(STDIN_FILENO, &orig) < 0)
    {
        perror("tcgetattr");
        return nullptr;
    }

    struct sigaction sa;
    sa.sa_handler = signal_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESETHAND;
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);
    sigaction(SIGQUIT, &sa, NULL);

    atexit(restore_terminal);

    struct termios tc = orig;
    tc.c_lflag &= ~ECHO;
    if (tcsetattr(STDIN_FILENO, TCSANOW, &tc) < 0)
    {
        perror("tcsetattr");
        return nullptr;
    }

    static char buf[BUFSIZ];
    char* rc = fgets(buf, sizeof(buf), stdin);

    restore_terminal();
    putchar('\n');

    if (!rc)
    {
        perror("fgets");
        return nullptr;
    }

    char* newline = strrchr(rc, '\n');
    if (newline) *newline = 0;

    return buf;
}

Result<int> luna_main(int argc, char** argv)
{
    StringView name;
    bool prompt_password;
    bool login;

    if (geteuid() != 0)
    {
        fprintf(stderr, "%s must be setuid root!\n", argv[0]);
        return 1;
    }

    os::ArgumentParser parser;
    parser.add_description("Switch to a different user (by default, root)."_sv);
    parser.add_system_program_info("su"_sv);
    parser.add_positional_argument(name, "name"_sv, "root"_sv);
    parser.add_switch_argument(prompt_password, 'p', "prompt", "prompt for a password even if running as root");
    parser.add_switch_argument(login, 'l', "login"_sv, "change directory to the user's home and start a login shell");
    parser.parse(argc, argv);

    struct passwd* entry = getpwnam(name.chars());
    if (!entry)
    {
        fprintf(stderr, "%s: user %s not found!\n", argv[0], name.chars());
        return 1;
    }

    endpwent();

    if ((prompt_password || getuid() != geteuid()) && *entry->pw_passwd)
    {
        signal(SIGTTOU, SIG_IGN);

        char* pass = getpass();
        if (!pass) return 1;

        if (strcmp(pass, entry->pw_passwd))
        {
            fprintf(stderr, "%s: wrong password!\n", argv[0]);
            return 1;
        }

        memset(pass, 0, strlen(pass));
    }

    setgid(entry->pw_gid);
    setuid(entry->pw_uid);

    if (login)
    {
        chdir(entry->pw_dir);
        clearenv();
        setenv("PATH", "/bin:/sbin", 1);
        setpgid(0, 0);
    }

    if (login || entry->pw_uid != 0) setenv("USER", entry->pw_name, 1);

    setenv("HOME", entry->pw_dir, 1);
    setenv("SHELL", entry->pw_shell, 1);

    execl(entry->pw_shell, entry->pw_shell, NULL);

    perror("execl");
    return 1;
}