/********************************************************************
 * File: DisplayWin32.cpp
 * Author: Keith Schwarz (htiek@cs.stanford.edu)
 *
 * Implementation of the Display class for the Win32 API.  This ends up
 * being a bit complicated due to the fact that we have to explicitly
 * set up a message loop.  To get around this, the implementation fires
 * off a thread for each window that enters a message loop.
 */
#include "Display.h"
#include "grid.h"
#include <windows.h>
#include <stdexcept>
#include <limits>
#include <tchar.h>
#include <iostream>
#include <ddraw.h>
using namespace std;
using namespace Copper3D;

/* Win32 headers define this incorrectly. */
#undef max

/****** Module Constants ******/
const TCHAR* const kDisplayClassName = _T("Copper3D_Display"); // Window class name
const UINT kTerminateMessage = WM_APP + 1; // Message meaning "kill the window."
typedef void (*PixelPlotFunction)(UCHAR* surface, size_t pitch, size_t x, size_t y, const Copper3D::Color& c);

/* Pixel-plotting functions for different color depths and formats. */
static void PlotPixel32(UCHAR* surface, size_t pitch, size_t x, size_t y, const Copper3D::Color& color) {
	/* Bump surface pointer to correct location. */
	surface += x * 4 + y * pitch; // Four bytes per pixel, plus the skip for the pitch.

	/* Write RBG value. */
	*surface++ = (unsigned char)(color.b() * 255);
	*surface++ = (unsigned char)(color.g() * 255);
	*surface++ = (unsigned char)(color.r() * 255);
}

static void PlotPixel24(UCHAR* surface, size_t pitch, size_t x, size_t y, const Copper3D::Color& color) {
	/* Bump surface pointer to correct location. */
	surface += x * 3 + y * pitch; // Three bytes per pixel, plus the pitch.

	/* Write RBG value. */
	*surface++ = (unsigned char)(color.r() * 255);
	*surface++ = (unsigned char)(color.g() * 255);
	*surface++ = (unsigned char)(color.b() * 255);
}

/* Utility function to zero memory in a type-safe manner. */
template <typename DDType> static void InitDDObject(DDType& dd) {
	memset(&dd, 0, sizeof(dd));
	dd.dwSize = sizeof(dd);
}

/* Utility class that represents a generalized semaphore. */
class Semaphore {
public:
	explicit Semaphore(LONG initialCount);
	~Semaphore();

	void acquire();
	void release();

private:
	/* The actual semaphore. */
	HANDLE semaphore;

	/* No copying. */
	Semaphore(const Semaphore&);
	Semaphore& operator= (const Semaphore&);
};

/* Actual display class. */
class DisplayWin32: public Display {
public:
	DisplayWin32(size_t width, size_t height);
	~DisplayWin32();

	virtual void plotPixel(size_t x, size_t y, Color c);
	virtual void clear();
	virtual void display();

private:
	/* Thread function that manages the window. */
	static DWORD WINAPI ManagerThread(LPVOID me);

	/* Utility function to set up the window class. */
	static bool RegisterWindowClass();

	/* Window procedure. */
	static LRESULT CALLBACK WindowProcedure(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

	HWND hWnd;              // Handle to the window
	Semaphore windowReady;  // Semaphore for the main thread to communicate back that
	                        //   the window has been created.
	HANDLE processThread;   // Handle to the thread that processes messages.

	Semaphore   imageSemaphore;
	grid<Copper3D::Color> image;

	LPDIRECTDRAWSURFACE7 surface; // The surface used for off-screen rendering
	LPDIRECTDRAWCLIPPER  clipper; // The clipper used to prevent DirectDraw from bleeding
	                              // off the edge of the screen.
	PixelPlotFunction    plotFn;  // The function used to draw points.
};

/* Helper class to ensure that DirectDraw is set up and cleaned up appropriately. */
static class DirectDrawWrapper {
public:
	DirectDrawWrapper();
	~DirectDrawWrapper();

