/*
 * main.cpp
 * MontaukOS Screenshot Tool
 * Captures the screen framebuffer and saves JPEG screenshots through a GUI flow.
 * Copyright (c) 2026 Daniel Hammer
 */

#include <montauk/syscall.h>
#include <montauk/string.h>
#include <montauk/heap.h>
#include <gui/dialogs.hpp>
#include <gui/mtk.hpp>
#include <gui/standalone.hpp>
#include <gui/truetype.hpp>

extern "C" {
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

#define STBI_WRITE_NO_STDIO
#include <gui/stb_image_write.h>
}

using namespace gui;

// ============================================================================
// Constants
// ============================================================================

static constexpr int INIT_W = 420;
static constexpr int INIT_H = 230;
static constexpr int PAD = 18;
static constexpr int OPTION_H = 44;
static constexpr int OPTION_GAP = 8;
static constexpr int FOOTER_H = 56;
static constexpr int BUTTON_W = 132;
static constexpr int BUTTON_H = 30;
static constexpr int MIN_AREA_SIZE = 4;

// ============================================================================
// App state
// ============================================================================

enum CaptureMode : uint8_t {
    CAPTURE_SCREEN = 0,
    CAPTURE_AREA = 1,
};

struct ScreenBuffer {
    uint32_t* pixels;
    int width;
    int height;
    int pitch;
};

struct JpegBuffer {
    uint8_t* data;
    uint64_t size;
    uint64_t capacity;
    bool failed;
};

static int g_win_w = INIT_W;
static int g_win_h = INIT_H;
static int g_mouse_x = -1;
static int g_mouse_y = -1;
static CaptureMode g_mode = CAPTURE_SCREEN;

// ============================================================================
// Small helpers
// ============================================================================

static void safe_copy(char* dst, int dst_len, const char* src) {
    if (!dst || dst_len <= 0) return;
    if (!src) src = "";
    montauk::strncpy(dst, src, dst_len);
    dst[dst_len - 1] = '\0';
}

static bool str_ends_with_ci(const char* s, const char* suffix) {
    if (!s || !suffix) return false;
    int slen = montauk::slen(s);
    int suflen = montauk::slen(suffix);
    if (suflen > slen) return false;

    const char* tail = s + slen - suflen;
    for (int i = 0; i < suflen; i++) {
        char a = tail[i];
        char b = suffix[i];
        if (a >= 'A' && a <= 'Z') a += 32;
        if (b >= 'A' && b <= 'Z') b += 32;
        if (a != b) return false;
    }
    return true;
}

static void show_error(const char* message) {
    if (message && message[0]) {
        montauk::print("screenshot: ");
        montauk::print(message);
        montauk::print("\n");
        gui::dialogs::message_box("Screenshot", message, gui::dialogs::MESSAGE_BOX_OK);
    }
}

static void free_screen(ScreenBuffer* screen) {
    if (!screen) return;
    if (screen->pixels) montauk::mfree(screen->pixels);
    screen->pixels = nullptr;
    screen->width = 0;
    screen->height = 0;
    screen->pitch = 0;
}

static Rect normalized_rect(int x0, int y0, int x1, int y1) {
    int left = x0 < x1 ? x0 : x1;
    int top = y0 < y1 ? y0 : y1;
    int right = x0 > x1 ? x0 : x1;
    int bottom = y0 > y1 ? y0 : y1;
    return {left, top, right - left + 1, bottom - top + 1};
}

static Rect clamp_rect_to_screen(const Rect& rect, int w, int h) {
    Rect bounds = {0, 0, w, h};
    return rect.intersect(bounds);
}

// ============================================================================
// MTK picker window
// ============================================================================

static mtk::Theme app_theme() {
    mtk::Theme theme = mtk::make_theme();
    theme.window_bg = colors::WHITE;
    theme.surface = Color::from_rgb(0xF5, 0xF5, 0xF5);
    theme.surface_alt = Color::from_rgb(0xFA, 0xFA, 0xFA);
    theme.surface_hover = mtk::mix(theme.surface, theme.accent, 18);
    theme.border = Color::from_rgb(0xD0, 0xD0, 0xD0);
    theme.text = Color::from_rgb(0x2E, 0x2E, 0x2E);
    theme.text_muted = Color::from_rgb(0x74, 0x74, 0x74);
    theme.text_subtle = Color::from_rgb(0x88, 0x88, 0x88);
    return theme;
}

