#define CLIENT_IMPLEMENTATION
#include "Client.h"
#include "IPC.h"
#include "Keyboard.h"
#include "Mouse.h"
#include "Screen.h"
#include "Window.h"
#include <errno.h>
#include <grp.h>
#include <moon/Keyboard.h>
#include <os/ArgumentParser.h>
#include <os/File.h>
#include <os/LocalServer.h>
#include <os/Process.h>
#include <os/Security.h>
#include <pwd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/poll.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

static void debug(const Vector<OwnedPtr<Client>>& clients)
{
    os::println("--- wind: DEBUG OUTPUT ---");

    os::println("-- wind: Listing clients --");

    for (const auto& client : clients)
    {
        os::println("Client with fd %d, owns %zu windows", client->conn.fd(), client->windows.size());
    }

    os::println("-- wind: Listing windows --");

    for (const auto& window : g_windows)
    {
        os::println("Window of client (fd %d), id %d, %sdecorated, %sdirty (\"%s\") (%d,%d,%d,%d)",
                    window->client->conn.fd(), window->id, window->type == ui::WindowType::Normal ? "" : "not ",
                    window->dirty ? "" : "not ", window->name.chars(), window->surface.pos.x, window->surface.pos.y,
                    window->surface.width, window->surface.height);
    }

    os::println("-- wind: Listing processes --");

    system("ps");

    os::println("-- wind: Listing memory usage --");

    system("free -h");

    os::println("--- wind: END DEBUG OUTPUT ---");
}

Result<void> set_supplementary_groups(const char* name)
{
    Vector<gid_t> extra_groups;

    setgrent();
    group* grp;
    while ((grp = getgrent()))
    {
        for (char** user = grp->gr_mem; *user; user++)
        {
            if (!strcmp(*user, name))
            {
                os::println("Adding supplementary group: %d", grp->gr_gid);
                TRY(extra_groups.try_append(grp->gr_gid));
                break;
            }
        }
    }
    endgrent();

    if (setgroups(static_cast<int>(extra_groups.size()), extra_groups.data()) < 0) return err(errno);

    return {};
}

Result<int> luna_main(int argc, char** argv)
{
    srand((unsigned)time(NULL));

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

    StringView socket_path = "/tmp/wind.sock";
    StringView user;

    os::ArgumentParser parser;
    parser.add_description("The display server for Luna's graphical user interface."_sv);
    parser.add_system_program_info("wind"_sv);
    parser.add_value_argument(socket_path, 's', "socket"_sv, "the path for the local IPC socket"_sv);
    parser.add_value_argument(user, 'u', "user"_sv, "the user to run as"_sv);
    parser.parse(argc, argv);

    if (geteuid() != 0)
    {
        os::eprintln("error: wind must be run as root to initialize resources, run with --user=<USERNAME> to drop "
                     "privileges afterwards");
        return 1;
    }

    auto mouse = TRY(os::File::open("/dev/mouse", os::File::ReadOnly));
    mouse->set_buffer(os::File::NotBuffered);
    mouse->set_close_on_exec();

    auto keyboard = TRY(os::File::open("/dev/kbd", os::File::ReadOnly));
    keyboard->set_buffer(os::File::NotBuffered);
    keyboard->set_close_on_exec();

    TRY(Screen::open());
    auto& screen = Screen::the();

    Mouse mouse_pointer { screen.canvas() };

    int fd = open("/dev/null", O_RDONLY);
    if (fd >= 0)
    {
        dup2(fd, STDIN_FILENO);
        close(fd);
    }

    setegid(2);
    seteuid(2);

    if (setsid() < 0) perror("setsid");

    mode_t mask = umask(0002);

    auto server = TRY(os::LocalServer::create(socket_path, false));
    TRY(server->listen(20));

    umask(mask);

    seteuid(0);

    clearenv();

    pid_t child = TRY(os::Process::fork());
    if (!child)
    {
        if (!user.is_empty())
        {
            auto* pwd = getpwnam(user.chars());
            if (pwd)
            {
                TRY(set_supplementary_groups(user.chars()));
                setgid(pwd->pw_gid);
                setuid(pwd->pw_uid);
            }
        }

        StringView args[] = { "/usr/bin/init"_sv, "--user"_sv };
        TRY(os::Process::exec("/usr/bin/init"_sv, Slice<StringView> { args, 2 }, false));
    }

    umask(0002);

    setuid(2);
    setgid(2);

    ui::Color background = ui::Color::from_rgb(0x10, 0x10, 0x10);

    Vector<OwnedPtr<Client>> clients;
    Vector<struct pollfd> fds;
    TRY(fds.try_append({ .fd = mouse->fd(), .events = POLLIN, .revents = 0 }));
    TRY(fds.try_append({ .fd = keyboard->fd(), .events = POLLIN, .revents = 0 }));
    TRY(fds.try_append({ .fd = server->fd(), .events = POLLIN, .revents = 0 }));

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

    while (1)
    {
        screen.canvas().fill(background);
        for (auto* window : g_windows) window->draw(screen.canvas());
        mouse_pointer.draw(screen.canvas());
        screen.sync();

        for (auto& pfd : fds) { pfd.revents = 0; }

        int rc = poll(fds.data(), fds.size(), 1000);
        if (!rc) continue;
        if (rc < 0 && errno != EINTR) { os::println("poll: error: %s", strerror(errno)); }

        if (fds[0].revents & POLLIN)
        {
            moon::MousePacket packet;
            TRY(mouse->read_typed(packet));
            mouse_pointer.update(packet);
        }
        if (fds[1].revents & POLLIN)
        {
            moon::KeyboardPacket packet;
            TRY(keyboard->read_typed(packet));
            if (!packet.released && packet.key == moon::K_Tab) debug(clients);
            auto request = wind::Keyboard::decode_keyboard_event((moon::KeyCode)packet.key, packet.released);
            if (g_windows.last().has_value())
            {
                auto* window = g_windows.last().value();
                request.window = window->id;
                os::IPC::send_async(window->client->conn, request);
            }
        }
        for (usize i = 0; i < clients.size(); i++)
        {
            if (fds[i + 3].revents & POLLIN) wind::handle_ipc(*clients[i]);
            if (fds[i + 3].revents & POLLHUP)
            {
                os::println("wind: Client %d disconnected", i);
                fds.remove_at(i + 3);
                auto client = clients.remove_at(i);
                client->conn.disconnect();
                for (auto& window : client->windows)
                {
                    if (window)
                    {
                        g_windows.remove(window);
                        mouse_pointer.window_did_close(window);
                        delete window;
                    }
                }
            }
        }
        if (fds[2].revents & POLLIN)
        {
            auto client = TRY(server->accept());
            os::println("wind: New client connected!");
            TRY(fds.try_append({ .fd = client.fd(), .events = POLLIN, .revents = 0 }));
            OwnedPtr<Client> c = TRY(adopt_owned_if_nonnull(new Client(move(client))));
            TRY(clients.try_append(move(c)));
        }
    }
}