/*
    * login_shutdown.cpp
    * Graceful shutdown view and staged power-off for the MontaukOS login screen
    * Copyright (c) 2026 Daniel Hammer
*/

#include "login.hpp"
#include <montauk/thread.h>

using namespace gui;

bool volume0_is_persistent() {
    if (montauk::drivekind(0) == 0) return false;
    return true;
}

// Minimum time each shutdown stage stays on screen, so the status message is
// readable even when the underlying work completes near-instantly.
static constexpr uint64_t STAGE_MIN_VISIBLE_MS = 600;

// Per-stage watchdog budgets. Each stage's blocking work runs on a disposable
// worker thread; if it overruns its budget we abandon the worker and move on,
// so an unresponsive device can never wedge the power-off.
//  - Bluetooth: a disconnect is normally sub-second; 4 s tolerates a slow
//    controller. Bailing here is safe -- it does not risk data.
//  - Filesystem flush: kept generous (10 s) so a legitimately slow flush on
//    real hardware is never cut short and writes are not lost; only a truly
//    wedged storage controller hits this bound.
static constexpr uint64_t BT_STAGE_TIMEOUT_MS = 4000;
static constexpr uint64_t FS_STAGE_TIMEOUT_MS = 10000;
static constexpr uint64_t LOG_WRITE_TIMEOUT_MS = 10000;

void draw_shutdown_screen(LoginState* ls, const char* heading, const char* status) {
    Framebuffer& fb = ls->fb;
    int sw = ls->screen_w;
    int sh = ls->screen_h;

    // Same packed/scratch handling as draw_login_screen: Canvas assumes a
    // tightly packed buffer, so compose into the scratch buffer when the
    // framebuffer pitch is not width*4.
    bool packed = fb.pitch() == sw * (int)sizeof(uint32_t);
    uint32_t* target = packed ? fb.buffer() : ls->compose;
    if (!target) return;
    Canvas c(target, sw, sh);

    const mtk::Theme& th = ls->theme;
    int sfh = system_font_height();

    // ==== Layout (mirrors the login/setup card chrome, minus the footer) ====
    // Titlebar holds the app title; the body holds a heading + a status line,
    // matching the login/setup heading + subtitle. The card is sized to fit the
    // content with symmetric padding -- no empty footer band.
    const char* window_title = "MontaukOS";
    int content_x = CONTENT_PAD_X;

    int y = TITLEBAR_H + CONTENT_TOP_PAD;
    int heading_y = y;
    y += sfh + 6;
    int status_y = y;
    y += sfh;

    int card_h = y + CONTENT_TOP_PAD + 10;
    int card_x = (sw - CARD_W) / 2;
    int card_y = (sh - card_h) / 2;

    Rect card = {card_x, card_y, CARD_W, card_h};
    Rect titlebar = {card_x, card_y, CARD_W, TITLEBAR_H};

    // ==== Background ====
    if (ls->has_wallpaper && ls->bg_wallpaper) {
        montauk::memcpy(target, ls->bg_wallpaper,
                        (uint64_t)sw * sh * sizeof(uint32_t));
    } else {
        c.fill(BG_COLOR);
    }

    // ==== Card ====
    int so = 4; // drop-shadow offset
    c.fill_rect_alpha(card.x + so, card.y + card.h, card.w, so, colors::SHADOW);
    c.fill_rect_alpha(card.x + card.w, card.y + so, so, card.h, colors::SHADOW);
    c.fill_rect_alpha(card.x + card.w, card.y + card.h, so, so, colors::SHADOW);

    c.fill_rect(card.x, card.y, card.w, card.h, CARD_BG);
    c.fill_rect(titlebar.x, titlebar.y, titlebar.w, titlebar.h, TITLEBAR_BG);
    c.rect(card.x, card.y, card.w, card.h, CARD_BORDER);
    c.hline(titlebar.x, titlebar.y + titlebar.h - 1, titlebar.w, CARD_BORDER);

    // ==== Title + heading + status ====
    int window_tw = text_width(window_title);
    c.text(titlebar.x + (titlebar.w - window_tw) / 2,
           titlebar.y + (TITLEBAR_H - sfh) / 2, window_title, th.text);
    c.text(card_x + content_x, card_y + heading_y, heading, th.text);
    c.text(card_x + content_x, card_y + status_y, status, th.text_subtle);

    // ==== Present ====
    if (!packed) fb.copy_from(ls->compose, sw * (int)sizeof(uint32_t));
    fb.flip();
}

// Show a stage message, hold it on screen for a moment, then return so the
// caller can perform the stage's work.
static void show_stage(LoginState* ls, const char* heading, const char* status) {
    draw_shutdown_screen(ls, heading, status);
    montauk::sleep_ms(STAGE_MIN_VISIBLE_MS);
}