static Rect option_rect(int index) {
    int top = 48 + index * (OPTION_H + OPTION_GAP);
    return {PAD, top, gui_max(g_win_w - PAD * 2, 0), OPTION_H};
}

static Rect capture_button_rect() {
    int y = g_win_h - FOOTER_H + (FOOTER_H - BUTTON_H) / 2;
    int w = gui_min(BUTTON_W, gui_max((g_win_w - PAD * 3) / 2, 72));
    return {g_win_w - PAD - w, y, w, BUTTON_H};
}

static Rect cancel_button_rect() {
    Rect capture = capture_button_rect();
    int w = capture.w;
    return {capture.x - 8 - w, capture.y, w, BUTTON_H};
}

static bool mouse_in_rect(const Rect& rect) {
    return rect.contains(g_mouse_x, g_mouse_y);
}

static void draw_option(Canvas& c,
                        const Rect& rect,
                        CaptureMode mode,
                        const char* label,
                        const mtk::Theme& theme) {
    if (rect.empty()) return;

    bool selected = g_mode == mode;
    bool hovered = mouse_in_rect(rect);
    if (selected) {
        Color fill = hovered ? mtk::mix(theme.accent_soft, theme.accent, 14)
                             : theme.accent_soft;
        mtk::draw_rounded_frame(c, rect, theme.radius_md, fill, fill);
    } else if (hovered) {
        mtk::draw_rounded_frame(c, rect, theme.radius_md,
                                theme.surface_hover, theme.surface_hover);
    }

    Rect radio = {rect.x + 12, rect.y + (rect.h - 24) / 2, rect.w - 24, 24};
    mtk::draw_radio(c, radio, label, selected, theme);
}

static void render_picker(Canvas& c) {
    mtk::Theme theme = app_theme();
    c.fill(theme.window_bg);

    c.text(PAD, 16, "Capture mode", theme.text);

    draw_option(c, option_rect(0), CAPTURE_SCREEN, "Whole screen", theme);
    draw_option(c, option_rect(1), CAPTURE_AREA, "Selected area", theme);

    int footer_y = gui_max(g_win_h - FOOTER_H, 0);
    c.fill_rect(0, footer_y, g_win_w, gui_max(g_win_h - footer_y, 0), theme.surface);
    mtk::draw_separator(c, 0, footer_y, g_win_w, theme);

    Rect cancel = cancel_button_rect();
    Rect capture = capture_button_rect();
    mtk::draw_button(c, cancel, "Cancel", mtk::BUTTON_SECONDARY,
                     mtk::widget_state(false, mouse_in_rect(cancel)), theme);
    mtk::draw_button(c, capture, "Capture", mtk::BUTTON_PRIMARY,
                     mtk::widget_state(false, mouse_in_rect(capture)), theme);
}

