/*
    * login_input.cpp
    * Input and session flow for the MontaukOS login screen
    * Copyright (c) 2026 Daniel Hammer
*/

#include "login.hpp"

namespace {

void transition_to_login_mode(LoginState* ls) {
    ls->mode = MODE_LOGIN;
    ls->active_field = 0;
    // Scrub the full buffers, not just the first byte: plaintext residue must
    // not linger in the heap after the credentials are consumed.
    montauk::memset(ls->password, 0, sizeof(ls->password));
    ls->password_len = 0;
    montauk::memset(ls->confirm, 0, sizeof(ls->confirm));
    ls->confirm_len = 0;
    gui::mtk::text_input_reset(ls->password_input);
    gui::mtk::text_input_reset(ls->confirm_input);
    gui::mtk::text_input_set_cursor(ls->username_input, ls->username_len,
                                    ls->username_len, false);
    ls->show_error = false;
}

void finish_desktop_session(LoginState* ls) {
    montauk::user::clear_session();
    montauk::memset(ls->password, 0, sizeof(ls->password));
    ls->password_len = 0;
    gui::mtk::text_input_reset(ls->password_input);
    ls->active_field = 1;
    ls->show_error = false;
}

// Build "--console 0:/users/<username>" so the full-screen Terminal starts in
// the user's home directory, mirroring how the desktop is handed the home dir.
void build_console_args(LoginState* ls, char* out, int cap) {
    const char* prefix = "--console 0:/users/";
    int i = 0;
    while (prefix[i] && i < cap - 1) { out[i] = prefix[i]; i++; }
    int j = 0;
    while (ls->username[j] && i < cap - 1) out[i++] = ls->username[j++];
    out[i] = '\0';
}

void launch_session(LoginState* ls) {
    // Authentication already succeeded and the session is set, so the password
    // is no longer needed. Wipe it before the session runs rather than waiting
    // for finish_desktop_session, so it does not sit in the heap for the whole
    // session.
    montauk::memset(ls->password, 0, sizeof(ls->password));
    ls->password_len = 0;

    int pid;
    if (ls->session_mode == SESSION_CONSOLE) {
        char console_args[96];
        build_console_args(ls, console_args, (int)sizeof(console_args));
        pid = montauk::spawn("0:/apps/terminal/terminal.elf", console_args);
    } else {
        pid = montauk::spawn("0:/os/desktop.elf", ls->username);
    }
    if (pid >= 0) {
        montauk::setuser(pid, ls->username);
        montauk::waitpid(pid);
    }

    // The desktop hands off a graceful power-off by posting a request and
    // exiting (see desktop_request_power). Pick it up now that waitpid has
    // returned and run the shutdown stages; perform_graceful_shutdown never
    // returns. A normal logout (or the console session exiting) leaves no
    // request, so we fall through and return to the login screen.
    int req = montauk::power_request(montauk::abi::POWER_REQ_QUERY);
    if (req == montauk::abi::POWER_REQ_SHUTDOWN || req == montauk::abi::POWER_REQ_REBOOT) {
        perform_graceful_shutdown(ls, req);
    }

    finish_desktop_session(ls);
}

gui::Rect login_field_rect(const LoginLayout& lo, LoginState* ls, int idx) {
    if (ls->mode == MODE_FIRST_BOOT) {
        switch (idx) {
        case 0: return lo.username_field;
        case 1: return lo.display_name_field;
        case 2: return lo.password_field;
        case 3: return lo.confirm_field;
        default: break;
        }
    } else {
        switch (idx) {
        case 0: return lo.username_field;
        case 1: return lo.password_field;
        default: break;
        }
    }
    return {0, 0, 0, 0};
}

} // namespace

LoginField get_field(LoginState* ls, int field) {
    if (ls->mode == MODE_FIRST_BOOT) {
        switch (field) {
        case 0: return {ls->username, &ls->username_len, (int)sizeof(ls->username), false, &ls->username_input};
        case 1: return {ls->display_name, &ls->display_name_len, (int)sizeof(ls->display_name), false, &ls->display_name_input};
        case 2: return {ls->password, &ls->password_len, (int)sizeof(ls->password), true, &ls->password_input};
        case 3: return {ls->confirm, &ls->confirm_len, (int)sizeof(ls->confirm), true, &ls->confirm_input};
        default: break;
        }
    } else {
        switch (field) {
        case 0: return {ls->username, &ls->username_len, (int)sizeof(ls->username), false, &ls->username_input};
        case 1: return {ls->password, &ls->password_len, (int)sizeof(ls->password), true, &ls->password_input};
        default: break;
        }
    }
    return {nullptr, nullptr, 0, false, nullptr};
}

int max_fields(LoginState* ls) {
    return ls->mode == MODE_FIRST_BOOT ? 4 : 2;
}

