#include <bits/termios.h>
#include <os/ArgumentParser.h>
#include <pwd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <unistd.h>

static struct termios orig;

void restore_terminal()
{
    ioctl(fileno(stdin), TCSETS, &orig);
}

char* getpass()
{
    fputs("Password: ", stdout);

    if (ioctl(fileno(stdin), TCGETS, &orig) < 0)
    {
        perror("ioctl(TCGETS)");
        return nullptr;
    }

    atexit(restore_terminal);

    struct termios tc = orig;
    tc.c_lflag &= ~ECHO;
    if (ioctl(fileno(stdin), TCSETS, &tc) < 0)
    {
        perror("ioctl(TCSETS)");
        return nullptr;
    }

    static char buf[1024];
    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;
    }

    if ((prompt_password || getuid() != geteuid()) && *entry->pw_passwd)
    {
        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);
    }

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