#include "stdafx.h"

namespace MSTSCLib
{
#include "mstsclib_h.h"
};

namespace MSTSCLib_Redist
{
// extremely ew, but actually the cleanest way to import the alternate UUIDs
#include "mstsclib_redist_i.c"
};

#include "rdesktop/rdesktop.h"
#include "rdesktop/proto.h"

namespace
{
#ifdef _MSC_VER
extern "C" char __ImageBase;
#endif

HMODULE GetCurrentModule()
{
	return reinterpret_cast<HMODULE>(&__ImageBase);
}

}

namespace
{

LONG g_moduleRefCount = 0;

void lockServer()
{
	InterlockedIncrement(&g_moduleRefCount);
}

void unlockServer()
{
	InterlockedDecrement(&g_moduleRefCount);
}

bool canUnloadServer()
{
	return g_moduleRefCount == 0;
}

}

namespace
{

void FreeLpsz(LPSTR lpsz)
{
	if(lpsz)
		delete[] lpsz;
}

LPSTR AllocLpsz(const CHAR * lpsz, size_t cb)
{
	LPSTR lpszNew = new CHAR[cb + 1];

	if(lpszNew == NULL)
		return NULL;

	CopyMemory(lpszNew, lpsz, cb);
	lpszNew[cb] = 0;

	return lpszNew;
}

LPSTR AllocLpsz(const WCHAR * lpwsz, int cchIn)
{
	int cch = WideCharToMultiByte(CP_ACP, 0, lpwsz, cchIn, NULL, 0, NULL, NULL);

	if(cch <= 0)
		return NULL;

	LPSTR lpsz = new CHAR[cch];

	if(lpsz == NULL)
		return NULL;

	cch = WideCharToMultiByte(CP_ACP, 0, lpwsz, cchIn, lpsz, cch, NULL, NULL);

	if(cch <= 0)
	{
		FreeLpsz(lpsz);
		return NULL;
	}

	return lpsz;
}

LPSTR BstrToLpsz(BSTR bstr)
{
	return AllocLpsz(bstr, SysStringLen(bstr));
}

BSTR LpszToBstr(LPSTR lpsz)
{
	int cch = MultiByteToWideChar(CP_ACP, 0, lpsz, -1, NULL, 0);

	if(cch <= 0)
		return NULL;

	BSTR bstr = SysAllocStringLen(NULL, cch);

	if(bstr == NULL)
		return NULL;

	cch = MultiByteToWideChar(CP_ACP, 0, lpsz, -1, bstr, cch);

	if(cch <= 0)
	{
		SysFreeString(bstr);
		return NULL;
	}

	return bstr;
}

}

