#include <iostream>

/* Xlib include files */
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <X11/extensions/XShm.h>

#include <time.h> // clock_gettime

#include <ffmpeg/avcodec.h>

using namespace std;

// Frames per second
#define fps	25

// Delay for a frame
#define frame_delay (1000000 / fps)

// Taken from xdtv
#define Y(r,g,b) (((16829*r+31913*g+6416*b)>>16)+16)
#define U(r,g,b) (((-9750*r-19148*g+28898*b)>>16)+128)
#define V(r,g,b) (((28898*r-24199*g-4699*b)>>16)+128)

#define RGBDO {b=(*s++); g=*(s++); r=*(s++); s++;}

// Converts a RGB32 img to YUV4:2:0
static void rgb32_to_yuv420(unsigned char *d, unsigned char *s, int w, int h) {
  int aa, bb;
  unsigned char *y, *u, *v, r, g, b;
  y = d;
  u = y + w * h;
  v = u + w * h / 4;
  for (aa = h; aa > 0; aa -= 2) {
    for (bb = w; bb > 0; bb -= 2) {
      RGBDO; *(y++)=Y(r,g,b); *(u++)=U(r,g,b);
      RGBDO; *(y++)=Y(r,g,b);
    }
    for (bb = w; bb > 0; bb -= 2) {
      RGBDO; *(y++)=Y(r,g,b);
      RGBDO; *(y++)=Y(r,g,b); *(v++)=V(r,g,b);
    }
  }
}

// Gets the time im microsseconds
unsigned long long usecs(void) {
	struct timespec x;
	clock_gettime(CLOCK_REALTIME, &x);
	return (x.tv_nsec + x.tv_sec * 1000000000) / 1000;
}

int XShmAlloc(Display * const display, XShmSegmentInfo & shminfo, const unsigned int size) {
	shminfo.readOnly = False;

	// MIT-SHM stuff
	shminfo.shmid = shmget(IPC_PRIVATE, size, IPC_CREAT|0777);
	if (shminfo.shmid == -1)
		return -1;

	// MIT-SHM stuff
	shminfo.shmaddr = (char *)shmat (shminfo.shmid, 0, 0);

	// MIT-SHM stuff
	XShmAttach(display, &shminfo);

	// MIT-SHM stuff
	XSync(display, False);
	
	return 0;
}

void XShmFree(Display * const display, XShmSegmentInfo & shminfo) {
	XShmDetach (display, &shminfo);
	shmdt (shminfo.shmaddr);
	shmctl (shminfo.shmid, IPC_RMID, 0);
}

// Approach 1
void plan_a() {
}

// Approach 2
void plan_b() {
}

// Approach 3
void plan_c() {
}

