#include "stdafx.h"
#include "screencapture.h"

#ifdef RVC_OS_WIN
static HMONITOR GetPrimaryMonitorHandle()
{
	const POINT ptZero = {0, 0};
	return MonitorFromPoint(ptZero, MONITOR_DEFAULTTOPRIMARY);
}

static HBITMAP CopyScreenToBitmap(LPRECT lpRect)
{
	HDC         hScrDC, hMemDC;         // screen DC and memory DC     
	//	HBITMAP     hBitmap; //, 
	//	HBITMAP     hBitmap;
	//	HBITMAP     hOldBitmap;    // handles to deice-dependent bitmaps     
	int         nX, nY, nX2, nY2;       // coordinates of rectangle to grab     
	int         nWidth, nHeight;        // DIB width and height     
	int         xScrn, yScrn;           // screen resolution      

	HGDIOBJ     hOldBitmap , hBitmap;

	// check for an empty rectangle 
	if (IsRectEmpty(lpRect))       
		return NULL;      
	// create a DC for the screen and create     
	// a memory DC compatible to screen DC          

	MONITORINFOEX mi;
	mi.cbSize = sizeof(mi);
	GetMonitorInfoA(GetPrimaryMonitorHandle(), &mi);
	//hScrDC = CreateDCA("DISPLAY", NULL, NULL, NULL);     
	hScrDC = CreateDCA(mi.szDevice, NULL, NULL, NULL);     
	hMemDC = CreateCompatibleDC(hScrDC);      // get points of rectangle to grab  

	nX = lpRect->left;     
	nY = lpRect->top;     
	nX2 = lpRect->right;     
	nY2 = lpRect->bottom;      // get screen resolution      

	xScrn = GetDeviceCaps(hScrDC, HORZRES);     
	yScrn = GetDeviceCaps(hScrDC, VERTRES);      

	//make sure bitmap rectangle is visible      

	if (nX < 0)         
		nX = 0;     

	if (nY < 0)         
		nY = 0;     

	if (nX2 > xScrn)         
		nX2 = xScrn;     

	if (nY2 > yScrn)         
		nY2 = yScrn;      

	nWidth = nX2 - nX;     
	nHeight = nY2 - nY;      

	// create a bitmap compatible with the screen DC     

	hBitmap = CreateCompatibleBitmap(hScrDC, nWidth, nHeight);      

	// select new bitmap into memory DC     

	hOldBitmap =   SelectObject (hMemDC, hBitmap);      

	// bitblt screen DC to memory DC     

	BitBlt(hMemDC, 0, 0, nWidth, nHeight, hScrDC, nX, nY, SRCCOPY);     

	// select old bitmap back into memory DC and get handle to     
	// bitmap of the screen          

	hBitmap = SelectObject(hMemDC, hOldBitmap);      

	// clean up      

	DeleteDC(hScrDC);     
	DeleteDC(hMemDC);      

	// return handle to the bitmap      

	return (HBITMAP)hBitmap; 

}


static HPALETTE GetSystemPalette() 
{ 
	HDC hDC;                // handle to a DC 
	static HPALETTE hPal = NULL;   // handle to a palette 
	HANDLE hLogPal;         // handle to a logical palette 
	LPLOGPALETTE lpLogPal;  // pointer to a logical palette 
	int nColors;            // number of colors 

	// Find out how many palette entries we want. 

	MONITORINFOEX mi;
	mi.cbSize = sizeof(mi);
	GetMonitorInfoA(GetPrimaryMonitorHandle(), &mi);
	hDC = CreateDCA(mi.szDevice, NULL, NULL, NULL);     
	//hDC = CreateDC ( TEXT("DISPLAY"), NULL, NULL, NULL );

	if (!hDC) 
		return NULL; 

	nColors = 256; //PalEntriesOnDevice(hDC);   // Number of palette entries 

	// Allocate room for the palette and lock it. 

	hLogPal = GlobalAlloc(GHND, sizeof(LOGPALETTE) + nColors * 
		sizeof(PALETTEENTRY)); 

	// if we didn't get a logical palette, return NULL 

	if (!hLogPal) 
		return NULL; 

	// get a pointer to the logical palette 

	lpLogPal = (LPLOGPALETTE)GlobalLock(hLogPal); 

	// set some important fields 

	lpLogPal->palVersion = 0x300; //PALVERSION; 
	lpLogPal->palNumEntries = nColors; 

	// Copy the current system palette into our logical palette 

	GetSystemPaletteEntries(hDC, 0, nColors, 
		(LPPALETTEENTRY)(lpLogPal->palPalEntry)); 

	// Go ahead and create the palette.  Once it's created, 
	// we no longer need the LOGPALETTE, so free it.     

	hPal = CreatePalette(lpLogPal); 

	// clean up 

	GlobalUnlock(hLogPal); 
	GlobalFree(hLogPal); 

	DeleteDC(hDC);

	return hPal; 
} 

