// STL includes
#include <iostream>
#include <cstdint>

// X11Grabber includes
#include <grabber/X11Grabber.h>

X11Grabber::X11Grabber(bool useXGetImage, int cropLeft, int cropRight, int cropTop, int cropBottom, int horizontalPixelDecimation, int verticalPixelDecimation) :
    _imageResampler(),
    _useXGetImage(useXGetImage),
    _cropLeft(cropLeft),
    _cropRight(cropRight),
    _cropTop(cropTop),
    _cropBottom(cropBottom),
    _x11Display(nullptr),
    _pixmap(None),
    _srcFormat(nullptr),
    _dstFormat(nullptr),
    _srcPicture(None),
    _dstPicture(None),
    _screenWidth(0),
    _screenHeight(0),
    _croppedWidth(0),
    _croppedHeight(0),
    _image(0,0)
{
    _imageResampler.setHorizontalPixelDecimation(horizontalPixelDecimation);
    _imageResampler.setVerticalPixelDecimation(verticalPixelDecimation);
    _imageResampler.setCropping(0, 0, 0, 0); // cropping is performed by XShmGetImage or XGetImage
    memset(&_pictAttr, 0, sizeof(_pictAttr));
    _pictAttr.repeat = RepeatNone;
}

X11Grabber::~X11Grabber()
{
    if (_x11Display != nullptr)
    {
	freeResources();
        XCloseDisplay(_x11Display);
    }
}

void X11Grabber::freeResources()
{
    // Cleanup allocated resources of the X11 grab
    XDestroyImage(_xImage);
    if(_XShmAvailable && !_useXGetImage)   {
	XShmDetach(_x11Display, &_shminfo);
	shmdt(_shminfo.shmaddr);
	shmctl(_shminfo.shmid, IPC_RMID, 0);
    }
    if (_XRenderAvailable && !_useXGetImage) {
	XRenderFreePicture(_x11Display, _srcPicture);
	XRenderFreePicture(_x11Display, _dstPicture);
	XFreePixmap(_x11Display, _pixmap);
    }
}

void X11Grabber::setupResources()
{
    if(_XShmAvailable && !_useXGetImage) {
	_xImage = XShmCreateImage(_x11Display, _windowAttr.visual,
	_windowAttr.depth, ZPixmap, NULL, &_shminfo,
	_croppedWidth, _croppedHeight);

	_shminfo.shmid = shmget(IPC_PRIVATE, _xImage->bytes_per_line * _xImage->height, IPC_CREAT|0777);
	_xImage->data = (char*)shmat(_shminfo.shmid,0,0);
	_shminfo.shmaddr = _xImage->data;
	_shminfo.readOnly = False;
    
	XShmAttach(_x11Display, &_shminfo);
    }
    if (_XRenderAvailable && !_useXGetImage) {
	if(_XShmPixmapAvailable) {
	    _pixmap = XShmCreatePixmap(_x11Display, _window, _xImage->data, &_shminfo, _croppedWidth, _croppedHeight, _windowAttr.depth);
	} else {
	    _pixmap = XCreatePixmap(_x11Display, _window, _croppedWidth, _croppedHeight, _windowAttr.depth);
	}
	_srcFormat = XRenderFindVisualFormat(_x11Display, _windowAttr.visual);
	_dstFormat = XRenderFindVisualFormat(_x11Display, _windowAttr.visual);
	_srcPicture = XRenderCreatePicture(_x11Display, _window, _srcFormat, CPRepeat, &_pictAttr);
	_dstPicture = XRenderCreatePicture(_x11Display, _pixmap, _dstFormat, CPRepeat, &_pictAttr);
	 XRenderSetPictureFilter(_x11Display, _srcPicture, "bilinear", NULL, 0);
    }
}

