#include "stdafx.h"
#include "Bitmap.h"

void DeleteScreen(Screen *);

Screen *
CreateScreen(void)
{
	Screen *scrn = GlobalAlloc(GPTR, sizeof(Screen));
	RECT rectScreen = { 0 };

	if (!scrn) {
		return NULL;
	}

	scrn->hdcScreen = CreateDC("DISPLAY", NULL, NULL, NULL);
	scrn->hdcBitmap = CreateCompatibleDC(scrn->hdcScreen);

	rectScreen.right = GetDeviceCaps(scrn->hdcScreen, HORZRES);
	rectScreen.bottom = GetDeviceCaps(scrn->hdcScreen, VERTRES);

	scrn->hScreenRgn = CreateRectRgnIndirect(&rectScreen);
	if (!scrn->hScreenRgn) {
		// Store the error to be bubbled up
		DWORD err = GetLastError();
		DeleteScreen(scrn);
		SetLastError(err);

		return NULL;
	}

	return scrn;
}

void
DeleteScreen(Screen *scrn)
{
	if (scrn) {
		if (scrn->hOldObj) SelectObject(scrn->hdcBitmap, scrn->hOldObj);
		if (scrn->hBitmap) DeleteObject(scrn->hBitmap);

		if (scrn->hScreenRgn) DeleteObject((HGDIOBJ) scrn->hScreenRgn);
		if (scrn->hdcBitmap) DeleteDC(scrn->hdcBitmap);
		if (scrn->hdcScreen) DeleteDC(scrn->hdcScreen);

		GlobalFree(scrn);
	}

	return;
}

#include <pshpack1.h>
typedef struct _FULLBMPHDR {
	BITMAPFILEHEADER file;
	BITMAPINFOHEADER info;
} FULLBMPHDR, *PFULLBMPHDR;
#include <poppack.h>

PFULLBMPHDR
GenerateBitmapInfo(Screen *scrn,
				   BITMAP *bmp)
{
	// GDI seems to modify the DWORD after the info header, even if we
	// don't need a palette (>8 bit BMPs). Thus, an extra RGBQUAD just
	// to be safe.
	HANDLE hHeap = GetProcessHeap();
	PFULLBMPHDR pBmpHdr = HeapAlloc(hHeap,
		HEAP_ZERO_MEMORY,
		sizeof(struct _FULLBMPHDR) + sizeof(RGBQUAD));
	RECT rcScrn = {0};
	WORD cClrBits = (WORD)(bmp->bmPlanes * bmp->bmBitsPixel);
	DWORD hdrSize;

	if (!pBmpHdr) return NULL;

	if (!GetRgnBox(scrn->hScreenRgn, &rcScrn)) {
		HeapFree(hHeap, 0, pBmpHdr);
		return NULL;
	}

	if (cClrBits <= 16) {
		cClrBits = 16;
	} else {
		cClrBits = 32;
	}

	pBmpHdr->info.biSize = sizeof(BITMAPINFOHEADER);
	pBmpHdr->info.biWidth = bmp->bmWidth;
	pBmpHdr->info.biHeight = bmp->bmHeight;
	pBmpHdr->info.biPlanes = bmp->bmPlanes;
	pBmpHdr->info.biBitCount = cClrBits;
	pBmpHdr->info.biCompression = BI_RGB;
	// According to MSDN docs, this alignment to WORD boundaries is
	// required for Win9x GDI.
	pBmpHdr->info.biSizeImage = (((bmp->bmWidth * cClrBits + 31) & ~31) >> 3)
		* bmp->bmHeight;

	pBmpHdr->file.bfType = 0x4d42; // "BM"
	// actual bits offset: file header + info header + colors used
	hdrSize = (DWORD) (sizeof(BITMAPFILEHEADER) +
		pBmpHdr->info.biSize +
			((cClrBits >= 16) ? 0 : (pBmpHdr->info.biClrUsed *
				sizeof(RGBQUAD))));
	// file size is header size (offset above) + bits size
	pBmpHdr->file.bfSize = hdrSize + pBmpHdr->info.biSizeImage;
	pBmpHdr->file.bfOffBits = hdrSize;

	return pBmpHdr;
}