static WORD DibNumColors (VOID FAR * pv)
{     
	int bits;     
	LPBITMAPINFOHEADER lpbi;     
	LPBITMAPCOREHEADER lpbc;      
	lpbi = ((LPBITMAPINFOHEADER)pv);     
	lpbc = ((LPBITMAPCOREHEADER)pv);      
	
	/*  With the BITMAPINFO format headers, the size of the palette      
	*  is in biClrUsed, whereas in the BITMAPCORE - style headers, it      
	*  is dependent on the bits per pixel ( = 2 raised to the power of      
	*  bits/pixel).      
	*/     
	
	if (lpbi->biSize != sizeof(BITMAPCOREHEADER))     
	{         
		if (lpbi->biClrUsed != 0)             
			return (WORD)lpbi->biClrUsed;         
	
		bits = lpbi->biBitCount;     
	}     
	else         
		bits = lpbc->bcBitCount;      

	switch (bits)     
	{     
		case 1:         
			return 2;     
		case 4:         

			return 16;     
		case 8:         
			return 256;     
		default:         
			//  A 24 bitcount DIB has no color table          
			return 0;     

	} 

	
} 
static WORD PaletteSize (VOID FAR * pv) 
{     
	LPBITMAPINFOHEADER lpbi;     
	WORD NumColors;      
	lpbi      = (LPBITMAPINFOHEADER)pv;     
	NumColors = DibNumColors(lpbi);      
	if (lpbi->biSize == sizeof(BITMAPCOREHEADER))         
		return NumColors * sizeof(RGBTRIPLE);     
	else         
		return NumColors * sizeof(RGBQUAD); 
} 

static HANDLE AllocRoomForDIB(BITMAPINFOHEADER bi, HBITMAP hBitmap) 
{ 
	DWORD               dwLen; 
	HANDLE              hDIB; 
	HDC                 hDC; 
	LPBITMAPINFOHEADER  lpbi; 
	HANDLE              hTemp; 

	// Figure out the size needed to hold the BITMAPINFO structure 
	// (which includes the BITMAPINFOHEADER and the color table). 

	dwLen = bi.biSize + PaletteSize((LPSTR) &bi); 
	hDIB  = GlobalAlloc(GHND,dwLen); 

	// Check that DIB handle is valid 

	if (!hDIB) 
		return NULL; 

	// Set up the BITMAPINFOHEADER in the newly allocated global memory, 
	// then call GetDIBits() with lpBits = NULL to have it fill in the 
	// biSizeImage field for us. 

	lpbi  = (LPBITMAPINFOHEADER)GlobalLock(hDIB); 
	*lpbi = bi; 

	MONITORINFOEX mi;
	mi.cbSize = sizeof(mi);
	GetMonitorInfoA(GetPrimaryMonitorHandle(), &mi);
	hDC = CreateDCA(mi.szDevice, NULL, NULL, NULL);     

	GetDIBits(hDC, hBitmap, 0, (UINT) bi.biHeight, NULL, (LPBITMAPINFO)lpbi, 
		DIB_RGB_COLORS); 
	DeleteDC(hDC);


	// If the driver did not fill in the biSizeImage field, 
	// fill it in -- NOTE: this is a bug in the driver! 

	if (lpbi->biSizeImage == 0) 
		lpbi->biSizeImage = (DWORD)(lpbi->biWidth * lpbi->biBitCount) * lpbi->biHeight / 8; 

	// Get the size of the memory block we need 

	dwLen = lpbi->biSize + PaletteSize((LPSTR) &bi) + lpbi->biSizeImage; 

	// Unlock the memory block 

	GlobalUnlock(hDIB); 

	// ReAlloc the buffer big enough to hold all the bits  

	if (hTemp = GlobalReAlloc(hDIB,dwLen,0)) 
		return hTemp; 
	else 
	{ 
		// Else free memory block and return failure 

		GlobalFree(hDIB); 
		return NULL; 
	} 
} 