int main(int argc, char *argv[]) {
	Display *display;

	// Conecta ao servidor X
	if ( (display = XOpenDisplay(NULL) ) == NULL) {
		cerr << argv[0] << ": não pude conectar ao servidor X "
		  << XDisplayName(NULL) << endl;
		exit(-1);
	}

	// Consulta se o servidor X tem a extensão MIT-SHM
	if (XShmQueryExtension(display) == 0) {
		cerr << argv[0] << ": Este servidor não tem a extensão MIT-SHM" << endl;
		exit(-1);
	}

	// Pega a janela raiz (root)
	Window root_window = XDefaultRootWindow(display);

	// Get the attributes of the root window
	XWindowAttributes root_attr;
	if (!XGetWindowAttributes(display, root_window, &root_attr ) ) {
		cerr << argv[0] << ": XGetWindowAttributes falhou" << endl;
		exit(-1);
	}

	int x, y;
	unsigned int width, height;
	unsigned int border_width;
	unsigned int depth;

	Window w;

	// Pega a informação sobre a geometria da janela raiz (root)
	if (XGetGeometry(display, root_window, &w,
	  &x, &y, &width, &height, &border_width, &depth) == 0) {
		cerr << argv[0] << ": XGetGeometry falhou" << endl;
		exit(-1);
	}

	// cap width and height
	width = 480;
	height = 360;

	// Procedimentos para obter a imagem pela extensão MIT-SHM
	XShmSegmentInfo shminfo;
#ifndef USE_PIXMAPS
	XImage * image = XShmCreateImage(display, CopyFromParent, depth, ZPixmap, NULL, &shminfo, width, height);
	if (image == NULL) {
		cerr << argv[0] << ": XShmCreateImage falhou" << endl;
		exit(-1);
	}

	if (XShmAlloc(display, shminfo, image->bytes_per_line * image->height) != 0) {
		cerr << argv[0] << ": XShmAlloc falhou" << endl;
		exit(-1);
	}
	image->data = shminfo.shmaddr;
#else
	if (XShmAlloc(display, shminfo, width * height * (depth / 8) ) != 0) {
		cerr << argv[0] << ": XShmAlloc falhou" << endl;
		exit(-1);
	}

	Pixmap image = XShmCreatePixmap(display, root_window, shminfo.shmaddr, &shminfo, width, height, depth);
	if (image == 0) {
		cerr << argv[0] << ": XShmCreatePixmap falhou" << endl;
		exit(-1);
	}

	// Create a graphics context
	GC gc = XCreateGC(display, root_window, 0, 0);
#endif

	// MIT-SHM stuff
	if (XShmAttach(display, &shminfo) == 0) {
		cerr << argv[0] << ": XShmAttach falhou" << endl;
		exit(-1);
	}

	AVCodec * codec;
	AVCodecContext * c = NULL;
	AVFrame *picture;

    // must be called before using avcodec lib
    avcodec_init();

    /* register all the codecs (you can also register only the codec
       you wish to have smaller code */
    avcodec_register_all();

	/* find the mpeg1 video encoder */
	codec = avcodec_find_encoder(CODEC_ID_MPEG2VIDEO);
	if (!codec) {
		cerr << argv[0] << "codec not found" << endl;
		exit(-1);
	}

    c = avcodec_alloc_context();
    picture = avcodec_alloc_frame();

	/* put sample parameters */
	c->bit_rate = 400000;
	/* resolution must be a multiple of two */
	c->width = width;
	c->height = height;
	/* frames per second */
	c->time_base= (AVRational){1, fps};
	c->gop_size = 10; /* emit one intra frame every ten frames */
	c->max_b_frames = 1;
	c->pix_fmt = PIX_FMT_YUV420P;

    /* open it */
	if (avcodec_open(c, codec) < 0) {
		cerr << argv[0] << "could not open codec" << endl;
		exit(-1);
	}

    /* the codec gives us the frame size, in samples */

	const char filename[] = "test.mpg";
    FILE * f = NULL;
    f = fopen(filename, "wb");
    if (!f) {
        cerr << argv[0] << "could not open " << filename << endl;
        exit(-1);
    }

	int i, out_size, outbuf_size, size;
	uint8_t * outbuf, * picture_buf;

	/* alloc image and output buffer */
	outbuf_size = 100000;
	outbuf = (uint8_t *)malloc(outbuf_size);
	size = c->width * c->height;
	picture_buf = (uint8_t *)malloc((size * 3) / 2); 	// Tamanho para YUV 4:2:0

	picture->data[0] = picture_buf;						// Seta o ponteiro Y
	picture->data[1] = picture->data[0] + size;			// Seta o ponteiro U
	picture->data[2] = picture->data[1] + size / 4;		// Seta o ponteiro V
	picture->linesize[0] = c->width;					// Seta o tamanho da linha Y
	picture->linesize[1] = c->width / 2;				// Seta o tamanho da linha U
	picture->linesize[2] = c->width / 2;				// Seta o tamanho da linha V

	cout << "Codificando vídeo..." << endl;

	unsigned int total_capture = 0, total_render = 0;
	unsigned long long time = usecs();
	unsigned long long time1 = time;
	unsigned long long start = time;

	/* encode 1 second of video */
	for(i = 0; i < 250; i++) {
#ifndef USE_PIXMAPS
		if (XShmGetImage(display, root_window, image, x, y, AllPlanes) == 0) {
			cerr << argv[0] << ": XShmGetImage falhou" << endl;
			exit(-1);
		}
#else
		if (XCopyArea(display, root_window, image, gc, 0, 0, width, height, 0, 0) == 0) {
			cerr << argv[0] << ": XCopyArea falhou" << endl;
			exit(-1);
		}

#endif

		time1 = usecs();
		total_capture += time1 - time;

		rgb32_to_yuv420((unsigned char *) picture->data[0], (unsigned char *)shminfo.shmaddr, width, height);

		/* encode the image */
		out_size = avcodec_encode_video(c, outbuf, outbuf_size, picture);
		fwrite(outbuf, 1, out_size, f);

		unsigned long long y;
		unsigned int x;
		if ( (x = (y = usecs() ) - time) < frame_delay)
			usleep(frame_delay - x);
		total_render += y - time1;
		time = usecs();
    }

    double run_time = (usecs() - start) / 1e6;

	cout << "OK" << endl;
	cout << "Levou " << run_time << " segundos" << endl;
	cout << "Média " << i / run_time << " operações por segundo" << endl;
	cout << "Tempo médio de captura: " << ( total_capture / i) / 1000.0 << "ms" << endl;
	cout << "Tempo médio de codificação: " << ( total_render / i) / 1000.0 << "ms" << endl;
	cout << "Tempo médio do processo: " << ( (total_capture + total_render) / i) / 1000.0 << "ms" << endl;

    /* get the delayed frames */
    for(; out_size; i++) {
        out_size = avcodec_encode_video(c, outbuf, outbuf_size, NULL);
        fwrite(outbuf, 1, out_size, f);
    }

    /* add sequence end code to have a real mpeg file */
    outbuf[0] = 0x00;
    outbuf[1] = 0x00;
    outbuf[2] = 0x01;
    outbuf[3] = 0xb7;
    fwrite(outbuf, 1, 4, f);
    fclose(f);
    free(picture_buf);
    free(outbuf);

	// Free video structures
    avcodec_close(c);
    av_free(c);
    av_free(picture);

#ifndef USE_PIXMAPS
	XDestroyImage (image);
#else
	XFreePixmap (display, image);
#endif
	// Destroy shared image
	XShmFree(display,shminfo);
	
	XCloseDisplay(display);
}