namespace {

// A unit of shutdown work plus a flag the worker sets when it finishes. `done`
// is polled across threads, so all access goes through atomics.
struct StageWork {
    void (*fn)();
    volatile bool done;
};

int stage_worker(void* arg) {
    StageWork* w = (StageWork*)arg;
    w->fn();
    __atomic_store_n(&w->done, true, __ATOMIC_RELEASE);
    return 0;
}

// Run fn() on a worker thread, waiting up to timeout_ms for it to finish. If a
// device wedges, only the disposable worker blocks (in its kernel syscall); the
// shutdown flow keeps moving. Returns true if fn() completed in time, false if
// it was abandoned. A timed-out worker is intentionally left running and never
// joined (joining would re-block us) -- power-off reclaims it in moments.
bool run_stage(void (*fn)(), uint64_t timeout_ms) {
    auto* w = (StageWork*)montauk::malloc(sizeof(StageWork));
    if (!w) { fn(); return true; }              // no memory: best-effort inline
    w->fn = fn;
    w->done = false;

    int tid = montauk::thread_spawn(stage_worker, w);
    if (tid < 0) { fn(); montauk::mfree(w); return true; } // can't isolate: inline

    const uint64_t step = 50;
    uint64_t waited = 0;
    while (!__atomic_load_n(&w->done, __ATOMIC_ACQUIRE) && waited < timeout_ms) {
        montauk::sleep_ms(step);
        waited += step;
    }

    if (__atomic_load_n(&w->done, __ATOMIC_ACQUIRE)) {
        montauk::thread_join(tid, nullptr);     // reclaim the worker's stack
        montauk::mfree(w);
        return true;
    }
    return false;
}

// True only if there is Bluetooth work to do: an adapter is present and at
// least one device is connected. bt_info/bt_list are non-blocking table reads
// that return immediately when no adapter exists (e.g. under QEMU), so this
// gate is effectively free.
bool bluetooth_has_active_connection() {
    montauk::abi::BtAdapterInfo adapter;
    if (montauk::bt_info(&adapter) != 0 || !adapter.initialized) {
        return false;  // no adapter
    }
    montauk::abi::BtDevInfo devs[8];
    int n = montauk::bt_list(devs, 8);
    for (int i = 0; i < n; i++) {
        if (devs[i].connected) return true;
    }
    return false;
}

// Stage bodies -- no captured state, so they double as bare thread entries.
void stage_disconnect_bluetooth() {
    montauk::abi::BtDevInfo devs[8];
    int n = montauk::bt_list(devs, 8);
    for (int i = 0; i < n; i++) {
        if (devs[i].connected) {
            montauk::bt_disconnect(devs[i].bdAddr);
        }
    }
}

static bool path_is_file(const char* path) {
    montauk::abi::FileStat st;
    if (montauk::stat(path, &st) < 0) return false;

    return !st.isDir;
}

#define LOG_READ_SIZE 65536

// Write the system log to a file on disk
void stage_save_log() {
    // Edge case - 0:/os/logs exists and is a file
    if (path_is_file("0:/os/logs"))
        return;

    montauk::fmkdir("0:/os/logs");

    // Clear any existing montaukos log 
    montauk::fdelete("0:/os/logs/montaukos");

    int handle = montauk::fcreate("0:/os/logs/montaukos");
    if (handle < 0)
        return;

    uint8_t* syslog_buf = (uint8_t*)montauk::malloc(LOG_READ_SIZE);
    int n = montauk::read_log((char *)syslog_buf, LOG_READ_SIZE);

    int result = montauk::fwrite(handle, syslog_buf, 0, n);
    
    montauk::close(handle);
    montauk::mfree(syslog_buf);
}

void stage_flush_filesystems() {
    montauk::fs_sync();
}

} // namespace

void perform_graceful_shutdown(LoginState* ls, int action) {
    const bool rebooting = (action == montauk::abi::POWER_REQ_REBOOT);
    const char* heading = rebooting ? "Restarting" : "Shutting Down";

    // ==== Stage 1: disconnect connected Bluetooth devices ====
    // Skip entirely when there is no adapter or nothing connected, so a typical
    // shutdown (and every QEMU run) does not pay for an empty stage. When there
    // is work, it is bounded so an unresponsive controller cannot block the
    // (critical) filesystem flush that follows.
    if (bluetooth_has_active_connection()) {
        show_stage(ls, heading, "Disconnecting Bluetooth devices...");
        if (!run_stage(stage_disconnect_bluetooth, BT_STAGE_TIMEOUT_MS)) {
            show_stage(ls, heading, "Bluetooth is unresponsive, continuing...");
        }
    }

    // No need to save log files or flush to disk if volume 0 is a ramdisk anyway
    if (volume0_is_persistent()) {
        // ==== Stage 2: write system log to disk ====
        show_stage(ls, heading, "Saving system log...");
        if (!run_stage(stage_save_log, LOG_WRITE_TIMEOUT_MS)) {
            show_stage(ls, heading, "Log save is unresponsive, continuing...");
        } 

        // ==== Stage 3: flush writes and unmount filesystems ====
        show_stage(ls, heading, "Flushing file systems...");
        if (!run_stage(stage_flush_filesystems, FS_STAGE_TIMEOUT_MS)) {
            show_stage(ls, heading, "Storage is unresponsive, continuing...");
        }
    }


    // ==== Stage 4: dispatch the ACPI power-off / reset ====
    show_stage(ls, heading, rebooting ? "Restarting now..." : "Powering off...");
    int rc = rebooting ? montauk::reset() : montauk::shutdown();
    // A successful power-control syscall never returns.  If it does return,
    // keep the trusted supervisor alive and make the authorization failure
    // visible instead of falling through an unreachable-code assumption.
    show_stage(ls, heading, rc == montauk::abi::SYS_ERR_PERMISSION
        ? "Power control permission denied."
        : "Power control failed.");
    for (;;) montauk::sleep_ms(1000);
}