static HANDLE ChangeBitmapFormat(HBITMAP hBitmap, WORD wBitCount, DWORD dwCompression, HPALETTE hPal) 
{ 
	HDC                hDC;          // Screen DC 
	HANDLE               hNewDIB=NULL; // Handle to new DIB 
	BITMAP             Bitmap;       // BITMAP data structure 
	BITMAPINFOHEADER   bi;           // Bitmap info. header 
	LPBITMAPINFOHEADER lpbi;         // Pointer to bitmap header 
	HPALETTE           hOldPal=NULL; // Handle to palette 
	WORD               NewBPP;       // New bits per pixel 
	DWORD              NewComp;      // New compression format 

	// Check for a valid bitmap handle 

	if (!hBitmap) 
		return NULL; 

	// Validate wBitCount and dwCompression 
	// They must match correctly (i.e., BI_RLE4 and 4 BPP or 
	// BI_RLE8 and 8BPP, etc.) or we return failure 

	if (wBitCount == 0) 
	{ 
		NewComp = dwCompression; 
		if (NewComp == BI_RLE4) 
			NewBPP = 4; 
		else if (NewComp == BI_RLE8) 
			NewBPP = 8; 
		else // Not enough info */ 
			return NULL; 
	} 
	else if (wBitCount == 1 && dwCompression == BI_RGB) 
	{ 
		NewBPP = wBitCount; 
		NewComp = BI_RGB; 
	} 
	else if (wBitCount == 4) 
	{ 
		NewBPP = wBitCount; 
		if (dwCompression == BI_RGB || dwCompression == BI_RLE4) 
			NewComp = dwCompression; 
		else 
			return NULL; 
	} 
	else if (wBitCount == 8) 
	{ 
		NewBPP = wBitCount; 
		if (dwCompression == BI_RGB || dwCompression == BI_RLE8) 
			NewComp = dwCompression; 
		else 
			return NULL; 
	} 
	else if (wBitCount == 24 && dwCompression == BI_RGB) 
	{ 
		NewBPP = wBitCount; 
		NewComp = BI_RGB; 
	} 
	else 
		return NULL; 

	// Get info about the bitmap 

	GetObject(hBitmap, sizeof(BITMAP), (LPSTR)&Bitmap); 

	// Fill in the BITMAPINFOHEADER appropriately 

	bi.biSize               = sizeof(BITMAPINFOHEADER); 
	bi.biWidth              = Bitmap.bmWidth; 
	bi.biHeight             = Bitmap.bmHeight; 
	bi.biPlanes             = 1; 
	bi.biBitCount           = NewBPP; 
	bi.biCompression        = NewComp; 
	bi.biSizeImage          = 0; 
	bi.biXPelsPerMeter      = 0; 
	bi.biYPelsPerMeter      = 0; 
	bi.biClrUsed            = 0; 
	bi.biClrImportant       = 0; 

	// Go allocate room for the new DIB 

	hNewDIB = AllocRoomForDIB(bi, hBitmap); 
	if (!hNewDIB) 
		return NULL; 

	// Get a pointer to the new DIB 

	lpbi = (LPBITMAPINFOHEADER)GlobalLock(hNewDIB); 

	// If we have a palette, get a DC and select/realize it 

	if (hPal) 
	{ 
		MONITORINFOEX mi;
		mi.cbSize = sizeof(mi);
		GetMonitorInfoA(GetPrimaryMonitorHandle(), &mi);
		hDC = CreateDCA(mi.szDevice, NULL, NULL, NULL);    
		hOldPal = SelectPalette(hDC, hPal, FALSE); 
		RealizePalette(hDC); 
	} 

	// Call GetDIBits and get the new DIB bits 

	if (!GetDIBits(hDC, hBitmap, 0, (UINT) lpbi->biHeight, (LPSTR)lpbi + 
		(WORD)lpbi->biSize + PaletteSize((LPSTR)lpbi), (LPBITMAPINFO)lpbi, 
		DIB_RGB_COLORS)) 
	{ 
		GlobalUnlock(hNewDIB); 
		GlobalFree(hNewDIB); 
		hNewDIB = NULL; 
	} 

	// Clean up and return 

	if (hOldPal) 
	{ 
		SelectPalette(hDC, hOldPal, TRUE); 
		RealizePalette(hDC); 
		DeleteDC(hDC);
	} 

	// Unlock the new DIB's memory block 

	if (hNewDIB) 
		GlobalUnlock(hNewDIB); 

	return hNewDIB; 
}

