/* fbgrab-viewer - a remote viewer for fbgrab running in network mode
 *
 * (C) 2006 by Harald Welte <laforge@gnumonks.org>
 *
 * This program is Free Software and distributed under the terms of the GNU
 * General Public License, Version 2, as published by the Free Software Foundation.
 */

#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <netdb.h>
#include <dirent.h>

#include <png.h>

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

#define DEFAULT_WIN_WIDTH 240
#define DEFAULT_WIN_HEIGHT 320

const int LISTEN_QUEUE = 128;

int
listen_server(const char *hostname,
	      const char *service, int family, int socktype)
{
	struct addrinfo hints, *res, *ressave;
	int n, sockfd;

	memset(&hints, 0, sizeof(struct addrinfo));
	hints.ai_flags = AI_PASSIVE;
	hints.ai_family = family;
	hints.ai_socktype = socktype;

	n = getaddrinfo(hostname, service, &hints, &res);

	if (n < 0) {
		fprintf(stderr, "getaddrinfo error:: [%s]\n", gai_strerror(n));
		return -1;
	}

	ressave = res;

	sockfd = -1;
	while (res) {
		sockfd = socket(res->ai_family,
				res->ai_socktype, res->ai_protocol);

		if (!(sockfd < 0)) {
			if (bind(sockfd, res->ai_addr, res->ai_addrlen) == 0)
				break;

			close(sockfd);
			sockfd = -1;
		}
		res = res->ai_next;
	}

	if (sockfd < 0) {
		freeaddrinfo(ressave);
		fprintf(stderr, "socket error:: could not open socket\n");
		return -1;
	}

	listen(sockfd, LISTEN_QUEUE);

	freeaddrinfo(ressave);

	return sockfd;
}

#define NET_FILE_HDR_MAGIC	0x23421102

struct net_file_hdr {
	u_int32_t magic;
	u_int32_t size;
};


static int net_read_hdr(int fd)
{
	struct net_file_hdr hdr;
	int rc;

	rc = read(fd, &hdr, sizeof(hdr));
	if (rc < 0)
		return rc;

	if (ntohl(hdr.magic) != NET_FILE_HDR_MAGIC) {
		fprintf(stderr, "invalid magic\n");
		return -1;
	}

	return ntohl(hdr.size);
}

struct user_read_private {
	void *buf;
	int buf_size;
	int pos;
};

static void user_read_data(png_structp png_ptr,
			   png_bytep data, png_size_t length)
{
	struct user_read_private *urp;

	urp = png_get_io_ptr(png_ptr);
	if (length > urp->buf_size - urp->pos) {
		printf("shortening length from %u to %u\n",
			length, urp->buf_size - urp->pos);
		length = urp->buf_size - urp->pos;
	}

	memcpy(data, urp->buf+urp->pos, length);
	urp->pos += length;
}


/* The following two functions are from Xlib_JPEG_Example.c
 * by George Peter Staplin, thanks to him for making the code public domain.
 * We just made some few additions to handle RGBA buffers.
 *
 * You can read a useful tutorial here:
 * http://www.xmission.com/~georgeps/documentation/tutorials/Xlib_JPEG_Tutorial-5.html
 */
int get_byte_order (void) {
	union {
		char c[sizeof(short)];
		short s;
	} order;

	order.s = 1;
	if ((1 == order.c[0])) {
		return LSBFirst;
	} else {
		return MSBFirst;
	}
}