static bool run_picker(CaptureMode* out_mode) {
    if (!fonts::init()) {
        show_error("failed to initialize fonts");
        return false;
    }

    WsWindow win;
    if (!win.create("Screenshot", INIT_W, INIT_H)) {
        show_error("failed to create window");
        return false;
    }

    mtk::StandaloneHost host(&win);
    Canvas canvas = host.canvas();
    render_picker(canvas);
    host.present();

    bool accepted = false;
    while (true) {
        Montauk::WinEvent ev;
        int r = win.poll(&ev);

        if (r < 0) break;
        if (r == 0) {
            montauk::sleep_ms(16);
            continue;
        }

        bool redraw = false;
        if (ev.type == 3) {
            break;
        } else if (ev.type == 2) {
            g_win_w = win.width;
            g_win_h = win.height;
            redraw = true;
        } else if (ev.type == 4) {
            redraw = true;
        } else if (ev.type == 0 && ev.key.pressed) {
            if (ev.key.scancode == 0x01) break;
            if (ev.key.ascii == '\n' || ev.key.ascii == '\r') {
                accepted = true;
                break;
            }
            if (ev.key.ascii == '1') {
                g_mode = CAPTURE_SCREEN;
                redraw = true;
            } else if (ev.key.ascii == '2') {
                g_mode = CAPTURE_AREA;
                redraw = true;
            }
        } else if (ev.type == 1) {
            g_mouse_x = ev.mouse.x;
            g_mouse_y = ev.mouse.y;
            redraw = true;

            bool clicked = (ev.mouse.buttons & 1) && !(ev.mouse.prev_buttons & 1);
            if (clicked) {
                if (option_rect(0).contains(ev.mouse.x, ev.mouse.y)) {
                    g_mode = CAPTURE_SCREEN;
                } else if (option_rect(1).contains(ev.mouse.x, ev.mouse.y)) {
                    g_mode = CAPTURE_AREA;
                } else if (cancel_button_rect().contains(ev.mouse.x, ev.mouse.y)) {
                    break;
                } else if (capture_button_rect().contains(ev.mouse.x, ev.mouse.y)) {
                    accepted = true;
                    break;
                }
            }
        }

        if (redraw) {
            canvas = host.canvas();
            render_picker(canvas);
            host.present();
        }
    }

    if (out_mode) *out_mode = g_mode;
    win.destroy();
    return accepted;
}

// ============================================================================
// Framebuffer capture
// ============================================================================

static bool capture_screen(ScreenBuffer* out, char* err, int err_len) {
    if (!out) return false;
    out->pixels = nullptr;
    out->width = 0;
    out->height = 0;
    out->pitch = 0;

    Montauk::FbInfo fb;
    montauk::fb_info(&fb);

    int width = (int)fb.width;
    int height = (int)fb.height;
    int pitch = (int)fb.pitch;
    if (width <= 0 || height <= 0 || pitch <= 0) {
        safe_copy(err, err_len, "no framebuffer available");
        return false;
    }

    uint32_t* hw_fb = (uint32_t*)montauk::fb_map();
    if (!hw_fb) {
        safe_copy(err, err_len, "failed to map framebuffer");
        return false;
    }

    uint64_t size = (uint64_t)width * (uint64_t)height * sizeof(uint32_t);
    uint32_t* pixels = (uint32_t*)montauk::malloc(size);
    if (!pixels) {
        safe_copy(err, err_len, "out of memory");
        return false;
    }

    for (int y = 0; y < height; y++) {
        uint32_t* src = (uint32_t*)((uint8_t*)hw_fb + y * pitch);
        montauk::memcpy(pixels + y * width, src, (uint64_t)width * sizeof(uint32_t));
    }

    out->pixels = pixels;
    out->width = width;
    out->height = height;
    out->pitch = pitch;
    return true;
}

static uint32_t dim_pixel(uint32_t px) {
    uint8_t r = (uint8_t)((px >> 16) & 0xFF);
    uint8_t g = (uint8_t)((px >> 8) & 0xFF);
    uint8_t b = (uint8_t)(px & 0xFF);
    r = (uint8_t)((uint16_t)r * 3 / 5);
    g = (uint8_t)((uint16_t)g * 3 / 5);
    b = (uint8_t)((uint16_t)b * 3 / 5);
    return (px & 0xFF000000u) | ((uint32_t)r << 16) | ((uint32_t)g << 8) | b;
}

// ============================================================================
// Area selection window
// ============================================================================

struct AreaSelectorState {
    int win_w;
    int win_h;
    int mouse_x;
    int mouse_y;
    bool dragging;
    bool has_selection;
    bool fullscreen;
    int drag_start_x;
    int drag_start_y;
    Rect selection_view;
};

static int area_initial_w(const ScreenBuffer& screen) {
    int w = screen.width - 120;
    return gui_clamp(w, 360, 760);
}

static int area_initial_h(const ScreenBuffer& screen) {
    int h = screen.height - 140;
    return gui_clamp(h, 280, 560);
}

static Rect area_footer_rect(const AreaSelectorState& st) {
    return {0, st.win_h - FOOTER_H, st.win_w, FOOTER_H};
}