namespace
{

template<class T, class U> T aligndown(const T& X, const U& align)
{
	return X & ~(T(align) - 1);
}

template<class T, class U> T alignup(const T& X, const U& align)
{
	return aligndown(X + (align - 1), align);
}

/* Convert between bitmap formats */
uint8 * win32_convert_scanlines(int width, int height, int bitcount, int fromalign, int toalign, const uint8 * data, uint8 ** buffer)
{
	// TBD: profile & optimize the most common cases
	assert(width > 0);
	assert(height);
	assert(bitcount && bitcount <= 32);
	assert(fromalign <= toalign);
	assert(data);
	assert(buffer);

	bool flipped = height < 0;

	if(flipped)
		height = - height;

	int bytesperrow = alignup(width * bitcount, 8) / 8;
	int fromstride = alignup(bytesperrow, fromalign);
	int tostride = alignup(bytesperrow, toalign);
	assert(fromstride <= tostride);

	int datasize = tostride * height;

	uint8 * dibits = new uint8[datasize];

	const uint8 * src = data;
	uint8 * dest = dibits;

	const int pad = tostride - fromstride;

	assert(pad < 4);
	__assume(pad < 4);

	if(flipped)
	{
		dest += (height - 1) * tostride;
		tostride = - tostride;
	}

	for(int i = 0; i < height; ++ i)
	{
		memcpy(dest, src, fromstride);
		memset(dest + fromstride, 0, pad);
		src += fromstride;
		dest += tostride;
	}

	*buffer = dibits;
	return dibits;
}

/* Creates bitmaps */
HBITMAP win32_create_dib(LONG width, LONG height, WORD bitcount, const BYTE * data)
{
	struct b_
	{
		BITMAPINFO bmi;
		RGBQUAD colormap[256 - ARRAYSIZE(RTL_FIELD_TYPE(BITMAPINFO, bmiColors))];
	}
	b;

	b.bmi.bmiHeader.biSize = sizeof(b.bmi.bmiHeader);
	b.bmi.bmiHeader.biWidth = width;
	b.bmi.bmiHeader.biHeight = height;
	b.bmi.bmiHeader.biPlanes = 1;
	b.bmi.bmiHeader.biBitCount = bitcount;
	b.bmi.bmiHeader.biCompression = BI_RGB;
	b.bmi.bmiHeader.biSizeImage = 0;
	b.bmi.bmiHeader.biXPelsPerMeter = 0;
	b.bmi.bmiHeader.biYPelsPerMeter = 0;

	if(bitcount > 8)
	{
		b.bmi.bmiHeader.biClrUsed = 0;
		b.bmi.bmiHeader.biClrImportant = 0;
	}
	else
	{
		b.bmi.bmiHeader.biClrUsed = 2 << bitcount;
		b.bmi.bmiHeader.biClrImportant = 2 << bitcount;

		// TODO: palette
	}

	// FIXME: beyond ugly
	HDC hdc = CreateCompatibleDC(NULL);

	if(hdc == NULL)
		return NULL;

	HBITMAP hbm = CreateDIBitmap(hdc, &b.bmi.bmiHeader, CBM_INIT, data, &b.bmi, DIB_RGB_COLORS);

	if(hbm == NULL)
		error("CreateDIBitmap %dx%dx%d failed\n", width, height, bitcount);

	DeleteDC(hdc);
	return hbm;
}

/* Creates brushes */
HBRUSH win32_create_brush(BRUSH * brush, COLORREF fgcolour)
{
	if(brush == NULL)
		return (HBRUSH)GetStockObject(NULL_BRUSH);

	switch(brush->style)
	{
	case BS_SOLID:
	case BS_NULL:
	case BS_HATCHED:
	case BS_PATTERN:
	case BS_PATTERN8X8:
		break;

	default:
		return NULL;
	}

	switch(brush->style)
	{
	case BS_SOLID:
		return CreateSolidBrush(fgcolour);

	case BS_HATCHED:
		return CreateHatchBrush(brush->pattern[0], fgcolour);

	case BS_NULL:
		return (HBRUSH)GetStockObject(NULL_BRUSH);

	case BS_PATTERN:
	case BS_PATTERN8X8:
		{
			uint16 pattern[8];

			for(size_t i = 0; i < 8; ++ i)
				pattern[i] = brush->pattern[i];

			HBITMAP hpattern = CreateBitmap(8, 8, 1, 1, pattern);
			HBRUSH hbr = CreatePatternBrush(hpattern);
			DeleteObject(hpattern);
			return hbr;
		}

	DEFAULT_UNREACHABLE;
	}
}
};

/*
	"sealed" can improve optimizations by asserting a class cannot be derived
	from, optimizing out accesses to the v-table from inside the class
*/
#if defined(_MSC_VER) && _MSC_VER >= 1400
#define SEALED_ sealed
#else
#define SEALED_
#endif