/* NOTE: The last argument 'buffer_depth' is not present in the original code */
XImage *create_image_from_buffer (Display *dis, int screen, u_char *buf, int width, int height, int buffer_depth) {
	int depth;
	XImage *img = NULL;
	Visual *vis;
	double rRatio;
	double gRatio;
	double bRatio;
	int outIndex = 0;
	int i;
	int numBufBytes = (buffer_depth * (width * height));

	depth = DefaultDepth (dis, screen);
	vis = DefaultVisual (dis, screen);

	rRatio = vis->red_mask / 255.0;
	gRatio = vis->green_mask / 255.0;
	bRatio = vis->blue_mask / 255.0;

	if (depth >= 24) {
		size_t numNewBufBytes = (4 * (width * height));
    fprintf(stderr, "numNewBufBytes:%d\n", numNewBufBytes);
		u_int32_t *newBuf = malloc (numNewBufBytes);

		for (i = 0; i < numBufBytes; ++i) {
			unsigned int r, g, b;
			r = (buf[i] * rRatio);
			++i;
			g = (buf[i] * gRatio);
			++i;
			b = (buf[i] * bRatio);

			/* Not in the original code,
			 * consider the alpha component,
			 * for now just ignore it.
			 */
			if (buffer_depth > 3) ++i;

			r &= vis->red_mask;
			g &= vis->green_mask;
			b &= vis->blue_mask;

			newBuf[outIndex] = r | g | b;
			++outIndex;
		}

		img = XCreateImage (dis,
			CopyFromParent, depth,
			ZPixmap, 0,
			(char *) newBuf,
			width, height,
			32, 0
		);

	} else if (depth >= 15) {
		size_t numNewBufBytes = (2 * (width * height));
		u_int16_t *newBuf = malloc (numNewBufBytes);

		for (i = 0; i < numBufBytes; ++i) {
			unsigned int r, g, b;

			r = (buf[i] * rRatio);
			++i;
			g = (buf[i] * gRatio);
			++i;
			b = (buf[i] * bRatio);

			/* Not in the original code,
			 * consider the alpha component,
			 * for now just ignore it.
			 */
			if (buffer_depth > 3) ++i;

			r &= vis->red_mask;
			g &= vis->green_mask;
			b &= vis->blue_mask;

			newBuf[outIndex] = r | g | b;
			++outIndex;
		}

		img = XCreateImage (dis,
			CopyFromParent, depth,
			ZPixmap, 0,
			(char *) newBuf,
			width, height,
			16, 0
		);
	} else {
		fprintf (stderr, "This program does not support displays with a depth less than 15.");
		return NULL;
	}

	XInitImage (img);
	/*Set the client's byte order, so that XPutImage knows what to do with the data.*/
	/*The default in a new X image is the server's format, which may not be what we want.*/
	if ((LSBFirst == get_byte_order ())) {
		img->byte_order = LSBFirst;
	} else {
		img->byte_order = MSBFirst;
	}

	/*The bitmap_bit_order doesn't matter with ZPixmap images.*/
	img->bitmap_bit_order = MSBFirst;

	return img;
}

void usage(char *progname)
{
	printf("usage: %s [-m] [-d dir] [-g WxH]\n"
		"\t\t -m\tenables movie mode (save every frame as png)\n"
		"\t\t -d dir\tsave files in directory 'dir'\n"
		"\t\t -g WxH\tset window Width and Height in pixels\n"
		, progname);
}