static Rect area_cancel_button_rect(const AreaSelectorState& st) {
    Rect footer = area_footer_rect(st);
    int w = gui_min(BUTTON_W, gui_max((st.win_w - PAD * 3) / 2, 72));
    int y = footer.y + (footer.h - BUTTON_H) / 2;
    return {st.win_w - PAD - w * 2 - 8, y, w, BUTTON_H};
}

static Rect area_capture_button_rect(const AreaSelectorState& st) {
    Rect footer = area_footer_rect(st);
    int w = gui_min(BUTTON_W, gui_max((st.win_w - PAD * 3) / 2, 72));
    int y = footer.y + (footer.h - BUTTON_H) / 2;
    return {st.win_w - PAD - w, y, w, BUTTON_H};
}

static Rect area_preview_rect(const ScreenBuffer& screen, const AreaSelectorState& st) {
    if (st.fullscreen) {
        return {0, 0, gui_max(st.win_w, 1), gui_max(st.win_h, 1)};
    }

    Rect content = {PAD, PAD,
                    gui_max(st.win_w - PAD * 2, 1),
                    gui_max(st.win_h - FOOTER_H - PAD * 2, 1)};
    if (screen.width <= 0 || screen.height <= 0) return content;

    int preview_w = content.w;
    int preview_h = (int)((uint64_t)preview_w * (uint64_t)screen.height / (uint64_t)screen.width);
    if (preview_h > content.h) {
        preview_h = content.h;
        preview_w = (int)((uint64_t)preview_h * (uint64_t)screen.width / (uint64_t)screen.height);
    }
    if (preview_w < 1) preview_w = 1;
    if (preview_h < 1) preview_h = 1;

    return {content.x + (content.w - preview_w) / 2,
            content.y + (content.h - preview_h) / 2,
            preview_w,
            preview_h};
}

static bool area_mouse_in(const AreaSelectorState& st, const Rect& rect) {
    return rect.contains(st.mouse_x, st.mouse_y);
}

static Rect clamp_rect_to_preview(const Rect& rect, const Rect& preview) {
    return rect.intersect(preview);
}

static bool draw_fullscreen_preview(Canvas& c,
                                    const ScreenBuffer& screen,
                                    const Rect& preview,
                                    const Rect& selection,
                                    bool has_selection) {
    if (!screen.pixels || preview.x != 0 || preview.y != 0 ||
        preview.w != c.w || preview.h != c.h ||
        screen.width != c.w || screen.height != c.h) {
        return false;
    }

    uint64_t bytes = (uint64_t)c.w * (uint64_t)c.h * sizeof(uint32_t);
    (void)selection;
    (void)has_selection;
    montauk::memcpy(c.pixels, screen.pixels, bytes);

    return true;
}

static void draw_scaled_preview(Canvas& c,
                                const ScreenBuffer& screen,
                                const Rect& preview,
                                const Rect& selection,
                                bool has_selection) {
    if (!screen.pixels || preview.empty()) return;
    if (draw_fullscreen_preview(c, screen, preview, selection, has_selection))
        return;

    Rect sel = clamp_rect_to_preview(selection, preview);
    for (int y = 0; y < preview.h; y++) {
        int dst_y = preview.y + y;
        if (dst_y < 0 || dst_y >= c.h) continue;
        int src_y = (int)((uint64_t)y * (uint64_t)screen.height / (uint64_t)preview.h);
        uint32_t* src_row = screen.pixels + src_y * screen.width;
        uint32_t* dst_row = c.pixels + dst_y * c.w;

        bool y_in_sel = has_selection && dst_y >= sel.y && dst_y < sel.y + sel.h;
        for (int x = 0; x < preview.w; x++) {
            int dst_x = preview.x + x;
            if (dst_x < 0 || dst_x >= c.w) continue;
            int src_x = (int)((uint64_t)x * (uint64_t)screen.width / (uint64_t)preview.w);
            uint32_t px = src_row[src_x];
            bool in_sel = y_in_sel && dst_x >= sel.x && dst_x < sel.x + sel.w;
            dst_row[dst_x] = (!has_selection || in_sel) ? px : dim_pixel(px);
        }
    }
}

