#include <X11/Xlib.h>
#include <X11/Xutil.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

static unsigned long
lowbit(unsigned long x)
{
	return x & (~x + 1UL);
}

static unsigned char
chan(unsigned long pixel, unsigned long mask)
{
	unsigned long base;
	unsigned long max;
	unsigned long v;

	if (mask == 0)
	    return 0;
	base = lowbit(mask);
	max = mask / base;
	v = (pixel & mask) / base;
	if (max == 0)
	    return 0;
	return (unsigned char)((v * 255UL + (max / 2UL)) / max);
}

int
main(int argc, char **argv)
{
	Display *dpy;
	Window root;
	XImage *img;
	FILE *f;
	const char *geom;
	const char *outpath;
	unsigned int w;
	unsigned int h;
	int x;
	int y;
	int sw;
	int sh;
	int ix;
	int iy;
	char tail;
	char geombuf[128];

	if (argc == 3) {
	    geom = argv[1];
	    outpath = argv[2];
	} else if (argc == 2) {
	    size_t n;

	    outpath = argv[1];
	    if (!isatty(0) && fgets(geombuf, sizeof(geombuf), stdin) != NULL) {
	        n = 0;
	        while (geombuf[n] != '\0')
	            n++;
	        while (n > 0 && (geombuf[n - 1] == '\n' || geombuf[n - 1] == '\r')) {
	            geombuf[n - 1] = '\0';
	            n--;
	        }
	        geom = geombuf;
	    } else {
	        geom = NULL;
	    }
	} else {
	    fprintf(stderr, "usage: ss WIDTHxHEIGHT+X+Y file.ppm | ss file.ppm < geometry\n");
	    return 1;
	}
	if ((dpy = XOpenDisplay(NULL)) == NULL) {
	    fprintf(stderr, "failed to open X display\n");
	    return 1;
	}

	sw = DisplayWidth(dpy, DefaultScreen(dpy));
	sh = DisplayHeight(dpy, DefaultScreen(dpy));
	if (geom != NULL) {
	    if (sscanf(geom, "%ux%u+%d+%d%c", &w, &h, &x, &y, &tail) != 4 || w == 0 || h == 0) {
	        fprintf(stderr, "invalid geometry: %s\n", geom);
	        XCloseDisplay(dpy);
	        return 1;
	    }
	    if (x < 0 || y < 0 || x + (int)w > sw || y + (int)h > sh) {
	        fprintf(stderr, "geometry out of bounds (%dx%d screen)\n", sw, sh);
	        XCloseDisplay(dpy);
	        return 1;
	    }
	} else {
	    x = 0;
	    y = 0;
	    w = (unsigned int)sw;
	    h = (unsigned int)sh;
	}

	root = DefaultRootWindow(dpy);
	img = XGetImage(dpy, root, x, y, w, h, AllPlanes, ZPixmap);
	if (img == NULL) {
	    fprintf(stderr, "XGetImage failed\n");
	    XCloseDisplay(dpy);
	    return 1;
	}

	f = fopen(outpath, "wb");
	if (f == NULL) {
	    fprintf(stderr, "failed to open output: %s\n", outpath);
	    XDestroyImage(img);
	    XCloseDisplay(dpy);
	    return 1;
	}

	fprintf(f, "P6\n%u %u\n255\n", w, h);
	for (iy = 0; iy < (int)h; iy++) {
	    for (ix = 0; ix < (int)w; ix++) {
	        unsigned long p;
	        unsigned char rgb[3];

	        p = XGetPixel(img, ix, iy);
	        rgb[0] = chan(p, img->red_mask);
	        rgb[1] = chan(p, img->green_mask);
	        rgb[2] = chan(p, img->blue_mask);
	        if (fwrite(rgb, 1, 3, f) != 3) {
	            fprintf(stderr, "write failed: %s\n", outpath);
	            fclose(f);
	            XDestroyImage(img);
	            XCloseDisplay(dpy);
	            return 1;
	        }
	    }
	}

	fclose(f);
	XDestroyImage(img);
	XCloseDisplay(dpy);
	return 0;
}