int main(int argc, char **argv)
{
	int sockfd, listenfd =
			listen_server(NULL, "2342", AF_UNSPEC, SOCK_STREAM);
	int tmpfd;
	int rc, size;
	void *png;
	png_struct *pp;
	png_info *linfo;
	Display *xdisp;
	Window xwin;
	GC xgc;
	XGCValues xgcv;
	int movie_mode = 0;
	char frames_basepath[PATH_MAX] = "./\0";
	int win_width = DEFAULT_WIN_WIDTH;
	int win_height = DEFAULT_WIN_HEIGHT;

	char c;

	while( (c=getopt(argc, argv, "hmd:g:")) > 0) {
		switch(c) {

			case 'h':
				usage(argv[0]);
				exit(0);
				break;

			case 'm':
				movie_mode = 1;
				break;

			case 'd':
				snprintf(frames_basepath, PATH_MAX, "%s", optarg);
				DIR *d = opendir(frames_basepath);
				if (d == NULL) {
					perror(frames_basepath);
					exit(1);
				}
				closedir(d);
				break;

			case 'g':
				if (sscanf(optarg, "%dx%d\n", &win_width, &win_height) != 2) {
					usage(argv[0]);
					exit(1);
				}
				break;

			case '?':
				usage(argv[0]);
				exit(1);
				break;

			default:
				break;
		}
	}

	xdisp = XOpenDisplay(NULL);
	if (!xdisp) {
		fprintf(stderr, "can't open X11 display\n");
		exit(1);
	}

	xwin = XCreateSimpleWindow(xdisp, DefaultRootWindow(xdisp),
	                           0, 0, win_width, win_height, 0, 0, 0);
	XSelectInput(xdisp, xwin, StructureNotifyMask);

	XMapWindow(xdisp, xwin);

	memset(&xgcv, 0, sizeof(xgcv));
	xgcv.function = GXcopy;
	xgc = XCreateGC(xdisp, xwin, GCFunction, &xgcv);

	for(;;) {
		XEvent e;
		XNextEvent(xdisp, &e);
		if (e.type == MapNotify)
			break;
	}

	if (listenfd < 0)
		exit(1);

	tmpfd = creat("/tmp/viewer.png", 0600);
	if (tmpfd < 0)
		exit(1);

	printf("waiting for connection from fbgrab\n");

	sockfd = accept(listenfd, NULL, NULL);

	printf("connected!\n");

	while (1) {
		int bpp, pass, num_passes;
		int _read = 0, written = 0;
		unsigned char *pixel;
		struct user_read_private urp;
		XImage *ximg;

		printf("reading hdr\n");
		size = net_read_hdr(sockfd);
		if (size < 0)
			exit(1);

		png = malloc(size);
		if (!png)
			exit(1);
		while (_read != size) {
			rc = read(sockfd, png+_read, size-_read);
			if (rc < 0)
				exit(1);
			_read += rc;
		}

		lseek(tmpfd, SEEK_SET, 0);

		while (written != size) {
			rc = write(tmpfd, png+written, size-written);
			if (rc < 0)
				exit(1);
			written += rc;
		}

		printf("read png with %u bytes\n", size);

		// Save the image when in movie mode
		if (movie_mode) {
			char frame_filename[PATH_MAX];
			int framefd;

			snprintf(frame_filename, PATH_MAX, "%s/frame-%d.png", frames_basepath, (int) time(NULL));

			framefd = creat(frame_filename, S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH);
			if (framefd < 0) {
				perror(frame_filename);
				exit(1);
			}

			written = 0;
			while (written != size) {
				rc = write(framefd, png+written, size-written);
				if (rc < 0)
					exit(1);
				written += rc;
			}

			close(framefd);
			printf("Frame written to file %s\n", frame_filename);
		}

		urp.buf = png;
		urp.pos = 0;
		urp.buf_size = size;

		pp = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL,
					    NULL, NULL);
		linfo = png_create_info_struct(pp);
		png_set_read_fn(pp, &urp, &user_read_data);
		png_read_info(pp, linfo);

		printf("Image %ux%u\n", (int)linfo->width, (int)linfo->height);

		if (linfo->bit_depth < 8) {
			png_set_packing(pp);
			png_set_expand(pp);

			if (linfo->valid & PNG_INFO_sBIT)
				png_set_shift(pp, &(linfo->sig_bit));
		} else if (linfo->bit_depth == 16)
			png_set_strip_16(pp);

		switch (linfo->color_type) {
		case PNG_COLOR_TYPE_RGB:
			bpp = 3;
			break;
		case PNG_COLOR_TYPE_RGB_ALPHA:
			bpp = 4;
			break;
		case PNG_COLOR_TYPE_GRAY:
			bpp = 1;
			break;
		case PNG_COLOR_TYPE_GRAY_ALPHA:
			bpp = 2;
			break;
		case PNG_COLOR_TYPE_PALETTE:
			bpp = linfo->num_trans ? 4:3;
			break;
		}

		printf("allocating pixel\n");
		pixel = malloc(linfo->height * linfo->width * bpp);
		if (!pixel) {
			fprintf(stderr, "ENOMEM\n");
			break;
		}

		num_passes = 1;

		for (pass = 0; pass < num_passes; pass++) {
			int scanline;
			for (scanline = 0; scanline < linfo->height; scanline++) {
				if (linfo->color_type == PNG_COLOR_TYPE_PALETTE)
					png_set_expand(pp);

				png_read_row(pp, pixel+(scanline*linfo->width*bpp), NULL);
			}
		}

		//png_read_end(pp, linfo);

		ximg = create_image_from_buffer(xdisp, DefaultScreen(xdisp), pixel, linfo->width, linfo->height, bpp);

		XPutImage(xdisp, xwin, xgc, ximg,
			  0, 0, 0, 0, win_width, win_height);

		XFlush(xdisp);

		sleep(1);

		XDestroyImage(ximg);
		png_destroy_read_struct(&pp, &linfo, NULL);

		//free(pixel);
		free(png);
	}

	exit(0);
}