static void draw_area_selector(Canvas& c,
                               const ScreenBuffer& screen,
                               const AreaSelectorState& st) {
    mtk::Theme theme = app_theme();
    if (!st.fullscreen)
        c.fill(theme.window_bg);

    Rect preview = area_preview_rect(screen, st);
    if (!st.fullscreen)
        c.fill_rect(preview.x - 1, preview.y - 1, preview.w + 2, preview.h + 2, theme.border);
    draw_scaled_preview(c, screen, preview, st.selection_view, st.has_selection || st.dragging);

    if ((st.has_selection || st.dragging) && !st.selection_view.empty()) {
        Rect sel = clamp_rect_to_preview(st.selection_view, preview);
        if (!sel.empty()) {
            c.rect(sel.x, sel.y, sel.w, sel.h, colors::WHITE);
            if (sel.w > 2 && sel.h > 2) {
                Rect inner = {sel.x + 1, sel.y + 1, sel.w - 2, sel.h - 2};
                c.rect(inner.x, inner.y, inner.w, inner.h, theme.accent);
            }
        }
    }

    if (st.fullscreen)
        return;

    Rect footer = area_footer_rect(st);
    c.fill_rect(footer.x, footer.y, footer.w, footer.h, theme.surface);
    mtk::draw_separator(c, 0, footer.y, st.win_w, theme);

    if (st.has_selection) {
        Rect screen_sel = {0, 0, 0, 0};
        Rect sel = clamp_rect_to_preview(st.selection_view, preview);
        if (!sel.empty()) {
            int left = (int)((uint64_t)(sel.x - preview.x) * (uint64_t)screen.width / (uint64_t)preview.w);
            int top = (int)((uint64_t)(sel.y - preview.y) * (uint64_t)screen.height / (uint64_t)preview.h);
            int right = (int)((uint64_t)(sel.x + sel.w - preview.x) * (uint64_t)screen.width / (uint64_t)preview.w);
            int bottom = (int)((uint64_t)(sel.y + sel.h - preview.y) * (uint64_t)screen.height / (uint64_t)preview.h);
            screen_sel = clamp_rect_to_screen({left, top, right - left, bottom - top},
                                              screen.width, screen.height);
        }

        char size_label[48];
        snprintf(size_label, sizeof(size_label), "%d x %d", screen_sel.w, screen_sel.h);
        c.text(PAD, footer.y + (footer.h - system_font_height()) / 2,
               size_label, theme.text_muted);
    }

    Rect cancel = area_cancel_button_rect(st);
    Rect capture = area_capture_button_rect(st);
    mtk::draw_button(c, cancel, "Cancel", mtk::BUTTON_SECONDARY,
                     mtk::widget_state(false, area_mouse_in(st, cancel)), theme);
    mtk::draw_button(c, capture, "Capture", mtk::BUTTON_PRIMARY,
                     mtk::widget_state(false, area_mouse_in(st, capture), st.has_selection), theme);
}

static Point clamp_point_to_rect(int x, int y, const Rect& rect) {
    return {gui_clamp(x, rect.x, rect.x + rect.w - 1),
            gui_clamp(y, rect.y, rect.y + rect.h - 1)};
}

static Rect view_selection_to_screen(const ScreenBuffer& screen,
                                     const Rect& preview,
                                     const Rect& selection) {
    Rect sel = clamp_rect_to_preview(selection, preview);
    if (sel.empty()) return {0, 0, 0, 0};

    int left = (int)((uint64_t)(sel.x - preview.x) * (uint64_t)screen.width / (uint64_t)preview.w);
    int top = (int)((uint64_t)(sel.y - preview.y) * (uint64_t)screen.height / (uint64_t)preview.h);
    int right = (int)((uint64_t)(sel.x + sel.w - preview.x) * (uint64_t)screen.width / (uint64_t)preview.w);
    int bottom = (int)((uint64_t)(sel.y + sel.h - preview.y) * (uint64_t)screen.height / (uint64_t)preview.h);
    if (right <= left) right = left + 1;
    if (bottom <= top) bottom = top + 1;

    return clamp_rect_to_screen({left, top, right - left, bottom - top},
                                screen.width, screen.height);
}

