//#bx pkgs := libjpeg x11 xext

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/shm.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/extensions/XShm.h>
#include <jpeglib.h>
#include <jerror.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <string.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>

static void *imgbuf = NULL;
static size_t imgbuflen = 0;
static size_t imgbufsize = 0;

static void mem_jpeg_init_destination(j_compress_ptr cinfo) {
	if (imgbufsize == 0) {
		imgbufsize = 4 * 1024;
		imgbuf = malloc(imgbufsize);
	}

	imgbuflen = 0;
	cinfo->dest->next_output_byte = imgbuf;
	cinfo->dest->free_in_buffer = imgbufsize;
}

static boolean mem_jpeg_empty_output_buffer(j_compress_ptr cinfo) {
	imgbuflen = imgbufsize;
	imgbufsize *= 2;
	imgbuf = realloc(imgbuf, imgbufsize);

	cinfo->dest->next_output_byte = imgbuf + imgbuflen;
	cinfo->dest->free_in_buffer = imgbufsize - imgbuflen;
	printf("grew buffer to %zu\n", imgbufsize);
	return TRUE;
}

static void mem_jpeg_term_destination(j_compress_ptr cinfo) {
	imgbuflen = imgbufsize - cinfo->dest->free_in_buffer;
}

static void write_jpeg(XImage *img) {
	struct jpeg_compress_struct cinfo = {0};
	jpeg_create_compress(&cinfo);
	struct jpeg_destination_mgr dest = {0};
	dest.init_destination = mem_jpeg_init_destination;
	dest.empty_output_buffer = mem_jpeg_empty_output_buffer;
	dest.term_destination = mem_jpeg_term_destination;
	cinfo.dest = (struct jpeg_destination_mgr *)&dest;

	struct jpeg_error_mgr jerr;
	cinfo.err = jpeg_std_error(&jerr);

	cinfo.image_width = img->width;
	cinfo.image_height = img->height;
	cinfo.input_components = 4;
	cinfo.in_color_space = JCS_EXT_BGRX;
	jpeg_set_defaults(&cinfo);
	jpeg_start_compress(&cinfo, TRUE);

	for (int i = 0; i < cinfo.image_height; ++i) {
		JSAMPROW row = (unsigned char *)img->data + i * img->bytes_per_line;
		jpeg_write_scanlines(&cinfo, &row, 1);
	}

	jpeg_finish_compress(&cinfo);
	jpeg_destroy_compress(&cinfo);
}

static void upload_jpeg(XImage *img, int fd) {
	write_jpeg(img);
	uint32_t len = htonl((uint32_t)imgbuflen);
	write(fd, &len, sizeof(len));
	write(fd, imgbuf, imgbuflen);
}

int main(int argc, char **argv) {
	if (argc <= 1) {
		printf("Usage: %s <host> [port] [sleep time]\n", argv[0]);
		return 1;
	}

	char *host = argv[1];
	char *port;
	int sleeptime;
	if (argc >= 3) {
		port = argv[2];
	} else {
		port = "8099";
	}

	if (argc >= 4) {
		sleeptime = atoi(argv[3]);
	} else {
		sleeptime = 200;
	}

	printf("Connect %s:%s...\n", host, port);
	struct addrinfo hints = {
		.ai_family = AF_UNSPEC,
		.ai_socktype = SOCK_STREAM,
		.ai_protocol = 0,
		.ai_flags = AI_ADDRCONFIG,
	};
	struct addrinfo *info;
	if (getaddrinfo(host, port, &hints, &info) < 0) {
		perror("getaddrinfo");
		return 1;
	}

	int sockfd = socket(info->ai_family, info->ai_socktype, info->ai_protocol);
	if (sockfd < 0) {
		perror("socket");
		return 1;
	}

	if (connect(sockfd, info->ai_addr, info->ai_addrlen) < 0) {
		perror("connect");
		return 1;
	}

	freeaddrinfo(info);

	// This is useful
	if (getenv("DISPLAY") == NULL) {
		setenv("DISPLAY", ":0", 1);
	}

	Display *display = XOpenDisplay(NULL);
	if (display == NULL) {
		fprintf(stderr, "Failed to open display.\n");
		return 1;
	}

	Window root = XDefaultRootWindow(display);
	XWindowAttributes gwa;
	XGetWindowAttributes(display, root, &gwa);

	// Create shm image
	XShmSegmentInfo shminfo;
	XImage *image = XShmCreateImage(
			display, DefaultVisual(display, DefaultScreen(display)),
			32, ZPixmap, NULL, &shminfo, gwa.width, gwa.height);

	// Attach shm image
	shminfo.shmid = shmget(IPC_PRIVATE,
			image->bytes_per_line * image->height,
			IPC_CREAT|0777);
	shminfo.shmaddr = image->data = shmat(shminfo.shmid, 0, 0);
	shminfo.readOnly = False;
	XShmAttach(display, &shminfo);

	while (1) {
		if (!XShmGetImage(display, root, image, 0, 0, AllPlanes)) {
			fprintf(stderr, "Failed to get image.\n");
			exit(EXIT_FAILURE);
		}

		upload_jpeg(image, sockfd);
		usleep(sleeptime * 1000);
	}

	XShmDetach(display, &shminfo);
	XDestroyImage(image);
}