/*
 * xscreenshot.c -- simple program to output a screenshot of an X11
 * display as a PPM image.  Only supports 32bpp TrueColor displays.
 *
 * Written by Andrew Church <achurch@achurch.org>
 * This source code is public domain.
 *
 * ==============================================
 *
 * This source code is written to the C99 standard.  Some compilers
 * (such as GCC) may require a command-line flag to enable C99 support.
 *
 * ==============================================
 *
 * Version history:
 *
 * Version 1.0, 2013-04-15
 *   - Initial release.
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>

int main(int argc, char **argv)
{
    const char *display_name;
    if (argc > 2 || (argc == 2 && argv[1][0] == '-')) {
        fprintf(stderr, "Usage: %s [display] >screenshot.ppm\n", argv[0]);
        return 2;
    }
    if (argc >= 2) {
        display_name = argv[1];
    } else {
        display_name = getenv("DISPLAY");
        if (!display_name || !*display_name) {
            fprintf(stderr, "DISPLAY environment variable not set, can't take"
                    " screenshot\n");
            return 1;
        }
    }

    Display *display = XOpenDisplay(display_name);
    if (!display) {
        fprintf(stderr, "Unable to open display %s\n", display_name);
        return 1;
    }
    const int screen = XDefaultScreen(display);
    const int width = DisplayWidth(display, screen);
    const int height = DisplayHeight(display, screen);
    const Window window = RootWindow(display, screen);

    unsigned long red_mask, green_mask, blue_mask;
    int red_shift, green_shift, blue_shift;
    int red_bits, green_bits, blue_bits;
    XWindowAttributes window_attr;
    if (!XGetWindowAttributes(display, window, &window_attr)) {
        fprintf(stderr, "Unable to retrieve root window attributes\n");
        XCloseDisplay(display);
        return 1;
    }
    if (window_attr.visual->class != TrueColor) {
        fprintf(stderr, "Visual class %d unsupported\n",
                window_attr.visual->class);
    }
    red_mask = window_attr.visual->red_mask;
    green_mask = window_attr.visual->green_mask;
    blue_mask = window_attr.visual->blue_mask;
    for (red_shift = 0; !(red_mask & 1); red_shift++, red_mask >>= 1) {}
    for (green_shift = 0; !(green_mask & 1); green_shift++, green_mask >>= 1) {}
    for (blue_shift = 0; !(blue_mask & 1); blue_shift++, blue_mask >>= 1) {}
    for (red_bits = 0; red_mask & (1 << red_bits); red_bits++) {}
    for (green_bits = 0; green_mask & (1 << green_bits); green_bits++) {}
    for (blue_bits = 0; blue_mask & (1 << blue_bits); blue_bits++) {}

    XImage *image =
        XGetImage(display, window, 0, 0, width, height, AllPlanes, ZPixmap);
    if (!image) {
        fprintf(stderr, "Unable to copy root window image\n");
        XCloseDisplay(display);
        return 1;
    }
    if (image->bitmap_unit != 32) {
        fprintf(stderr, "%dbpp display not supported\n", image->bitmap_unit);
        XDestroyImage(image);
        XCloseDisplay(display);
        return 1;
    }

    const long npixels = (long)width * (long)height;
    unsigned char *pixels = malloc(npixels*6);
    if (!pixels) {
        fprintf(stderr, "Unable to allocate memory for pixel data (%dx%d)\n",
                width, height);
        XDestroyImage(image);
        XCloseDisplay(display);
        return 1;
    }
    unsigned int maxval;
    if (red_bits > 8 || green_bits > 8 || blue_bits > 8) {
        maxval = 65535;
        const uint32_t *in = (const uint32_t *)image->data;
        unsigned char *out = pixels;
        for (int y = 0; y < height; y++, in += image->bytes_per_line/4) {
            for (int x = 0; x < width; x++, out += 6) {
                const uint32_t pixel = in[x];
                unsigned long red = (pixel >> red_shift) & red_mask;
                unsigned long green = (pixel >> green_shift) & green_mask;
                unsigned long blue = (pixel >> blue_shift) & blue_mask;
                red = (red << (16-red_bits)) | (red >> (2*red_bits-16));
                green = (green << (16-green_bits)) | (green >> (2*green_bits-16));
                blue = (blue << (16-blue_bits)) | (blue >> (2*blue_bits-16));
                out[0] = red >> 8;
                out[1] = red >> 0;
                out[2] = green >> 8;
                out[3] = green >> 0;
                out[4] = blue >> 8;
                out[5] = blue >> 0;
            }
        }
    } else {
        maxval = 255;
        const uint32_t *in = (const uint32_t *)image->data;
        unsigned char *out = pixels;
        for (int y = 0; y < height; y++, in += image->bytes_per_line/4) {
            for (int x = 0; x < width; x++, out += 3) {
                const uint32_t pixel = in[x];
                unsigned long red = (pixel >> red_shift) & red_mask;
                unsigned long green = (pixel >> green_shift) & green_mask;
                unsigned long blue = (pixel >> blue_shift) & blue_mask;
                red = (red << (8-red_bits)) | (red >> (2*red_bits-8));
                green = (green << (8-green_bits)) | (green >> (2*green_bits-8));
                blue = (blue << (8-blue_bits)) | (blue >> (2*blue_bits-8));
                out[0] = red;
                out[1] = green;
                out[2] = blue;
            }
        }
    }

    XDestroyImage(image);
    XCloseDisplay(display);

    if (printf("P6\n%d %d 255\n", width, height) < 0 ) {
        perror("Unable to write PPM header");
        free(pixels);
        return 1;
    }
    if (fwrite(pixels, maxval==65535 ? 6 : 3, npixels, stdout) != npixels) {
        perror("Unable to write image data");
        free(pixels);
        return 1;
    }

    free(pixels);
    return 0;
}