enum AreaEventResult : uint8_t {
    AREA_EVENT_CONTINUE = 0,
    AREA_EVENT_ACCEPT,
    AREA_EVENT_CANCEL,
};

static AreaEventResult handle_area_selector_event(WsWindow& win,
                                                  const ScreenBuffer& screen,
                                                  AreaSelectorState& st,
                                                  const Montauk::WinEvent& ev,
                                                  bool* redraw) {
    if (!redraw) return AREA_EVENT_CONTINUE;

    if (ev.type == 3) {
        return AREA_EVENT_CANCEL;
    }

    if (ev.type == 2) {
        st.win_w = win.width;
        st.win_h = win.height;
        st.selection_view = clamp_rect_to_preview(st.selection_view,
                                                  area_preview_rect(screen, st));
        *redraw = true;
        return AREA_EVENT_CONTINUE;
    }

    if (ev.type == 4) {
        *redraw = true;
        return AREA_EVENT_CONTINUE;
    }

    if (ev.type == 0 && ev.key.pressed) {
        if (ev.key.scancode == 0x01)
            return AREA_EVENT_CANCEL;
        if ((ev.key.ascii == '\n' || ev.key.ascii == '\r') && st.has_selection)
            return AREA_EVENT_ACCEPT;
        return AREA_EVENT_CONTINUE;
    }

    if (ev.type != 1)
        return AREA_EVENT_CONTINUE;

    int old_x = st.mouse_x;
    int old_y = st.mouse_y;
    st.mouse_x = ev.mouse.x;
    st.mouse_y = ev.mouse.y;
    if (old_x != st.mouse_x || old_y != st.mouse_y)
        *redraw = true;

    Rect preview = area_preview_rect(screen, st);
    bool left_pressed = (ev.mouse.buttons & 1) && !(ev.mouse.prev_buttons & 1);
    bool left_held = (ev.mouse.buttons & 1) != 0;
    bool left_released = !(ev.mouse.buttons & 1) && (ev.mouse.prev_buttons & 1);

    if (left_pressed) {
        if (!st.fullscreen) {
            Rect cancel = area_cancel_button_rect(st);
            Rect capture = area_capture_button_rect(st);
            if (cancel.contains(ev.mouse.x, ev.mouse.y))
                return AREA_EVENT_CANCEL;
            if (capture.contains(ev.mouse.x, ev.mouse.y) && st.has_selection)
                return AREA_EVENT_ACCEPT;
        }
        if (preview.contains(ev.mouse.x, ev.mouse.y)) {
            Point p = clamp_point_to_rect(ev.mouse.x, ev.mouse.y, preview);
            st.dragging = true;
            st.has_selection = false;
            st.drag_start_x = p.x;
            st.drag_start_y = p.y;
            st.selection_view = {p.x, p.y, 1, 1};
            *redraw = true;
        }
    } else if (st.dragging && left_held) {
        Point p = clamp_point_to_rect(ev.mouse.x, ev.mouse.y, preview);
        st.selection_view = clamp_rect_to_preview(
            normalized_rect(st.drag_start_x, st.drag_start_y, p.x, p.y),
            preview);
        *redraw = true;
    } else if (st.dragging && left_released) {
        Point p = clamp_point_to_rect(ev.mouse.x, ev.mouse.y, preview);
        st.selection_view = clamp_rect_to_preview(
            normalized_rect(st.drag_start_x, st.drag_start_y, p.x, p.y),
            preview);
        st.dragging = false;
        st.has_selection = st.selection_view.w >= MIN_AREA_SIZE
                        && st.selection_view.h >= MIN_AREA_SIZE;
        if (st.fullscreen && st.has_selection)
            return AREA_EVENT_ACCEPT;
        *redraw = true;
    }

    return AREA_EVENT_CONTINUE;
}

