/*
 * orionshot — X11 screenshot tool for OrionOS
 * Usage: orionshot full        full-screen PNG capture
 *        orionshot region      interactive rubber-band region select
 *
 * Saves to ~/Pictures/Screenshots/orionos_YYYYMMDD_HHMMSS.png
 * Sends a desktop notification on success.
 */
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/cursorfont.h>
#include <png.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <sys/stat.h>

#ifndef MIN
#define MIN(a,b) ((a)<(b)?(a):(b))
#endif

/* ── PNG writer ──────────────────────────────────────────────────────── */

static int get_shift(unsigned long mask) {
    if (!mask) return 0;
    int s = 0;
    while (!(mask & 1)) { mask >>= 1; s++; }
    return s;
}

static int write_png(const char *path, XImage *img, int ox, int oy, int w, int h) {
    FILE *fp = fopen(path, "wb");
    if (!fp) { perror("orionshot: fopen"); return -1; }

    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!png) { fclose(fp); return -1; }
    png_infop info = png_create_info_struct(png);
    if (!info) { png_destroy_write_struct(&png, NULL); fclose(fp); return -1; }
    if (setjmp(png_jmpbuf(png))) {
        png_destroy_write_struct(&png, &info);
        fclose(fp);
        return -1;
    }

    png_init_io(png, fp);
    png_set_IHDR(png, info, (unsigned)w, (unsigned)h, 8,
                 PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
    png_write_info(png, info);

    int r_shift = get_shift(img->red_mask);
    int g_shift = get_shift(img->green_mask);
    int b_shift = get_shift(img->blue_mask);
    int r_scale = (img->red_mask   >> r_shift) < 0xFF ? 2 : 1;
    int g_scale = (img->green_mask >> g_shift) < 0xFF ? 2 : 1;
    int b_scale = (img->blue_mask  >> b_shift) < 0xFF ? 2 : 1;
    (void)r_scale; (void)g_scale; (void)b_scale;

    unsigned char *row = malloc((size_t)w * 3);
    if (!row) { png_destroy_write_struct(&png, &info); fclose(fp); return -1; }

    for (int j = 0; j < h; j++) {
        for (int i = 0; i < w; i++) {
            unsigned long px = XGetPixel(img, ox + i, oy + j);
            row[i*3+0] = (unsigned char)((px & img->red_mask)   >> r_shift);
            row[i*3+1] = (unsigned char)((px & img->green_mask) >> g_shift);
            row[i*3+2] = (unsigned char)((px & img->blue_mask)  >> b_shift);
        }
        png_write_row(png, row);
    }
    free(row);
    png_write_end(png, NULL);
    png_destroy_write_struct(&png, &info);
    fclose(fp);
    return 0;
}

/* ── output path ─────────────────────────────────────────────────────── */

static void make_save_path(char *buf, size_t bufsz) {
    const char *home = getenv("HOME");
    if (!home) home = "/tmp";
    char dir[512];
    snprintf(dir, sizeof(dir), "%s/Pictures/Screenshots", home);
    mkdir(dir, 0755);

    time_t t = time(NULL);
    struct tm *tm = localtime(&t);
    char ts[32];
    strftime(ts, sizeof(ts), "%Y%m%d_%H%M%S", tm);
    snprintf(buf, bufsz, "%s/orionos_%s.png", dir, ts);
}

/* ── region selection ────────────────────────────────────────────────── */

