/*
   This program captures images from a pwc video device and display the images
   on the screen using Xlib and XShm (MIT X Shared memory extension). 
   Feel free to use it and improve it.

Author: Vincent Hourdin     (origin: Edgar A. Sotter S. (GTK+ && ov511 && !shm))
Date: April 2004

Compile command :
gcc -O3 -o webcam webcam-xlib-shm-pwc.c -L/usr/X11R6/lib -lX11 -I/usr/src/linux/drivers/usb/media/pwc -lXext

*/

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <linux/types.h> 
#include <linux/videodev.h>
#include <errno.h>
#include "pwc-ioctl.h"	/* in kernel source tree or pwc package */
#include <X11/Xlib.h>

#include <sys/ipc.h>
#include <sys/shm.h>
#include <X11/extensions/XShm.h>


#define COLS      640
#define ROWS      480
#define FRAMERATE 10

#define DUMP_BASENAME	"dump"
#define VERSION_STRING	"0.1.1"

/* ioctl ressources */
struct video_window win;
struct video_mbuf vidbuf;

/* images buffers */
unsigned char *bigbuf;
unsigned char *buf;

/* X ressources */
GC	gc;
Display *display;
Window	window, root;
unsigned long white_pixel, black_pixel;
XImage *ximage;
int nplanes;
XShmSegmentInfo	shm_info;

/* color conversion variables */
int Ubegin, Vbegin;
int VtoR[256];
int UtoG[256];
int VtoG[256];
int UtoB[256];


/* functions */
void yuv420p2rgbpixbuf();
void yuv420p2graypixbuf();
void init_conversion();
void dump_frame();