int
SnapScreen(Screen *scrn)
{
	RECT rcScreen = {0};
	int iScrWidth,
		iScrHeight;
	if (!scrn) return FALSE;

	// Don't make a new hBitmap if we have one
	if (scrn->hBitmap) return TRUE;

	GetRgnBox(scrn->hScreenRgn, &rcScreen);

	iScrWidth = rcScreen.right - rcScreen.left;
	iScrHeight = rcScreen.bottom - rcScreen.top;

	scrn->hBitmap = CreateCompatibleBitmap(scrn->hdcScreen,
		iScrWidth,
		iScrHeight);
	if (!scrn->hBitmap) return FALSE;

	scrn->hOldObj = SelectObject(scrn->hdcBitmap, scrn->hBitmap);
	if (!scrn->hOldObj) {
		DWORD err = GetLastError();
		DeleteObject(scrn->hBitmap);
		scrn->hBitmap = NULL;
		SetLastError(err);
		return FALSE;
	}

	if (!BitBlt(scrn->hdcBitmap,
		0, 0,
		iScrWidth,
		iScrHeight,
		scrn->hdcScreen,
		rcScreen.left, rcScreen.top,
		SRCCOPY))
	{
		DWORD err = GetLastError();
		SelectObject(scrn->hdcBitmap, scrn->hOldObj);
		DeleteObject(scrn->hBitmap);
		scrn->hBitmap = NULL;
		SetLastError(err);

		return FALSE;
	}

	return TRUE;
}