static int select_region(Display *dpy, Window root,
                          int *rx, int *ry, int *rw, int *rh) {
    Cursor cross = XCreateFontCursor(dpy, XC_crosshair);
    if (XGrabPointer(dpy, root, False,
                     ButtonPressMask | ButtonReleaseMask | PointerMotionMask,
                     GrabModeAsync, GrabModeAsync,
                     root, cross, CurrentTime) != GrabSuccess) {
        XFreeCursor(dpy, cross);
        return -1;
    }
    XGrabKeyboard(dpy, root, False, GrabModeAsync, GrabModeAsync, CurrentTime);

    GC gc;
    XGCValues gcv;
    gcv.function       = GXxor;
    gcv.foreground     = 0xFFFFFF;
    gcv.subwindow_mode = IncludeInferiors;
    gcv.line_width     = 1;
    gc = XCreateGC(dpy, root,
                   GCFunction | GCForeground | GCSubwindowMode | GCLineWidth, &gcv);

    int sx = 0, sy = 0, ex = 0, ey = 0;
    int started = 0, drawn = 0, done = 0, cancelled = 0;

    while (!done) {
        XEvent e;
        XNextEvent(dpy, &e);

        if (e.type == ButtonPress && e.xbutton.button == Button1) {
            sx = e.xbutton.x_root; sy = e.xbutton.y_root;
            ex = sx; ey = sy; started = 1;

        } else if (e.type == MotionNotify && started) {
            /* erase previous rubber band */
            if (drawn) {
                int dx = MIN(sx,ex), dy = MIN(sy,ey);
                int dw = abs(ex-sx), dh = abs(ey-sy);
                if (dw > 0 && dh > 0)
                    XDrawRectangle(dpy, root, gc, dx, dy, (unsigned)dw, (unsigned)dh);
                XFlush(dpy); drawn = 0;
            }
            ex = e.xmotion.x_root; ey = e.xmotion.y_root;
            {
                int dx = MIN(sx,ex), dy = MIN(sy,ey);
                int dw = abs(ex-sx), dh = abs(ey-sy);
                if (dw > 0 && dh > 0) {
                    XDrawRectangle(dpy, root, gc, dx, dy, (unsigned)dw, (unsigned)dh);
                    XFlush(dpy); drawn = 1;
                }
            }

        } else if (e.type == ButtonRelease && e.xbutton.button == Button1) {
            if (drawn) {
                int dx = MIN(sx,ex), dy = MIN(sy,ey);
                int dw = abs(ex-sx), dh = abs(ey-sy);
                if (dw > 0 && dh > 0)
                    XDrawRectangle(dpy, root, gc, dx, dy, (unsigned)dw, (unsigned)dh);
                XFlush(dpy);
            }
            ex = e.xbutton.x_root; ey = e.xbutton.y_root;
            done = 1;

        } else if (e.type == KeyPress) {
            cancelled = 1; done = 1;
        }
    }

    XUngrabPointer(dpy, CurrentTime);
    XUngrabKeyboard(dpy, CurrentTime);
    XFreeGC(dpy, gc);
    XFreeCursor(dpy, cross);
    XSync(dpy, False);

    if (cancelled) return -1;

    *rx = MIN(sx,ex); *ry = MIN(sy,ey);
    *rw = abs(ex-sx); *rh = abs(ey-sy);
    if (*rw < 2 || *rh < 2) return -1;
    return 0;
}

/* ── main ────────────────────────────────────────────────────────────── */

int main(int argc, char *argv[]) {
    int do_region = (argc > 1 && strcmp(argv[1], "region") == 0);

    Display *dpy = XOpenDisplay(NULL);
    if (!dpy) { fprintf(stderr, "orionshot: cannot open display\n"); return 1; }
    int screen = DefaultScreen(dpy);
    Window root = RootWindow(dpy, screen);
    int sw = DisplayWidth(dpy, screen);
    int sh = DisplayHeight(dpy, screen);

    int cx = 0, cy = 0, cw = sw, ch = sh;

    if (do_region) {
        if (select_region(dpy, root, &cx, &cy, &cw, &ch) < 0) {
            XCloseDisplay(dpy);
            return 0; /* cancelled — silently exit */
        }
        /* tiny pause so rubber-band XOR erasure is fully flushed */
        usleep(60000);
    }

    XImage *img = XGetImage(dpy, root, cx, cy, (unsigned)cw, (unsigned)ch,
                             AllPlanes, ZPixmap);
    XCloseDisplay(dpy);

    if (!img) { fprintf(stderr, "orionshot: XGetImage failed\n"); return 1; }

    char path[1024];
    make_save_path(path, sizeof(path));

    if (write_png(path, img, 0, 0, cw, ch) < 0) {
        XDestroyImage(img);
        fprintf(stderr, "orionshot: failed to write %s\n", path);
        return 1;
    }
    XDestroyImage(img);

    /* extract just the filename for the notification body */
    const char *fname = strrchr(path, '/');
    fname = fname ? fname + 1 : path;

    char notif[512];
    snprintf(notif, sizeof(notif),
             "notify-send -t 3000 'Screenshot' '%s' 2>/dev/null &", fname);
    system(notif);

    fprintf(stderr, "orionshot: saved %s\n", path);
    return 0;
}