/* Class that implements the RDP client GUI */
class RdpClientUI
{
public:
	// TODO: pass the client settings relevant to the GUI here
	HRESULT Initialize(HWND hwndParent)
	{
		// TODO: create the various windows
		// TODO: create display window thread
		// TODO: create input thread
		return E_FAIL;
	}

public:
	static BOOL Startup()
	{
		WNDCLASSEX wcexUI = { sizeof(wcexUI) };
		WNDCLASSEX wcexConsole = { sizeof(wcexConsole) };
		WNDCLASSEX wcexDisplay = { sizeof(wcexDisplay) };
		WNDCLASSEX wcexInput = { sizeof(wcexInput) };

		HBRUSH nullBrush = (HBRUSH)GetStockObject(HOLLOW_BRUSH);

		wcexUI.lpfnWndProc = NULL; // TODO
		wcexUI.hInstance = GetCurrentModule();
		wcexUI.hCursor = LoadCursor(NULL, IDC_ARROW);
		wcexUI.hbrBackground = nullBrush;
		wcexUI.lpszClassName = TEXT("MissTosca_UI");

		wcexConsole.style = CS_VREDRAW | CS_HREDRAW;
		wcexConsole.lpfnWndProc = NULL; // TODO
		wcexConsole.hInstance = GetCurrentModule();
		wcexConsole.hCursor = LoadCursor(NULL, IDC_ARROW);
		wcexConsole.hbrBackground = nullBrush;
		wcexConsole.lpszClassName = TEXT("MissTosca_Console");

		wcexDisplay.style = CS_VREDRAW | CS_HREDRAW;
		wcexDisplay.lpfnWndProc = NULL; // TODO
		wcexDisplay.hInstance = GetCurrentModule();
		wcexDisplay.hCursor = LoadCursor(NULL, IDC_ARROW);
		wcexDisplay.hbrBackground = nullBrush;
		wcexDisplay.lpszClassName = TEXT("MissTosca_Display");
		
		wcexInput.style = CS_VREDRAW | CS_HREDRAW;
		wcexInput.lpfnWndProc = NULL; // TODO
		wcexInput.hInstance = GetCurrentModule();
		wcexInput.hCursor = NULL;
		wcexInput.hbrBackground = nullBrush;
		wcexInput.lpszClassName = TEXT("MissTosca_Input");

		return
			RegisterClassEx(&wcexUI) &&
			RegisterClassEx(&wcexConsole) &&
			RegisterClassEx(&wcexDisplay) &&
			RegisterClassEx(&wcexInput);
	}

	static void Shutdown()
	{
		// TODO
	}

	/*
		This is the main UI window. It's the direct child of the control
		window, it fills its whole extent and it contains the scrollbars.
		When activated, it will move keyboard focus to the input window
	*/
private:
	HWND m_uiWindow;
	LONG m_scrollHPos;
	LONG m_scrollVPos;