bool try_first_boot_submit(LoginState* ls) {
    ls->show_error = false;

    if (ls->username_len == 0) {
        montauk::strcpy(ls->error_msg, "Username is required");
        ls->show_error = true;
        return false;
    }
    if (ls->password_len == 0) {
        montauk::strcpy(ls->error_msg, "Password is required");
        ls->show_error = true;
        return false;
    }
    if (!montauk::streq(ls->password, ls->confirm)) {
        montauk::strcpy(ls->error_msg, "Passwords do not match");
        ls->show_error = true;
        return false;
    }

    montauk::fmkdir("0:/users");

    const char* dname = ls->display_name_len > 0 ? ls->display_name : ls->username;
    if (!montauk::user::create_user(ls->username, dname, ls->password, "admin")) {
        montauk::strcpy(ls->error_msg, "Failed to create user");
        ls->show_error = true;
        return false;
    }

    return true;
}

bool try_login(LoginState* ls) {
    ls->show_error = false;

    if (ls->username_len == 0) {
        montauk::strcpy(ls->error_msg, "Username is required");
        ls->show_error = true;
        return false;
    }

    if (!montauk::user::authenticate(ls->username, ls->password)) {
        montauk::strcpy(ls->error_msg, "Invalid username or password");
        ls->show_error = true;
        return false;
    }

    montauk::user::set_session(ls->username);
    return true;
}

void handle_key(LoginState* ls, const montauk::abi::KeyEvent& key) {
    if (!key.pressed) return;

    // Tab / Shift+Tab cycles fields.
    if (key.scancode == 0x0F) {
        int mf = max_fields(ls);
        if (key.shift) ls->active_field = (ls->active_field + mf - 1) % mf;
        else ls->active_field = (ls->active_field + 1) % mf;
        LoginField f = get_field(ls, ls->active_field);
        if (f.state && f.len)
            gui::mtk::text_input_set_cursor(*f.state, *f.len, *f.len, false);
        ls->show_error = false;
        return;
    }

    // Enter submits.
    if (key.ascii == '\n' || key.ascii == '\r') {
        if (ls->mode == MODE_FIRST_BOOT) {
            if (try_first_boot_submit(ls)) transition_to_login_mode(ls);
        } else if (try_login(ls)) {
            launch_session(ls);
        }
        return;
    }

    // Everything else is text editing, handled by the toolkit (cursor moves,
    // selection, word-jump, clipboard, backspace/delete, insertion).
    LoginField f = get_field(ls, ls->active_field);
    if (!f.text || !f.state || !f.len) return;
    int result = gui::mtk::text_input_key(*f.state, f.text, f.cap, key,
                                          nullptr, nullptr, f.masked);
    if (result & gui::mtk::TEXT_INPUT_CHANGED)
        *f.len = gui::mtk::text_input_len(f.text, f.cap);
}

void handle_mouse(LoginState* ls) {
    int mx = ls->mouse.x;
    int my = ls->mouse.y;
    uint8_t buttons = ls->mouse.buttons;
    uint8_t prev = ls->prev_buttons;
    bool left_pressed = (buttons & 0x01) && !(prev & 0x01);

    LoginLayout lo = layout_login_screen(ls);

    if (left_pressed) {
        // Focus a field, or act on a button.
        int clicked = -1;
        if (lo.username_field.contains(mx, my)) {
            clicked = 0;
        } else if (ls->mode == MODE_FIRST_BOOT) {
            if (lo.display_name_field.contains(mx, my)) clicked = 1;
            else if (lo.password_field.contains(mx, my)) clicked = 2;
            else if (lo.confirm_field.contains(mx, my)) clicked = 3;
        } else {
            if (lo.password_field.contains(mx, my)) clicked = 1;
        }

        if (clicked >= 0) {
            ls->active_field = clicked;
            ls->show_error = false;
        } else if (ls->mode == MODE_LOGIN && lo.desktop_radio.contains(mx, my)) {
            ls->session_mode = SESSION_DESKTOP;
            ls->show_error = false;
        } else if (ls->mode == MODE_LOGIN && lo.console_radio.contains(mx, my)) {
            ls->session_mode = SESSION_CONSOLE;
            ls->show_error = false;
        } else if (lo.submit_button.contains(mx, my)) {
            if (ls->mode == MODE_FIRST_BOOT) {
                if (try_first_boot_submit(ls)) transition_to_login_mode(ls);
            } else if (try_login(ls)) {
                launch_session(ls);
            }
            return;
        } else if (lo.shutdown_button.contains(mx, my)) {
            perform_graceful_shutdown(ls, montauk::abi::POWER_REQ_SHUTDOWN);
            return;
        } else if (lo.reboot_button.contains(mx, my)) {
            perform_graceful_shutdown(ls, montauk::abi::POWER_REQ_REBOOT);
            return;
        }
    }

    // Route the pointer to the active field for cursor placement and
    // drag-selection (press, drag, release are all handled by the toolkit).
    LoginField f = get_field(ls, ls->active_field);
    if (f.text && f.state) {
        gui::Rect r = login_field_rect(lo, ls, ls->active_field);
        gui::mtk::text_input_mouse(*f.state, r, f.text, mx, my, buttons, prev,
                                   true, f.masked);
    }
}
