/*
 * Compile with:
 * gcc `libpng-config --cflags` -o pngvnc pngvnc.c `libpng-config --libs`
 *
 * Under public domain. Original by Nishi <nishi@nishi.boats>
 */

#include <stdio.h>
#include <stdlib.h>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/tcp.h>
#include <netinet/in.h>
#include <signal.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>

#include <png.h>

#define NAME "PNG-VNC (http://nishi.boats/pngvnc)"

png_structp png;
png_infop info;

int width;
int height;

png_bytep* rows;
uint8_t** frows;

#define strsend(s,str) send(s, str, strlen(str), 0)

void send32(int sock, uint32_t num){
	unsigned char c;
	int i;
	for(i = 0; i < 4; i++){
		c = (num >> 24) & 0xff;
		send(sock, &c, 1, 0);
		num = num << 8;
	}
}

void send16(int sock, uint16_t num){
	unsigned char c;
	int i;
	for(i = 0; i < 2; i++){
		c = (num >> 8) & 0xff;
		send(sock, &c, 1, 0);
		num = num << 8;
	}
}

void send8(int sock, uint8_t num){
	unsigned char c;
	int i;
	for(i = 0; i < 1; i++){
		c = (num >> 0) & 0xff;
		send(sock, &c, 1, 0);
		num = num << 8;
	}
}

int recv16(int sock, uint16_t* out){
	unsigned char c;
	int i;
	*out = 0;
	for(i = 0; i < 2; i++){
		if(recv(sock, &c, 1, 0) <= 0) return -1;
		*out <<= 8;
		*out |= c;
	}
	return 2;
}

int recv32(int sock, uint32_t* out){
	unsigned char c;
	int i;
	*out = 0;
	for(i = 0; i < 4; i++){
		if(recv(sock, &c, 1, 0) <= 0) return -1;
		*out <<= 8;
		*out |= c;
	}
	return 4;
}

void sendarea(int sock, int x, int y, int w, int h){
	int i, j;
	for(i = y; i < y + h; i++){
		send8(sock, 0);
		send8(sock, 0);
		send16(sock, 1);

		send16(sock, x);
		send16(sock, i);
		send16(sock, w);
		send16(sock, 1);
		send32(sock, 0);
		send(sock, frows[i] + x * 4, w * 4, 0);
	}
}

void vnc_handle(int sock){
	unsigned char c;
	strsend(sock, "RFB 003.003\n");
	if(recv(sock, &c, 1, 0) <= 0) return;
	send32(sock, 1);

	send16(sock, width);
	send16(sock, height);
	send8(sock, 32); /* bpp */
	send8(sock, 8); /* depth */
	send8(sock, 0); /* big endian */
	send8(sock, 0); /* true color */
	send16(sock, 8); /* red-max */
	send16(sock, 8); /* green-max */
	send16(sock, 8); /* blue-max */
	send8(sock, 0); /* red-shift */
	send8(sock, 8); /* green-shift */
	send8(sock, 16); /* blue-shift */

	send8(sock, 0);
	send8(sock, 0);
	send8(sock, 0);

	send32(sock, strlen(NAME));
	strsend(sock, NAME);

	/* start framebuffer crap */
	sendarea(sock, 0, 0, width, height);

	int cnt = 0;

	while(1){
		if(recv(sock, &c, 1, 0) <= 0) goto brk;
		int error = 0;
		socklen_t len = sizeof(error);
		int retval = getsockopt(sock, SOL_SOCKET, SO_ERROR, &error, &len);
		if(error != 0 || retval != 0) break;
		if(c == 0){
			unsigned char buffer[3 + 16];
			recv(sock, buffer, sizeof(buffer), 0);
		}else if(c == 1){
			if(recv(sock, &c, 1, 0) <= 0) goto brk;
			uint16_t out;
			if(recv16(sock, &out) <= 0) goto brk;
			if(recv16(sock, &out) <= 0) goto brk;
			int i;
			unsigned char buffer[6];
			for(i = 0; i < out; i++){
				if(recv(sock, buffer, 6, 0) <= 0) goto brk;
			}
		}else if(c == 2){
			if(recv(sock, &c, 1, 0) <= 0) goto brk;
			uint16_t out;
			if(recv16(sock, &out) <= 0) goto brk;
			int i;
			unsigned char buffer[4];
			for(i = 0; i < out; i++){
				if(recv(sock, buffer, 4, 0) <= 0) goto brk;
			}
		}else if(c == 3){
			/* update fb */
			if(recv(sock, &c, 1, 0) <= 0) goto brk;
			uint16_t x, y, w, h;
			if(recv16(sock, &x) <= 0) goto brk;
			if(recv16(sock, &y) <= 0) goto brk;
			if(recv16(sock, &w) <= 0) goto brk;
			if(recv16(sock, &h) <= 0) goto brk;
		}else if(c == 4){
			unsigned char buffer[1 + 2];
			if(recv(sock, buffer, 1 + 2, 0) <= 0) goto brk;
			uint32_t num;
			if(recv32(sock, &num) <= 0) goto brk;
			if(num == 0xff1b || num == 'q' || num == 'Q') goto brk;
		}else if(c == 5){
			unsigned char buffer[1 + 2 + 2];
			if(recv(sock, buffer, 1 + 2 + 2, 0) <= 0) goto brk;
		}
	}
brk:
	return;
}