bool X11Grabber::Setup()
{
    _x11Display = XOpenDisplay(NULL);
    if (_x11Display == nullptr)
    {
      std::cerr << "X11GRABBER ERROR: Unable to open display";
      if (getenv("DISPLAY"))
	std::cerr <<  " " << std::string(getenv("DISPLAY")) << std::endl;
      else
	std::cerr << ". DISPLAY environment variable not set" << std::endl;
      return false;
     }
     
    _window = DefaultRootWindow(_x11Display);

    int dummy, pixmaps_supported;
    
    _XRenderAvailable = XRenderQueryExtension(_x11Display, &dummy, &dummy);
    _XShmAvailable = XShmQueryExtension(_x11Display);
    XShmQueryVersion(_x11Display, &dummy, &dummy, &pixmaps_supported);
    _XShmPixmapAvailable = pixmaps_supported && XShmPixmapFormat(_x11Display) == ZPixmap;

    return true;
}

Image<ColorRgb> & X11Grabber::grab()
{
    if (_XRenderAvailable && !_useXGetImage) {
	XRenderComposite( _x11Display,		// *dpy,
			  PictOpSrc,		// op,
			  _srcPicture,		// src
			  None,			// mask
			  _dstPicture,		// dst
			  _cropLeft,		// src_x
			  _cropTop,		// src_y
			  0,			// mask_x
			  0,			// mask_y
			  0,			// dst_x
			  0,			// dst_y
			  _croppedWidth,	// width
			  _croppedHeight);	// height
	
	XSync(_x11Display, False);
      
	if (_XShmAvailable) {
	    XShmGetImage(_x11Display, _pixmap, _xImage, 0, 0, AllPlanes);
	} else {
	    _xImage = XGetImage(_x11Display, _pixmap, 0, 0, _croppedWidth, _croppedHeight, AllPlanes, ZPixmap);   
	}
    } else {
	if (_XShmAvailable && !_useXGetImage) {
	    XShmGetImage(_x11Display, _window, _xImage, _cropLeft, _cropTop, AllPlanes);
	} else {
	  _xImage = XGetImage(_x11Display, _window, _cropLeft, _cropTop, _croppedWidth, _croppedHeight, AllPlanes, ZPixmap);
	}
    }
    
    if (_xImage == nullptr)
    {
        std::cerr << "X11GRABBER ERROR: Grab failed" << std::endl;
        return _image;
    }

    _imageResampler.processImage(reinterpret_cast<const uint8_t *>(_xImage->data), _xImage->width, _xImage->height, _xImage->bytes_per_line, PIXELFORMAT_BGR32, _image);

    return _image;
}

int X11Grabber::updateScreenDimensions()
{
    const Status status = XGetWindowAttributes(_x11Display, _window, &_windowAttr);
    if (status == 0)
    {
        std::cerr << "X11GRABBER ERROR: Failed to obtain window attributes" << std::endl;
        return -1;
    }

    if (_screenWidth == unsigned(_windowAttr.width) && _screenHeight == unsigned(_windowAttr.height))
    {
        // No update required
        return 0;
    }
    
    std::cout << "X11GRABBER INFO: Update of screen resolution: [" << _screenWidth << "x" << _screenHeight <<"] => ";

    if (_screenWidth || _screenHeight)
      freeResources();
    
    _screenWidth  = _windowAttr.width;
    _screenHeight = _windowAttr.height;
    
    std::cout << "[" << _screenWidth << "x" << _screenHeight <<"]" << std::endl;
    
    if (_screenWidth > unsigned(_cropLeft + _cropRight))
      _croppedWidth  = _screenWidth - _cropLeft - _cropRight;
    else
      _croppedWidth  = _screenWidth;
    
    if (_screenHeight > unsigned(_cropTop + _cropBottom))
      _croppedHeight = _screenHeight - _cropTop - _cropBottom;
    else
      _croppedHeight = _screenHeight;
      
    std::cout << "X11GRABBER INFO: Using ";
    
    if (_XRenderAvailable && !_useXGetImage) {
	std::cout << "XRender for grabbing" << std::endl;
    } else {
	std::cout << "XGetImage for grabbing" << std::endl;
    }
    
    setupResources();

    return 0;
}