int main (int argc, char **argv){
	XSetWindowAttributes attr;
	XEvent ev;
	Visual *the_visual;
	struct video_capability cap;
	struct video_picture pic;
	struct video_mmap mapbuf;
	char *image_data, *opt_list = "f:ghv", *device;
	int screen, i, frame, cam, run, greyscale = 0, option, framerate = FRAMERATE;
	long event_mask = KeyPressMask /* | ButtonPressMask*/;

	while ((option = getopt (argc, argv, opt_list)) != -1){
		switch (option){
			case 'f':
				framerate = atoi(optarg);
				break;
			case 'g':
				greyscale = 1;
				break;
			case 'h':
				fprintf(stderr, "PWC webcam viewer.\n"
						"Usage: %s [-g] [-f framerate] [-h] [-v] [device]\n"
						"	-g : display pictures in greyscale\n"
						"	-f : set framerate to argument\n"
						"	-h : see this help\n"
						"	-v : version\n", *argv);
				exit(0);
			case 'v':
				fprintf(stderr, "vinvin's PWC viewer for linux version %s\n", VERSION_STRING);
				exit(0);
			case '?':
				fprintf(stderr, "Unknown argument : %c\n", optopt);
				fprintf(stderr, "Usage: %s [-g] [-f framerate] [-h] [-v] [device]\n", *argv);
				exit(1);
		}
	}

	/* opening device */
	if (optind != argc)
		device = argv[optind];
	else device = "/dev/video";
	
	cam = open(device, O_RDWR);

	if (cam < 0){
		fprintf(stderr, "ERROR: can't open %s. Make a link or connect the cam.\n", device);
		exit(-1);
	}

	//--------------------------------------------------------------------
	/* VIDEO CAPABILITIES*/
	//--------------------------------------------------------------------

	if (ioctl(cam, VIDIOCGCAP, &cap) == -1){
		perror("ioctl VIDIOCGCAP");
		exit(-1);
	}

	printf("------------------------------------\n");
	printf("------------------------------------\n");
	printf("name -> %s\n", cap.name);
	printf("type -> %d\n", cap.type);
	printf("channels -> %d\n", cap.channels);
	printf("audios -> %d\n", cap.audios);
	printf("maxwidth -> %d\n", cap.maxwidth);
	printf("maxheight -> %d\n", cap.maxheight);
	printf("minwidth -> %d\n", cap.minwidth);
	printf("minheight -> %d\n", cap.minheight);
	printf("------------------------------------\n");


	//-----------------------------------------------------------------------
	/* CAPTURE WINDOW */
	//-----------------------------------------------------------------------

	win.x = 0;
	win.y = 0;
	win.width = COLS;
	win.height = ROWS;
	win.clipcount = 0;
	win.chromakey = 1;
	win.flags = 0;
	win.flags &= ~PWC_FPS_FRMASK;
	win.flags |= (framerate << PWC_FPS_SHIFT);

	if (ioctl(cam, VIDIOCSWIN, &win) == -1){
		perror ("ioctl VIDIOCSWIN");
		exit (-1); 
	}

	ioctl( cam, VIDIOCGWIN, &win );

	printf("------------------------------------\n");
	printf("width -> %d\n", win.width);
	printf("height -> %d\n", win.height);
	if (win.flags & PWC_FPS_FRMASK)
		printf("framerate -> %d\n", (win.flags & PWC_FPS_FRMASK) >> PWC_FPS_SHIFT);
	else printf("no framerate !\n");


	//---------------------------------------------------------------------
	/* IMAGE PROPERTIES*/
	//---------------------------------------------------------------------
#if 0
	pic.palette = 0;
	pic.brightness = 50000;
	pic.whiteness = 42000;
	pic.contrast = 50000;
	pic.colour = 60000;

	if (ioctl( cam, VIDIOCSPICT, &pic ) == -1){
		perror("ioctl VIDIOCSPICT");
		exit(-1);
	}
#endif
	ioctl( cam, VIDIOCGPICT, &pic );

	printf("------------------------------------\n");
	printf("brightness -> %d \n", pic.brightness/256 );
	printf("colour -> %d\n", pic.colour/256 );
	printf("contrast -> %d \n", pic.contrast/256 );
	printf("whiteness (gamma) -> %d\n", pic.whiteness/256 );
	printf("depth -> %d\n", pic.depth );
	printf("palette -> %d\n", pic.palette );
	printf("------------------------------------\n");


	//-----------------------------------------------------------------------
	/* MAPPING BUFFER */
	//-----------------------------------------------------------------------

	if (ioctl(cam, VIDIOCGMBUF, &vidbuf) == -1){
		perror("ioctl VIDIOCGMBUF"); 
		exit(-1);
	}

	printf("------------------------------------\n");
	printf("size  -> %d\n", vidbuf.size);
	printf("frames -> %d\n", vidbuf.frames);
	printf("------------------------------------\n");

	bigbuf = (unsigned char *) mmap( 0, vidbuf.size, PROT_READ | PROT_WRITE, MAP_SHARED, cam, 0 );


	//------------------------------------------------------------------------------
	/* CREATE DISPLAY WINDOWS AND IMAGE AREA */
	//------------------------------------------------------------------------------

	if ((display = XOpenDisplay ("")) == NULL) {
		fprintf (stderr, "Can't open Display\n");
		exit (1);
	}
	
	if (XShmQueryExtension (display) == False){
		fprintf (stderr, "no shared memory support by X server\n");
		exit(1);
	}

	screen = DefaultScreen (display);
	gc = DefaultGC (display, screen);
	root = RootWindow (display, screen);
	white_pixel = WhitePixel (display, screen);
	black_pixel = BlackPixel (display, screen);
	nplanes = DisplayPlanes (display, screen);
	the_visual = DefaultVisual (display, screen);
	image_data = calloc (1, sizeof(unsigned int) * win.width * win.height);
	
	if (!image_data){
		perror("calloc");
		exit(1);
	}
	
	/* XShm code - create image, create shared memory */
	
	ximage = XShmCreateImage(display, the_visual, nplanes, ZPixmap,
			image_data, &shm_info, win.width, win.height);
	ximage->byte_order = LSBFirst;

	if ((shm_info.shmid = shmget(   IPC_PRIVATE,
					ximage->bytes_per_line * ximage->height,
					IPC_CREAT|0777 )) == -1){
		perror("shmget");
		exit(2);
	}
	
	shm_info.shmaddr = ximage->data = shmat (shm_info.shmid, 0, 0);
	shm_info.readOnly = False;
	
	if (XShmAttach (display, &shm_info) == False){
		perror ("attaching shared mem to the X server");
		exit(2);
	}
	XSync(display, 0);

	if (!XInitImage(ximage)){
		perror("XInitImage");
		exit(2);
	}

	
	attr.background_pixel = white_pixel;
	attr.border_pixel = black_pixel;

	window = XCreateWindow (display, root, 0, 0, win.width, win.height, 2, nplanes,
			CopyFromParent, the_visual, CWBackPixel | CWBorderPixel, &attr);

	XSelectInput (display, window, event_mask);



	//----------------------------------------------------------------------
	/* SET BUFFERS */
	//----------------------------------------------------------------------		 

	mapbuf.width = win.width;
	mapbuf.height = win.height;
	mapbuf.format = pic.palette;

	for (frame = 0; frame < vidbuf.frames; frame++){		// turn on both of the buffers
		mapbuf.frame = frame;					// to start capture process.
		if (ioctl(cam, VIDIOCMCAPTURE, &mapbuf) < 0){		// Now they can store images.
			perror("VIDIOCMCAPTURE");
			exit(-1);
		}
	}

	frame = 0;


	//---------------------------------------------------------------------
	/* CAPTURING */
	//---------------------------------------------------------------------		

	XStoreName (display, window, "cam view");
	XMapWindow (display, window);

	init_conversion();
	Ubegin = win.width * win.height + 1;
	Vbegin = Ubegin + (win.width * win.height) / 4;
	run = 1;

	while(run){
	/* keycodes : - 15,82   + 21,86    g 42   q 38    d 40     c 54    */  
		if (XCheckMaskEvent(display, event_mask, &ev)){
			switch (ev.type) {
				case KeyPress :
					switch (ev.xkey.keycode){
						case 40:		// d
							dump_frame();
							break;
						case 42:		// g
							greyscale ^= 1;
							break;
						case 38:		// q
							run = 0;
							break;
						default:
							break;
					}
					break;
					/*
				case ButtonPress :
					break;
					*/
				default :
					printf("not implemented event\n");
			}
		}

		i = -1;
		while(i<0){

			i = ioctl(cam, VIDIOCSYNC, &frame);	// Wait until the actual buffer
			if (i < 0 && errno == EINTR) continue;	// is full. When it happends 
			if (i < 0) {				// it start to capture to	
				perror ("VIDIOCSYNC");		// the other buffer.
				exit(-1);
			}
			break;
		}


		buf = bigbuf + vidbuf.offsets[frame];
		mapbuf.frame = frame;

		if (greyscale)
			yuv420p2graypixbuf();			// Convert colors and put into shm
		else yuv420p2rgbpixbuf();
		
		XSync(display, 0);
		XShmPutImage (display, window, gc, ximage, 0, 0, 0, 0, win.width, win.height, False);
		

		if (ioctl(cam, VIDIOCMCAPTURE, &mapbuf) < 0) {
			perror("VIDIOCMCAPTURE");		// Turn on the buffer that
			exit(-1);                               // was being used.
		}

		frame++;

		if (frame >= vidbuf.frames) frame=0;
	}
	

	XUnmapWindow (display, window);
	close(cam);
	XShmDetach (display, &shm_info);
	XDestroyWindow(display, window);
	shmdt (shm_info.shmaddr);
	shmctl (shm_info.shmid, IPC_RMID, 0);
	return 0;
}