int main(int argc, char** argv){
	if(argc < 3){
		printf("Usage: %s port png\n", argv[0]);
		return 1;
	}
	FILE* fp = fopen(argv[2], "rb");
	png = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	info = png_create_info_struct(png);

	if(setjmp(png_jmpbuf(png))){
		png_destroy_read_struct(&png, &info, NULL);
		fclose(fp);
		return 1;
	}
	png_init_io(png, fp);
	png_read_info(png, info);

	width = png_get_image_width(png, info);
	height = png_get_image_height(png, info);
	int depth = png_get_bit_depth(png, info);
	int type = png_get_color_type(png, info);
	printf("PNG %dx%d, depth %d, type %d\n", width, height, depth, type);

	if(depth == 16) png_set_strip_16(png);
	if(type == PNG_COLOR_TYPE_PALETTE) png_set_palette_to_rgb(png);
	if(type == PNG_COLOR_TYPE_GRAY && depth < 8) png_set_expand_gray_1_2_4_to_8(png);
	if(png_get_valid(png, info, PNG_INFO_tRNS)) png_set_tRNS_to_alpha(png);
	if(type == PNG_COLOR_TYPE_RGB || type == PNG_COLOR_TYPE_GRAY || type == PNG_COLOR_TYPE_PALETTE) png_set_filler(png, 0xFF, PNG_FILLER_AFTER);
	if(type == PNG_COLOR_TYPE_GRAY || type == PNG_COLOR_TYPE_GRAY_ALPHA) png_set_gray_to_rgb(png);
	png_read_update_info(png, info);

	rows = (png_bytep*)malloc(sizeof(*rows) * (height));
	frows = (uint8_t**)malloc(sizeof(*frows) * (height) * 4);
	int i, j;

	for(i = 0; i < height; i++){
		rows[i] = (png_byte*)malloc(png_get_rowbytes(png, info));
		frows[i] = malloc(4 * width);
	}

	png_read_image(png, rows);

	png_destroy_read_struct(&png, &info, NULL);
	fclose(fp);

	for(i = 0; i < height; i++){
		for(j = 0; j < width; j++){
			png_bytep byte = &(rows[i][j * 4]);
			uint32_t color = (byte[0] << 0) | (byte[1] << 8) | (byte[2] << 16);
			frows[i][j * 4 + 0] = (color >> 16) & 0xff;
			frows[i][j * 4 + 1] = (color >> 8) & 0xff;
			frows[i][j * 4 + 2] = (color >> 0) & 0xff;
			frows[i][j * 4 + 3] = 0;
		}
	}

	int sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	int yes = 1;
	if(setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (void*)&yes, sizeof(yes)) < 0){
		close(sock);
		return 1;
	}
	if(setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (void*)&yes, sizeof(yes)) < 0){
		close(sock);
		return 1;
	}
	struct sockaddr_in saddr;
	memset(&saddr, 0, sizeof(saddr));
	saddr.sin_family = AF_INET;
	saddr.sin_addr.s_addr = INADDR_ANY;
	saddr.sin_port = htons(atoi(argv[1]));
	if(bind(sock, (struct sockaddr*)&saddr, sizeof(saddr)) < 0){
		close(sock);
		return 1;
	}
	if(listen(sock, 128) < 0){
		close(sock);
		return 1;
	}
	signal(SIGPIPE, SIG_IGN);
	signal(SIGCHLD, SIG_IGN);
	printf("am ready\n");
	while(1){
		struct sockaddr_in caddr;
		int clen = sizeof(caddr);
		int csock = accept(sock, (struct sockaddr*)&caddr, &clen);
		if(csock < 0){
			break;
		}else{
			char address[513];
			struct sockaddr* sa = (struct sockaddr*)&caddr;
			getnameinfo(sa, sizeof(caddr), address, 512, NULL, 0, NI_NUMERICHOST);
			printf("got a connection from %s, am forking\n", address);
			if(fork() == 0){
				vnc_handle(csock);
				close(csock);
				printf("client died\n");
				_exit(0);
			}else{
				close(csock);
			}
		}
	}
	return 0;
}