/* 
 * Purpose: Make screenshot of X11 applications and output as png.
 * Usage:   xscreenshot [winid] > file.png
 * Author:  Hiltjo Posthuma <hiltjo@codemadness.org>
 * License: WTFPL
 * Compile: cc xscreenshot.c -o xscreenshot -lX11 -lpng
 * TODO:    I might have missed a few cases of color conversion (convertrow_*)
 */

#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

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

/* png */
#include <png.h>

static void
die(const char *s) {
	fprintf(stderr, "xscreenshot: %s\n", s);
	exit(EXIT_FAILURE);
}

/* LSBFirst: BGRA -> RGBA */
static void
convertrow_lsb(unsigned char *drow, unsigned char *srow, XImage *img) {
	int sx, dx;

	for(sx = 0, dx = 0; dx < img->bytes_per_line - 4; sx += 4) {
		drow[dx++] = srow[sx + 2]; /* B -> R */
		drow[dx++] = srow[sx + 1]; /* G -> G */
		drow[dx++] = srow[sx];     /* R -> B */
		if(img->depth == 32)
			drow[dx++] = srow[sx + 3]; /* A -> A */
		else
			drow[dx++] = 255;
	}
}

/* MSBFirst: ARGB -> RGBA */
static void
convertrow_msb(unsigned char *drow, unsigned char *srow, XImage *img) {
	int sx, dx;

	for(sx = 0, dx = 0; dx < img->bytes_per_line - 4; sx += 4) {
		drow[dx++] = srow[sx + 1]; /* G -> R */
		drow[dx++] = srow[sx + 2]; /* B -> G */
		drow[dx++] = srow[sx + 3]; /* A -> B */
		if(img->depth == 32)
			drow[dx++] = srow[sx]; /* R -> A */
		else
			drow[dx++] = 255;
	}
}

static void
pngstdout(XImage *img)
{
	png_structp png_struct_p;
	png_infop png_info_p;
	void (*convert)(unsigned char *, unsigned char *, XImage *);
	unsigned char *drow = NULL, *srow;
	int h;

	png_struct_p = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL,
	                                       NULL);
	png_info_p = png_create_info_struct(png_struct_p);

	if(!png_struct_p || !png_info_p || setjmp(png_jmpbuf(png_struct_p)))
		die("failed to initialize libpng");

	png_init_io(png_struct_p, stdout);
	png_set_IHDR(png_struct_p, png_info_p, img->width, img->height, 8,
	             PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
	             PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	png_write_info(png_struct_p, png_info_p);

	srow = (unsigned char *)img->data;
	drow = calloc(1, img->width * 4); /* output RGBA */
	if(!drow)
		die("Can't calloc");

	if(img->byte_order == LSBFirst)
		convert = convertrow_lsb;
	else
		convert = convertrow_msb;

	for(h = 0; h < img->height; h++) {
		convert(drow, srow, img);
		srow += img->bytes_per_line;
		png_write_row(png_struct_p, drow);
	}
	png_write_end(png_struct_p, NULL);

	free(drow);
	png_free_data(png_struct_p, png_info_p, PNG_FREE_ALL, -1);
	png_destroy_write_struct(&png_struct_p, NULL);
}

int
main(int argc, char *argv[])
{
	XImage *img;
	Display *dpy;
	Window win;
	XWindowAttributes attr;

	dpy = XOpenDisplay(NULL);
	if(!dpy)
		die("Can't XOpenDisplay");

	/* win */
	if(argc > 1) {
		errno = 0;
		win = (Window)strtol(argv[1], NULL, 0);
		if(errno != 0)
			die("input not a number");
	} else {
		win = RootWindow(dpy, 0);
	}

	XGrabServer(dpy);
	XGetWindowAttributes(dpy, win, &attr);

	img = XGetImage(dpy, win, 0, 0, attr.width, attr.height, 0xffffffff,
	                ZPixmap);
	if(!img)
		die("Can't XGetImage");
	pngstdout(img);
	XDestroyImage(img);
	XUngrabServer(dpy);
	XCloseDisplay(dpy);

	return EXIT_SUCCESS;
}