static bool run_area_selector(const ScreenBuffer& screen, Rect* out_area) {
    if (!screen.pixels || !out_area) return false;

    AreaSelectorState st = {};
    st.fullscreen = true;
    st.win_w = st.fullscreen ? screen.width : area_initial_w(screen);
    st.win_h = st.fullscreen ? screen.height : area_initial_h(screen);
    st.mouse_x = -1;
    st.mouse_y = -1;

    WsWindow win;
    if (!win.create("Select Area", st.win_w, st.win_h)) {
        show_error("failed to create area selection window");
        return false;
    }
    if (st.fullscreen)
        win.set_flags(Montauk::WIN_FLAG_FULLSCREEN);

    mtk::StandaloneHost host(&win);
    Canvas canvas = host.canvas();
    draw_area_selector(canvas, screen, st);
    host.present();

    bool accepted = false;
    while (true) {
        Montauk::WinEvent ev;
        int r = win.poll(&ev);

        if (r < 0) break;
        if (r == 0) {
            montauk::sleep_ms(16);
            continue;
        }

        bool redraw = false;
        AreaEventResult result = handle_area_selector_event(win, screen, st, ev, &redraw);

        if (result == AREA_EVENT_ACCEPT) {
            accepted = true;
            break;
        }
        if (result == AREA_EVENT_CANCEL)
            break;

        if (redraw) {
            canvas = host.canvas();
            draw_area_selector(canvas, screen, st);
            host.present();
        }
    }

    if (accepted) {
        Rect preview = area_preview_rect(screen, st);
        *out_area = view_selection_to_screen(screen, preview, st.selection_view);
        if (out_area->empty())
            accepted = false;
    }

    win.destroy();
    return accepted;
}

// ============================================================================
// JPEG encoding and saving
// ============================================================================

static void jpeg_write_callback(void* context, void* data, int size) {
    JpegBuffer* buf = (JpegBuffer*)context;
    if (!buf || size <= 0 || buf->failed) return;

    uint64_t needed = buf->size + (uint64_t)size;
    if (needed > buf->capacity) {
        uint64_t new_cap = buf->capacity * 2;
        if (new_cap < needed) new_cap = needed;
        uint8_t* new_data = (uint8_t*)montauk::realloc(buf->data, new_cap);
        if (!new_data) {
            buf->failed = true;
            return;
        }
        buf->data = new_data;
        buf->capacity = new_cap;
    }
    montauk::memcpy(buf->data + buf->size, data, size);
    buf->size += (uint64_t)size;
}

static bool encode_jpeg(const ScreenBuffer& screen,
                        const Rect& area,
                        JpegBuffer* out,
                        char* err,
                        int err_len) {
    if (!screen.pixels || !out) {
        safe_copy(err, err_len, "nothing to encode");
        return false;
    }

    Rect rect = clamp_rect_to_screen(area, screen.width, screen.height);
    if (rect.empty()) {
        safe_copy(err, err_len, "empty screenshot area");
        return false;
    }

    uint64_t rgb_size = (uint64_t)rect.w * (uint64_t)rect.h * 3;
    uint8_t* rgb = (uint8_t*)montauk::malloc(rgb_size);
    if (!rgb) {
        safe_copy(err, err_len, "out of memory");
        return false;
    }

    for (int y = 0; y < rect.h; y++) {
        uint32_t* src_row = screen.pixels + (rect.y + y) * screen.width + rect.x;
        uint8_t* dst_row = rgb + (uint64_t)y * rect.w * 3;
        for (int x = 0; x < rect.w; x++) {
            uint32_t px = src_row[x];
            dst_row[x * 3 + 0] = (uint8_t)((px >> 16) & 0xFF);
            dst_row[x * 3 + 1] = (uint8_t)((px >> 8) & 0xFF);
            dst_row[x * 3 + 2] = (uint8_t)(px & 0xFF);
        }
    }

    out->data = (uint8_t*)montauk::malloc(256 * 1024);
    out->size = 0;
    out->capacity = 256 * 1024;
    out->failed = false;
    if (!out->data) {
        montauk::mfree(rgb);
        safe_copy(err, err_len, "out of memory for JPEG buffer");
        return false;
    }

    int ok = stbi_write_jpg_to_func(jpeg_write_callback, out,
                                    rect.w, rect.h, 3, rgb, 100);
    montauk::mfree(rgb);

    if (!ok || out->failed || out->size == 0) {
        if (out->data) montauk::mfree(out->data);
        out->data = nullptr;
        out->size = 0;
        out->capacity = 0;
        safe_copy(err, err_len, "JPEG encoding failed");
        return false;
    }

    return true;
}