int
WriteRegionToFile(Screen *scrn,
				  RECT *rgn,
				  LPTSTR fname)
{
	BITMAP bmp = {0};			// hBitmap's metadata
	HANDLE hOutFile = NULL,		// Output .bmp
		hHeap = GetProcessHeap();
	LPBYTE lpBits = NULL;		// Actual bitmap bits

	HDC hdcRgn = NULL;			// HDC for defined region
	HBITMAP hRgnBitmap = NULL;	// Bitmap of defined region
	HGDIOBJ hOldObj = NULL;		// Old bitmap, if any
	PFULLBMPHDR pHdr = NULL;	// Full bitmap header
	PBITMAPINFO pbmi = NULL;	// BMI header, in case bits <= 8

	WORD cClrBits = 0;			// How many bits per pixel

	BOOL hasErr = FALSE;		// If something went wrong, set this
	DWORD err = 0,				// Used to store LastError
		dwTmp = 0;				// Used for written bytes storage

	if (!scrn || (!rgn || !fname)) {
		SetLastError(ERROR_INVALID_PARAMETER);
		return -1;
	}

	if (!RectInRegion(scrn->hScreenRgn, rgn)) {
		// TODO: is there a better error we can set here?
		SetLastError(ERROR_INVALID_PARAMETER);
		return -1;
	}

	// Overlay already grabs the screen, so don't trample over it
	if (!scrn->hBitmap && !SnapScreen(scrn))
	{
		hasErr = TRUE;
		goto cleanup;
	}

	hdcRgn = CreateCompatibleDC(scrn->hdcScreen);
	if (!hdcRgn) {
		hasErr = TRUE;
		goto cleanup;
	}

	// Create a new bitmap for the specific region
	hRgnBitmap = CreateCompatibleBitmap(scrn->hdcScreen,
		rgn->right - rgn->left,
		rgn->bottom - rgn->top);
	if (!hRgnBitmap) {
		hasErr = TRUE;
		goto cleanup;
	}


	if (!GetObject(hRgnBitmap, sizeof(BITMAP), &bmp)) {
		hasErr = TRUE;
		goto cleanup;
	}

	hOldObj = SelectObject(hdcRgn, hRgnBitmap);
	if (!hOldObj) {
		hasErr = TRUE;
		goto cleanup;
	}

	if (!BitBlt(
		hdcRgn,
		0, 0,
		rgn->right - rgn->left,
		rgn->bottom - rgn->top,
		scrn->hdcBitmap,
		rgn->left, rgn->top,
		SRCCOPY))
	{
		hasErr = TRUE;
		goto cleanup;
	}

	pHdr = GenerateBitmapInfo(scrn,
		&bmp);
	if (!pHdr) {
		hasErr = TRUE;
		goto cleanup;
	}

	lpBits = HeapAlloc(hHeap,
		0,
		pHdr->info.biSizeImage);
	if (!lpBits) {
		hasErr = TRUE;
		goto cleanup;
	}

	if (pHdr->info.biBitCount <= 8) {
		pbmi = HeapAlloc(hHeap,
			HEAP_ZERO_MEMORY,
			sizeof(BITMAPINFOHEADER) +
				(pHdr->info.biClrUsed * sizeof(RGBQUAD)));
		if (!pbmi) {
			hasErr = TRUE;
			goto cleanup;
		}
		CopyMemory(&pbmi->bmiHeader, &pHdr->info, sizeof(BITMAPINFOHEADER));
	} else {
		// For all other bit sizes, we don't care about the color
		// table. Just use the info header directly.
		// (Yes this is pretty shady. Don't worry about it)
		pbmi = (PBITMAPINFO) &pHdr->info;
	}

	if (!GetDIBits(scrn->hdcBitmap,
		hRgnBitmap,
		0,
		(WORD) pHdr->info.biHeight,
		lpBits,
		pbmi,
		DIB_RGB_COLORS))
	{
		hasErr = TRUE;
		goto cleanup;
	}

	// Create and write to the output file

	hOutFile = CreateFile(fname,
		GENERIC_WRITE,
		0,
		NULL,
		CREATE_ALWAYS,
		FILE_ATTRIBUTE_NORMAL,
		NULL);
	if (hOutFile == INVALID_HANDLE_VALUE) {
		hasErr = TRUE;
		goto cleanup;
	}

	// First, write the header
	if (!WriteFile(hOutFile, pHdr, sizeof(FULLBMPHDR), &dwTmp, NULL)) {
		hasErr = TRUE;
		goto cleanup;
	}

	// Assuming we're only using 16- or 32-bit BMPs, we should never
	// need to write a color table. But since this could easily be
	// modified to support smaller bit widths (i.e., 256 or less),
	// let's make sure such cases don't crash.
	if (pHdr->info.biBitCount <= 8) {
		if (!WriteFile(hOutFile,
				pbmi,
				pbmi->bmiHeader.biClrUsed * sizeof(RGBQUAD),
				&dwTmp,
				NULL))
		{
			hasErr = TRUE;
			goto cleanup;
		}
	}

	// And last but not least, the pixels
	if (!WriteFile(hOutFile, lpBits, pHdr->info.biSizeImage, &dwTmp, NULL)) {
		hasErr = TRUE;
		goto cleanup;
	}

cleanup:
	// make sure to grab the error before we start freeing everything
	if (hasErr) {
		err = GetLastError();
	}

	if (hOutFile) {
		CloseHandle(hOutFile);
		if (hasErr) {
			// Since it's probably busted, try to delete it too
			DeleteFile(fname);
		}
	}

	if (lpBits) HeapFree(hHeap, 0, lpBits);

	// In case our bits are <= 8, pbmi will be an actual heap-alloc'd
	// pointer. Let's make sure to handle that.
	if (pbmi && pbmi != (PBITMAPINFO) &pHdr->info) HeapFree(hHeap, 0, pbmi);

	if (pHdr) HeapFree(hHeap, 0, pHdr);

	if (hOldObj) SelectObject(hdcRgn, hOldObj);
	if (hRgnBitmap) DeleteObject(hRgnBitmap);

	if (hdcRgn) DeleteDC(hdcRgn);

	if (hasErr) {
		SetLastError(err);
	}

	return hasErr ? -1 : 0;
}