	LRESULT UIWindowProc(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		switch(uMsg)
		{
			// Keep the keyboard focus on the input window
		case WM_ACTIVATE:
			switch(LOWORD(wParam))
			{
			case WA_INACTIVE:
				break;

			case WA_ACTIVE:
			case WA_CLICKACTIVE:
				if(!HIWORD(wParam))
					SetFocus(m_inputWindow);
			}

			return 0;

			// Resized: rearrange children windows, adjust scrollbars
		case WM_SIZE:
			{
				if(IsIconic(m_uiWindow))
					break;

				RECT rcClient;
				GetWindowRect(m_uiWindow, &rcClient);

				if(m_smartSizing)
				{
					// we are not supposed to maintain aspect ratio. Container has to do that
					m_consoleX = 0;
					m_consoleY = 0;
					m_consoleWidth = rcClient.right;
					m_consoleHeight = rcClient.bottom;
				}
				else
				{
					// center horizontally, no horizontal scrollbar
					if(rcClient.right >= m_consoleWidth)
						m_consoleX = (m_consoleWidth - rcClient.right) / 2;

					// center vertically, no vertical scrollbar
					if(rcClient.bottom >= m_consoleHeight)
						m_consoleY = (m_consoleHeight - rcClient.right) / 2;
				}

				SCROLLINFO scroll = { sizeof(scroll), SIF_ALL, 0 };
			
				// update the horizontal scrollbar
				scroll.nMax = m_consoleWidth;
				scroll.nPage = rcClient.right;
				scroll.nPos = 0 - m_consoleX;
				SetScrollInfo(m_uiWindow, SB_HORZ, &scroll, TRUE);

				// update the vertical scrollbar
				scroll.nMax = m_consoleHeight;
				scroll.nPage = rcClient.bottom;
				scroll.nPos = 0 - m_consoleY;
				SetScrollInfo(m_uiWindow, SB_VERT, &scroll, TRUE);

				// move/resize the console window
				MoveWindow(m_consoleWindow, m_consoleX, m_consoleY, m_consoleWidth, m_consoleHeight, TRUE);
			}

			return 0;

		case WM_HSCROLL:
			{
				SCROLLINFO scroll = { sizeof(scroll), SIF_TRACKPOS };
				GetScrollInfo(m_uiWindow, SB_HORZ, &scroll);
				m_consoleX = - scroll.nTrackPos;
				MoveWindow(m_consoleWindow, m_consoleX, m_consoleY, m_consoleWidth, m_consoleHeight, TRUE);
			}

			return 0;

		case WM_VSCROLL:
			{
				SCROLLINFO scroll = { sizeof(scroll), SIF_TRACKPOS };
				GetScrollInfo(m_uiWindow, SB_VERT, &scroll);
				m_consoleY = - scroll.nTrackPos;
				MoveWindow(m_consoleWindow, m_consoleX, m_consoleY, m_consoleWidth, m_consoleHeight, TRUE);
			}

			return 0;

		default:
			break;
		}

		return DefWindowProc(m_uiWindow, uMsg, wParam, lParam);
	}

	/*
		This is the full-screen title bar. It's displayed at the top of the
		main UI window while in full-screen mode, and it contains two toolbars
		with the pin, minimize, restore and close buttons
	*/
	HWND m_fullScreenBarWindow;

	/*
		This is the console window. It has the same extent as the display on
		the remote computer, or it fills the UI window in smart resizing mode,
		and it contains the input and display windows
	*/
private:
	HWND m_consoleWindow;
	LONG m_consoleX;
	LONG m_consoleY;
	LONG m_consoleWidth;
	LONG m_consoleHeight;
	bool m_smartSizing;

	LRESULT ConsoleWindowProc(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		switch(uMsg)
		{
		case WM_SIZE:
			{
				RECT rcClient;
				GetClientRect(m_consoleWindow, &rcClient);

				MoveWindow(m_inputWindow, 0, 0, rcClient.right, rcClient.bottom, TRUE);
				MoveWindow(m_displayWindow, 0, 0, rcClient.right, rcClient.bottom, TRUE);
			}

			return 0;

		default:
			break;
		}

		return DefWindowProc(m_consoleWindow, uMsg, wParam, lParam);
	}

	/*
		This is the display window. It represents the virtual display of the
		remote computer. It completely fills its parent, the console window,
		and it runs in its own thread for performance reasons
	*/
private:
	HWND m_displayWindow;
	LONG m_displayBufferWidth;
	LONG m_displayBufferHeight;
	HDC m_displayBuffer;
	void * m_displayBufferRaw;
	int m_displayBufferSave;
	int m_displayBufferBitDepth;
	int m_displayBufferByteDepth;
	int m_displayBufferStride;
	RECT m_displayBufferClip;
	CRITICAL_SECTION m_displayBufferMutex;

	LRESULT DisplayWindowProc(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		switch(uMsg)
		{
		case WM_DESTROY:
			PostQuitMessage(0);
			return 0;

		case WM_PRINTCLIENT:
			if(wParam == 0)
				break;

		case WM_PAINT:
			{
				HDC hdc = (HDC)wParam;

				EnterCriticalSection(&m_displayBufferMutex);

				if(hdc)