//------------------------------------------------------------------------
/* CONVERT COLORS */
//------------------------------------------------------------------------

/* converts yuv420p buffer to rgb shared memory buffer */
/* shared memory is at shm_info.shmaddr */
void yuv420p2rgbpixbuf(){
	int x, y;
	int position, uvpos, olduv;
	int r,g,b;
	unsigned char Y, U, V;
	unsigned char *ubuf, *vbuf;
	register unsigned char *pixbuf;

	position = 0;
	olduv = 0;
	uvpos = 0;
	ubuf = buf+Ubegin;
	vbuf = buf+Vbegin;
	pixbuf = (unsigned char *) shm_info.shmaddr;
	
	for (y = 0; y < win.height; y++){
		for (x = 0; x < win.width; x++){

			Y = buf[position++];
			U = ubuf[uvpos];
			V = vbuf[uvpos];

			r = Y + VtoR[V];
			g = Y - UtoG[U] - VtoG[V];
			b = Y + UtoB[U];

			*pixbuf++ = (unsigned char) (b > 255) ? 255 : ((b < 0) ? 0 : b);
			*pixbuf++ = (unsigned char) (g > 255) ? 255 : ((g < 0) ? 0 : g);
			*pixbuf++ = (unsigned char) (r > 255) ? 255 : ((r < 0) ? 0 : r);
			pixbuf++;
			
			if (x&1)
				uvpos++;
		}
		if (y&1)
			olduv = uvpos;

		else 	uvpos = olduv;
	}
}

/* converts yuv420p buffer to grayscale rgb shared memory buffer */
void yuv420p2graypixbuf(){
	int x, y;
	int position;
	unsigned char Y;
	register unsigned char *pixbuf;

	position = 0;
	pixbuf = (unsigned char *)shm_info.shmaddr;

	for (y = 0; y < win.height; y++){
		for (x = 0; x < win.width; x++){
			Y = buf[position++];
			*pixbuf++ = Y;
			*pixbuf++ = Y;
			*pixbuf++ = Y;
			pixbuf++;
		}
	}
}



/* creates a ppm file with current frame snapshot
 * unique file is created each time with time in name */

void dump_frame(){
	time_t t;
	int size, filesize, fd;
	register int i;
	char filename[64];
	register unsigned int *ptr;
	unsigned char *dest, *dst;
	
	time(&t);
	sprintf(filename, "%s%d.ppm", DUMP_BASENAME, (int) t);

	if ((fd = open(filename, O_WRONLY|O_CREAT, S_IRUSR| S_IWUSR)) < 0){
		perror("open");
		return;
	}

	size = win.width * win.height;
	filesize = size * 3 + 15;
	dest = malloc (filesize);
	if (!dest){
		perror("unable to dump video frame : malloc");
		return;
	}
	
	sprintf((char*)dest, "P6 %d %d 255 ", win.width, win.height);
	dst = dest+15;
	
	ptr = (unsigned int *) shm_info.shmaddr;
	
	for (i=0; i<size; i++, ptr++){
		*dst++ = ((*ptr) & 0x00ff0000) >> 16;
		*dst++ = ((*ptr) & 0x0000ff00) >> 8;
		*dst++ = ((*ptr) & 0x000000ff);
	}
	
	write(fd, dest, filesize);
	printf("file %s created\n", filename);

	close(fd);
}


/* speed up further computations */
void init_conversion(){
	int i;
	
	for (i=0; i<=255; i++){
		VtoR[i] = 1.4016868*(i-128);
		UtoG[i] = 0.3436954*(i-128);
		VtoG[i] = 0.7141690*(i-128);
		UtoB[i] = 1.7721604*(i-128);
	}
}