	LPDIRECTDRAW7 directDraw();
	LPDIRECTDRAWSURFACE7 primarySurface();

private:
	LPDIRECTDRAW7 mDD;
	LPDIRECTDRAWSURFACE7 mPrimary;
} dd;

/****** DirectDrawWrapper Implementation ******/
DirectDrawWrapper::DirectDrawWrapper() {
	/* Create the main DirectDraw object. */
	HRESULT hr;
	if (FAILED(hr = DirectDrawCreateEx(NULL, (void **)&mDD, IID_IDirectDraw7, NULL)))
		throw runtime_error("Couldn't initialize the DirectDraw objects.");

	/* Say we want windowed mode. */
	if (FAILED(hr = mDD->SetCooperativeLevel(NULL, DDSCL_NORMAL)))
		throw runtime_error("Couldn't set the cooperative level to normal.");

	/* Create the primary surface. */
	DDSURFACEDESC2 surfaceDescriptor;
	InitDDObject(surfaceDescriptor);

	/* Want a primary surface. */
	surfaceDescriptor.dwFlags = DDSD_CAPS;
	surfaceDescriptor.ddsCaps.dwCaps = DDSCAPS_PRIMARYSURFACE;

	/* Make the surface. */
	if (FAILED(hr = mDD->CreateSurface(&surfaceDescriptor, &mPrimary, NULL))) {
		/* Need to release DD. */
		mDD->Release();
		throw runtime_error("Could not create primary surface.");
	}
}

DirectDrawWrapper::~DirectDrawWrapper() {
	/* Release our hold on these objects. */
	mPrimary->Release();
	mDD->Release();
}

/* Accessors are pretty straightforward. */
LPDIRECTDRAW7 DirectDrawWrapper::directDraw() {
	return mDD;
}
LPDIRECTDRAWSURFACE7 DirectDrawWrapper::primarySurface() {
	return mPrimary;
}

/****** Semaphore Implementation *******/

/* Semaphore constructor bends over backwards ensuring that the semaphore is properly
 * initialized.
 */
Semaphore::Semaphore(LONG initialCount) {
	SetLastError(0);
	semaphore = CreateSemaphore(NULL, initialCount, numeric_limits<LONG>::max(), NULL);
	if (!semaphore || GetLastError())
		throw runtime_error("Error creating semaphore.");
}

/* Cleans up the semaphore.  It had better not still be in use! */
Semaphore::~Semaphore() {
	CloseHandle(semaphore);
}

/* Acquires or releases the semaphore. */
void Semaphore::acquire() {
	WaitForSingleObject(semaphore, INFINITE);
}
void Semaphore::release() {
	ReleaseSemaphore(semaphore, 1, NULL);
}

/****** DisplayWin32 Implementation ******/

/* Constructor forwards the request, then fires off the thread responsible for
 * maintaining the window.
 */
DisplayWin32::DisplayWin32(size_t width, size_t height) 
  : Display(width, height), windowReady(0), imageSemaphore(1) {
	hWnd = NULL;
	clipper = NULL;
	surface = NULL;
	image.resize(height, width);
	fill(image.begin(), image.end(), Copper3D::Color(0.0f, 0.0f, 0.0f));

	/* Fire off the loading thread. */
	processThread = CreateThread(NULL, 0, ManagerThread, this, 0, NULL);
	if (!processThread)
		throw runtime_error("Couldn't create the Display manager thread.");

	/* Wait for the thread to set up the window. */
	windowReady.acquire();

	/* Check that the window is good. */
	if (!hWnd)
		throw runtime_error("Couldn't create the Display.");

	clear();
}

/* Destructor kills the window and waits for the thread to return. */
DisplayWin32::~DisplayWin32() {
	PostMessage(hWnd, kTerminateMessage, 0, 0);
	WaitForSingleObject(processThread, INFINITE);
	if (clipper) clipper->Release();
	if (surface) surface->Release();
}

/* Registers the window class if it hasn't already been registered. */
bool DisplayWin32::RegisterWindowClass() {
	/* Ensure we only call this once. */
	static bool registered = false;
	if (registered) return true;

	WNDCLASSEX windowClass;
	windowClass.cbSize = sizeof(WNDCLASSEX);
	windowClass.style = CS_OWNDC | CS_DBLCLKS;
	windowClass.lpfnWndProc = WindowProcedure;
	windowClass.cbWndExtra = 0;
	windowClass.cbClsExtra = 0;
	windowClass.hInstance = GetModuleHandle(NULL);

	windowClass.hIcon = LoadIcon(NULL, IDI_APPLICATION);
	windowClass.hCursor = LoadCursor(NULL, IDC_CROSS);
	windowClass.hbrBackground = (HBRUSH)GetStockObject(BLACK_BRUSH);
	windowClass.lpszMenuName = NULL;
	windowClass.lpszClassName = kDisplayClassName;
	windowClass.hIconSm = LoadIcon(NULL, IDI_APPLICATION);
	
	/* Make sure registration worked! */
	if (!RegisterClassEx(&windowClass))
		return false;

	registered = true;
	return true;
}

/* Manager thread needs to create the window and forward messages to it. */
DWORD DisplayWin32::ManagerThread(LPVOID voidMe) {
	/* Recover the window. */
	DisplayWin32* me = static_cast<DisplayWin32*>(voidMe);

	/* Ensure that the window class is registered. */
	if (!RegisterWindowClass()) {
		me->windowReady.release();
		return 0;
	}

	/* Create the window. */
	HWND hWnd = CreateWindowEx(0, kDisplayClassName, _T("Copper3D"), WS_VISIBLE | WS_BORDER, 0, 0, me->width(), me->height(), NULL, NULL, GetModuleHandle(NULL), NULL);
	if (!hWnd) {
		me->windowReady.release();
		return 0;
	}

	/* Create the clipper. */
	if (FAILED(dd.directDraw()->CreateClipper(0, &me->clipper, NULL))) {
		DestroyWindow(hWnd);
		me->windowReady.release();
		return 0;
	}

	/* Attach the clipper to the window. */
	if (FAILED(me->clipper->SetHWnd(0, hWnd))) {
		DestroyWindow(hWnd);
		me->clipper->Release();
		me->windowReady.release();
		return 0;
	}

	/* Create the surface. */
	DDSURFACEDESC2 ddsd;
	InitDDObject(ddsd);
	
	ddsd.dwFlags = DDSD_CAPS | DDSD_HEIGHT | DDSD_WIDTH;
	ddsd.ddsCaps.dwCaps = DDSCAPS_SYSTEMMEMORY;
	ddsd.dwHeight = me->width();
	ddsd.dwWidth  = me->height();

	/* Try making the surface. */
	if (FAILED(dd.directDraw()->CreateSurface(&ddsd, &me->surface, NULL))) {
		DestroyWindow(hWnd);
		me->clipper->Release();
		me->windowReady.release();
		return 0;
	}

	/* Figure out which pixel-plotting function we're using. */
	DDPIXELFORMAT pixelFormat;
	memset(&pixelFormat, 0, sizeof(pixelFormat));
	pixelFormat.dwSize = sizeof(pixelFormat);

	/* Query pixel format. */
	if (FAILED(me->surface->GetPixelFormat(&pixelFormat))) {
		DestroyWindow(hWnd);
		me->clipper->Release();
		me->surface->Release();
		me->windowReady.release();
		return 0;
	}

	switch (pixelFormat.dwRGBBitCount) {
	case 32:
		me->plotFn = PlotPixel32;
		break;
	case 24:
		me->plotFn = PlotPixel24;
		break;
	default:
		DestroyWindow(hWnd);
		me->clipper->Release();
		me->surface->Release();
		me->windowReady.release();
		return 0;
	}

	/* Set the object to be its private data. */
	SetLastError(0);
	if (!SetWindowLongPtr(hWnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(me)) && GetLastError()) {
		DestroyWindow(hWnd);
		me->clipper->Release();
		me->surface->Release();
		me->windowReady.release();
		return 0;
	}

	/* Window is ready! */
	me->hWnd = hWnd;
	me->windowReady.release();

	/* Main event loop. */
	while (true) {
		MSG message;
		BOOL result = GetMessage(&message, NULL, 0, 0);

		/* If the message is bad, get out of here. */
		if (result == -1)
			return 0;

		/* If the message is a quit, we're done. */
		if (result == 0)
			return 0;

		/* Otherwise process the message. */
		TranslateMessage(&message);
		DispatchMessage(&message);
	}
}

/* Uses DirectDraw to clear a surface to pure black. */
static void ClearSurface(LPDIRECTDRAWSURFACE7 surface) {
	/* Use Blt for this one. */
	DDBLTFX fx;
	InitDDObject(fx);
	fx.dwFillColor = 0;
	surface->Blt(NULL, NULL, NULL, DDBLT_COLORFILL | DDBLT_WAIT, &fx);
}

/* Locks the DirectDraw surface. */
static DDSURFACEDESC2 LockSurface(LPDIRECTDRAWSURFACE7 surface) {
	DDSURFACEDESC2 result;
	InitDDObject(result);

	surface->Lock(NULL, &result, DDLOCK_SURFACEMEMORYPTR | DDLOCK_WAIT, NULL);

	return result;
}

/* Renders an image. */
static void PlotImage(const DDSURFACEDESC2& data, LPDIRECTDRAWSURFACE7 target,
					  const grid<Copper3D::Color>& image, PixelPlotFunction plotFn) {
	for (size_t x = 0; x < image.numCols(); ++x)
		for (size_t y = 0; y < image.numRows(); ++y)
			plotFn((UCHAR*)data.lpSurface, data.lPitch, x, y, image[y][x]);
}

/* Unlocks a locked surface. */
static void UnlockSurface(LPDIRECTDRAWSURFACE7 surface) {
	surface->Unlock(NULL);
}

/* Attaches a clipper to a surface. */
static void SetClipper(LPDIRECTDRAWCLIPPER clipper) {
	dd.primarySurface()->SetClipper(clipper);
}

/* Gets the rectangle for a window after factoring in other considerations. */
static RECT GetCorrectWindowRect(HWND window) {
	RECT rc;
	GetWindowRect(window, &rc);
	return rc;
}

/* Given a surface, renders that image to the primary window surface. */
static void BltImage(LPDIRECTDRAWSURFACE7 source, RECT where, RECT sourceRect) {
	dd.primarySurface()->Blt(&where, source, &sourceRect, DDBLT_WAIT, NULL);
}

/* Window procedure to handle events. */
LRESULT CALLBACK DisplayWin32::WindowProcedure(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam) {
	if (msg == kTerminateMessage) {
		DestroyWindow(hWnd);
		return 0;
	}
	if (msg == WM_DESTROY) {
		PostQuitMessage(0);
		return 0;
	}
	if (msg == WM_PAINT) {
		/* Get the caller and abort if one isn't set. */
		DisplayWin32* me = reinterpret_cast<DisplayWin32*>(GetWindowLongPtr(hWnd, GWLP_USERDATA));
		if (!me || !me->surface || !me->clipper) 
			return DefWindowProc(hWnd, msg, wParam, lParam);

		PAINTSTRUCT ps;
		BeginPaint(hWnd, &ps);
		
		me->imageSemaphore.acquire();

		/* Okay... fun times.  We need to:
		 * 1. Clear the surface.
		 * 2. Lock the surface.
		 * 3. Plot each pixel on the surface.
		 * 4. Unlock the surface.
		 * 5. Set the clipper
		 * 6. Get the window position
		 * 7. Blt the image.
		 */
		ClearSurface(me->surface);
		DDSURFACEDESC2 data = LockSurface(me->surface);
		PlotImage(data, me->surface, me->image, me->plotFn);
		UnlockSurface(me->surface);
		SetClipper(me->clipper);
		RECT targetRect = GetCorrectWindowRect(hWnd);
		RECT sourceRect = {0, 0, me->width(), me->height()};
		BltImage(me->surface, targetRect, sourceRect);

		me->imageSemaphore.release();

		EndPaint(hWnd, &ps);
	}
	return DefWindowProc(hWnd, msg, wParam, lParam);
}

/* Plotting a pixel just sets the proper entry in the image grid. */
void DisplayWin32::plotPixel(size_t x, size_t y, Color c) {
	image[y][x] = c;
}

void DisplayWin32::clear() {
	/* Wipe out the grid. */
	fill(image.begin(), image.end(), Copper3D::Color(0.0f, 0.0f, 0.0f));
}

/* Displaying works by invalidating the window; the WM_PAINT message actually
 * does the work.
 */
void DisplayWin32::display() {
	RECT rc;
	SetRect(&rc, 0, 0, width(), height());
	InvalidateRect(hWnd, &rc, FALSE);
	UpdateWindow(hWnd);
}

/* Factory function implementation. */
Display* Display::New(size_t width, size_t height) {
	return new DisplayWin32(width, height);
}