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

LRESULT CALLBACK OverlayWndProc(HWND, UINT, WPARAM, LPARAM);

#define OVERLAY_CLASS _T("GrabbyOverlay")
#define OVERLAY_TITLE _T("Grabby Overlay")

// Static function to store and retrieve the overlay-chosen rect.
static void
SetGetRect(BOOL set, RECT *r)
{
	static RECT rcMain = {0};
	if (r) {
		if (set) {
			CopyMemory(&rcMain, r, sizeof(RECT));
		} else {
			CopyMemory(r, &rcMain, sizeof(RECT));
		}
	} 
	return;
}

void
GetChosenRect(RECT *rcOut)
{
	SetGetRect(FALSE, rcOut);
	return;
}

int
CreateOverlay(HINSTANCE hInstance,
				Screen *scrn)
{
	HWND hWnd = NULL;
	RECT rcScreen = {0};
	MSG msg = {0};
	BOOL bRet;
	WNDCLASSEX wcClass = {
		sizeof(WNDCLASSEX),
		CS_HREDRAW | CS_VREDRAW,
		OverlayWndProc,
		0,
		0,
		hInstance,
		LoadIcon(NULL, IDI_APPLICATION),
		LoadCursor(NULL, IDC_CROSS),
		(HBRUSH)( COLOR_WINDOW + 1 ),
		OVERLAY_TITLE,
		OVERLAY_CLASS,
		NULL
	};

	if (!RegisterClassEx(&wcClass)) {
		return -1;
	}

	if (!GetRgnBox(scrn->hScreenRgn, &rcScreen)) {
		UnregisterClass(OVERLAY_CLASS, hInstance);
		return -1;
	}

	hWnd = CreateWindowEx(
		WS_EX_TOPMOST,
		OVERLAY_CLASS,
		OVERLAY_TITLE,
		WS_VISIBLE | WS_POPUP,
		rcScreen.left,
		rcScreen.top,
		rcScreen.right - rcScreen.left,
		rcScreen.bottom - rcScreen.top,
		NULL,
		NULL,
		hInstance,
		scrn);

	if (!hWnd) {
		UnregisterClass(OVERLAY_CLASS, hInstance);
		return -1;
	}

	ShowWindow(hWnd, SW_SHOWNORMAL);

	while ((bRet = GetMessage(&msg, hWnd, 0, 0)) != 0) {
		if (bRet == -1) {
			// TODO: error handling
			break;
		} else {
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}

	if (hWnd) DestroyWindow(hWnd);
	UnregisterClass(OVERLAY_CLASS, hInstance);

	return (int)msg.wParam;
}

// Left, Top, Right, Bottom
#define POINT2RECT(a,b) {	\
	min(a.x, b.x),			\
	min(a.y, b.y),			\
	max(a.x, b.x),			\
	max(a.y, b.y)			\
}

LRESULT CALLBACK
OverlayWndProc(HWND hWnd,
			   UINT uMsg,
			   WPARAM wParam,
			   LPARAM lParam)
{
	HDC hDC = NULL;
	PAINTSTRUCT ps = {0};
	static POINTS ptStart = {0},		// Starting drag point
		ptCur = {0};					// Current/end drag point
	static BOOL isDrag = FALSE;			// Are we mid-drag?
	static Screen *scrn = NULL;			// Screen data
	static RECT rcScreen = {0};
	static HGDIOBJ hOldObj = NULL;

	switch (uMsg) {
	case WM_CREATE:
		{
			HDC dcWnd = GetDC(hWnd);
			// Get the Screen object passed in from CreateWindowEx
			HBITMAP hOldBmp = NULL;
			LPCREATESTRUCT csCreate = (LPCREATESTRUCT)lParam;
			if (csCreate && csCreate->lpCreateParams) {
				scrn = (Screen *)csCreate->lpCreateParams;
				if (!scrn) {
					// uhhhhhhhhh
					MessageBox(hWnd,
						_T("Couldn't get screen info for mysterious reasons !?"),
						_T("Fatal Error"),
						MB_OK | MB_ICONERROR);
					PostQuitMessage(1);
					return 1;
				}

				GetRgnBox(scrn->hScreenRgn, &rcScreen);

				if (!SnapScreen(scrn)) {
					MessageBox(hWnd,
						_T("Couldn't store screen data !!"),
						_T("Fatal Error"),
						MB_OK | MB_ICONERROR);
					PostQuitMessage(1);
					return 1;
				}
			}
		}
		/* PASSTHRU */
	case WM_PAINT:
		hDC = BeginPaint(hWnd, &ps);

		if (scrn->hBitmap) {
			// Draw the "background"
			if (!BitBlt(hDC,
				ps.rcPaint.left,
				ps.rcPaint.top,
				ps.rcPaint.right - ps.rcPaint.left,
				ps.rcPaint.bottom - ps.rcPaint.top,
				scrn->hdcBitmap,
				rcScreen.left + ps.rcPaint.left,
				rcScreen.top + ps.rcPaint.top,
				SRCCOPY))
			{
				// walk it off, I guess...
				OutputDebugString(_T("Failed to blit background to DC..."));
			}
		}

		if (isDrag && hDC) {
			RECT rcOuter = POINT2RECT(ptStart, ptCur);
			DrawFocusRect(hDC, &rcOuter);
		}

		EndPaint(hWnd, &ps);
		break;
	case WM_ERASEBKGND:
		// Do nothing
		return 1;
	case WM_MOUSEMOVE:
		if (isDrag && (wParam & MK_LBUTTON)) {
			POINTS ptPrev = ptCur;
			ptCur = MAKEPOINTS(lParam);
			if ((ptCur.x != ptStart.x) &&
				(ptCur.y != ptStart.y))
			{
				RECT rcPrevDmg = POINT2RECT(ptStart, ptPrev);
				// Expand out a little bit to contain the focus rect
				rcPrevDmg.left--;
				rcPrevDmg.right++;
				rcPrevDmg.top--;
				rcPrevDmg.bottom++;
				InvalidateRect(hWnd, &rcPrevDmg, TRUE);
			}

		}
		break;
	case WM_LBUTTONDOWN:
		if (!isDrag) {
			// Start the drag operation and set the origin
			ptStart = ptCur = MAKEPOINTS(lParam);
			isDrag = TRUE;
		}
		break;
	case WM_LBUTTONUP:
		if (isDrag) {
			// Done dragging, store the rect and get out of here
			RECT rcFinal = POINT2RECT(ptStart, ptCur);

			SetGetRect(TRUE, &rcFinal);
			isDrag = FALSE;
			PostQuitMessage(0);
		}
		break;
	case WM_KEYUP:
		if (wParam != VK_ESCAPE) {
			break;
		}
		/* PASSTHRU */
	case WM_RBUTTONUP:
		// We treat right click as Escape
		PostQuitMessage(1);
		break;
	default:
		return DefWindowProc(hWnd, uMsg, wParam, lParam);
	}

	return 0;
}