static int SaveDib(HANDLE hDib, void *buf)
{
	BITMAPFILEHEADER    bmfHdr;     // Header for Bitmap file     
	LPBITMAPINFOHEADER  lpBI;       // Pointer to DIB info structure     
	DWORD               dwDIBSize;

	lpBI = (LPBITMAPINFOHEADER)GlobalLock(hDib);     
	if (!lpBI)     
	{         
		return -1;     
	}
	if (lpBI->biSize != sizeof(BITMAPINFOHEADER))     
	{         
		GlobalUnlock(hDib);         
		return -1;     
	}
	memcpy(buf, (LPSTR)lpBI + (WORD)lpBI->biSize + PaletteSize((LPSTR)lpBI), lpBI->biSizeImage);
	GlobalUnlock(hDib);
	return 0;
}

int screencapture_capture1(RECT *lprc, void *buf, int *size)
{
	int width = lprc->right - lprc->left;
	int height = lprc->bottom - lprc->top;
	int linesize = (width * 3 + 3) & 0xfffffffc;

	if (buf == NULL) {
		*size = linesize * height;
		return 0;
	}

	HBITMAP hBitmap = CopyScreenToBitmap(lprc);
	HPALETTE hPal = GetSystemPalette(); 
	HANDLE hDIB = ChangeBitmapFormat(hBitmap, 24, BI_RGB, hPal);
	SaveDib(hDIB, buf);
	GlobalFree(hDIB);
	DeleteObject(hPal);
	DeleteObject(hBitmap);
	return 0;
}