static void build_default_save_dir(char* out, int out_len) {
    char username[32] = {};
    if (montauk::getuser(username, sizeof(username)) > 0 && username[0]) {
        char home[128];
        snprintf(home, sizeof(home), "0:/users/%s", username);
        snprintf(out, out_len, "%s/Pictures", home);
        montauk::fmkdir(out);
        return;
    }
    safe_copy(out, out_len, "0:/");
}

static void build_suggested_name(char* out, int out_len) {
    Montauk::DateTime dt;
    montauk::gettime(&dt);
    snprintf(out, out_len,
             "screenshot_%04d%02d%02d_%02d%02d%02d.jpg",
             (int)dt.Year, (int)dt.Month, (int)dt.Day,
             (int)dt.Hour, (int)dt.Minute, (int)dt.Second);
}

static bool write_jpeg_file(const char* path, const JpegBuffer& jpeg, char* err, int err_len) {
    if (!path || !path[0]) {
        safe_copy(err, err_len, "no destination selected");
        return false;
    }
    if (!jpeg.data || jpeg.size == 0) {
        safe_copy(err, err_len, "no screenshot data to save");
        return false;
    }

    int fd = montauk::fcreate(path);
    if (fd < 0) {
        safe_copy(err, err_len, "failed to create file");
        return false;
    }

    int wrote = montauk::fwrite(fd, jpeg.data, 0, jpeg.size);
    montauk::close(fd);
    if (wrote < 0 || (uint64_t)wrote != jpeg.size) {
        safe_copy(err, err_len, "failed to write file");
        return false;
    }
    return true;
}

static bool save_with_dialog(const JpegBuffer& jpeg) {
    char save_dir[192];
    char suggested[128];
    char path[256];
    char message[160];
    char err[160];

    build_default_save_dir(save_dir, sizeof(save_dir));
    build_suggested_name(suggested, sizeof(suggested));
    path[0] = '\0';
    message[0] = '\0';
    err[0] = '\0';

    bool ok = gui::dialogs::save_file("Save Screenshot", save_dir, suggested,
                                      path, sizeof(path), message, sizeof(message));
    if (!ok) {
        if (message[0]) show_error(message);
        return false;
    }

    if (!str_ends_with_ci(path, ".jpg") && !str_ends_with_ci(path, ".jpeg")) {
        int len = montauk::slen(path);
        if (len + 4 < (int)sizeof(path)) {
            path[len + 0] = '.';
            path[len + 1] = 'j';
            path[len + 2] = 'p';
            path[len + 3] = 'g';
            path[len + 4] = '\0';
        }
    }

    if (!write_jpeg_file(path, jpeg, err, sizeof(err))) {
        show_error(err);
        return false;
    }

    char msg[320];
    snprintf(msg, sizeof(msg), "Screenshot saved: %s (%llu bytes)\n",
             path, (unsigned long long)jpeg.size);
    montauk::print(msg);
    return true;
}

// ============================================================================
// Entry point
// ============================================================================

extern "C" void _start() {
    CaptureMode mode = CAPTURE_SCREEN;
    if (!run_picker(&mode))
        montauk::exit(0);

    // Let the desktop remove the picker window before reading the screen.
    montauk::sleep_ms(140);

    ScreenBuffer screen = {};
    char err[160];
    err[0] = '\0';
    if (!capture_screen(&screen, err, sizeof(err))) {
        show_error(err);
        montauk::exit(1);
    }

    Rect area = {0, 0, screen.width, screen.height};
    if (mode == CAPTURE_AREA) {
        if (!run_area_selector(screen, &area)) {
            free_screen(&screen);
            montauk::exit(0);
        }
        // Let the desktop remove the selector window before opening the save dialog.
        montauk::sleep_ms(80);
    }

    JpegBuffer jpeg = {};
    if (!encode_jpeg(screen, area, &jpeg, err, sizeof(err))) {
        free_screen(&screen);
        show_error(err);
        montauk::exit(1);
    }
    free_screen(&screen);

    save_with_dialog(jpeg);
    if (jpeg.data) montauk::mfree(jpeg.data);
    montauk::exit(0);
}