int screencapture_capture(RECT *lprc, void *buf, int *size)
{
	int width = lprc->right - lprc->left;
	int height = lprc->bottom - lprc->top;
	int linesize = (width * 3 + 3) & 0xfffffffc;

	if (buf == NULL) {
		*size = linesize * height;
		return 0;
	}

	assert(*size >= linesize * height);

	HDC dc = NULL;
	HDC memdc = NULL;
	HBITMAP hbm = NULL;
	int rc;
	BITMAPINFO bmpInfo;
	//HWND hWndDeskTop = GetDesktopWindow();
	void *data;

	MONITORINFOEX mi;
	mi.cbSize = sizeof(mi);
	GetMonitorInfoA(GetPrimaryMonitorHandle(), &mi);
	dc = CreateDCA(mi.szDevice, NULL, NULL, NULL);     
	//dc = GetDC(hWndDeskTop);
	//dc = CreateDCA("DISPLAY", 0, 0, 0);

	if (!dc)
		return -1;

	memdc = CreateCompatibleDC(dc);
	if (memdc) {
		memset(&bmpInfo, 0, sizeof(bmpInfo));
		bmpInfo.bmiHeader.biSize=sizeof(BITMAPINFOHEADER);
		bmpInfo.bmiHeader.biPlanes = 1;
		bmpInfo.bmiHeader.biWidth = width;
		bmpInfo.bmiHeader.biHeight = height;
		bmpInfo.bmiHeader.biCompression = BI_RGB;
		bmpInfo.bmiHeader.biBitCount = 24;
		bmpInfo.bmiHeader.biSizeImage = width * height * 3;
		hbm = CreateDIBSection(dc, &bmpInfo, DIB_RGB_COLORS, &data, NULL, 0);
		//hbm = CreateCompatibleBitmap(memdc, width, height);
	}
	if (!memdc || !hbm) {
		rc = -1;
		goto on_error;
	}
	HGDIOBJ old = SelectObject(memdc, hbm);
	BOOL bret = BitBlt(memdc, 0, 0, width, height, dc, lprc ? lprc->left : 0, lprc ? lprc->top : 0, SRCCOPY);
	if (!bret) {
		rc = -1;
		goto on_error;
	}
	SelectObject(memdc, old);
	memcpy(buf, data, bmpInfo.bmiHeader.biSizeImage);

	*size = linesize * height;

	rc = 0;

on_error:

	if (memdc)
		DeleteDC(memdc);
	if (hbm)
		DeleteObject(hbm);
	if (dc)
		DeleteDC(dc);//

	return rc;
}

#else

#include <X11/Xutil.h>
#include <X11/Xlib.h>

int screencapture_capture(RECT* lprc, void* buf, int* size) {
	int width = lprc->right - lprc->left;
	int height = lprc->bottom - lprc->top;
	int linesize = (width * 3 + 3) & 0xfffffffc;

	if (buf == NULL) {
		*size = linesize * height;
		return 0;
	}

	Display* display = XOpenDisplay(NULL);
	Window root = DefaultRootWindow(display);

	XImage* image = XGetImage(display, root, lprc ? lprc->left : 0, lprc ? lprc->top : 0, width, height, AllPlanes, ZPixmap);

	unsigned long red_mask = image->red_mask;
	unsigned long green_mask = image->green_mask;
	unsigned long blue_mask = image->blue_mask;

	for (int x = 0; x < width; x++){
		for (int y = 0; y < height; y++)
		{
			unsigned long pixel = XGetPixel(image, x, y);

			unsigned char blue = pixel & blue_mask;
			unsigned char green = (pixel & green_mask) >> 8;
			unsigned char red = (pixel & red_mask) >> 16;

			((unsigned char *)buf)[(x + width * y) * 3] = red;
			((unsigned char*)buf)[(x + width * y) * 3 + 1] = green;
			((unsigned char*)buf)[(x + width * y) * 3 + 2] = blue;
		}
	}

	*size = linesize * height;

	if (image != NULL) {
		XDestroyImage(image);
		image = NULL;
	}
	if (display != NULL) {
		XCloseDisplay(display);
		display = NULL;
	}

	return 0;
}


int getScreenSize(int* width, int* height) {
	Display* display = XOpenDisplay(NULL);
	int default_id = DefaultScreen(display);
	Screen* screen = ScreenOfDisplay(display, default_id);

	*width = screen->width;
	*height = screen->height;

	if (display != NULL) {
		XCloseDisplay(display);
		display = NULL;
	}

	return 0;
}

#endif


int screencapture_clipoff(int width, int height, void* buf, int n_rc, RECT* rcs)
{
	int i, k;
	char* p;

	for (i = 0; i < n_rc; ++i) {
		RECT* lprc = &rcs[i];
		for (k = lprc->top, p = (char*)buf + lprc->top * width * 3 + lprc->left * 3; k < lprc->bottom; ++k) {
			memset(p, 0, (lprc->right - lprc->left) * 3);
			p += width * 3;
		}
	}

	return 0;
}