#include "FlowSshCppIncludes.h"
#include "FlowSshCppUtils.h"


namespace FlowSshCpp
{


// EnsureFailure_Throw, EnsureFailure_Abort

void EnsureFailure_Throw(char const* test, char const* funcsig, unsigned int line)
{
	enum { MaxModuleChars = 300 };
	wchar_t module[MaxModuleChars];
	{
		memset(module, 0, sizeof(module));
		GetModuleFileNameW(0, module, MaxModuleChars - 1);
		module[MaxModuleChars - 1] = 0;
	}

	enum { MaxBufChars = 1025 }; // Maximum supported by wsprintf() + 1
	wchar_t buf[MaxBufChars];
	{
		memset(buf, 0, sizeof(buf));
		wsprintfW(buf,
			L"Ensure check failed in module '%s', PID %u, function '%S', line %u:\r\n"
				L"%S",
			module, GetCurrentProcessId(), funcsig, line, test);
		buf[MaxBufChars - 1] = 0;
	}

	throw EnsureFailedException(buf);
}


void EnsureFailure_Abort(char const* test, char const* funcsig, unsigned int line)
{
	enum { MaxModuleChars = 300 };
	char module[MaxModuleChars];
	{
		memset(module, 0, sizeof(module));
		GetModuleFileNameA(0, module, MaxModuleChars - 1);
		module[MaxModuleChars - 1] = 0;
	}

	enum { MaxBufChars = 1025 }; // Maximum supported by wsprintf() + 1
	char buf[MaxBufChars];
	{
		memset(buf, 0, sizeof(buf));
		wsprintfA(buf,
			"\r\nEnsure check failed in module '%s', PID %u, TID %u, function '%s', line %u:\r\n"
				"%s\r\n",
			module, GetCurrentProcessId(), GetCurrentThreadId(), funcsig, line, test);
		buf[MaxBufChars - 1] = 0;
	}

	OutputDebugStringA(buf);

	HANDLE outHandle = GetStdHandle(STD_OUTPUT_HANDLE);
	HANDLE conOutHandle = 0;
	if (!outHandle || outHandle == INVALID_HANDLE_VALUE)
	{
		conOutHandle = CreateFileA("CONOUT$", GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, 0, OPEN_EXISTING, 0, 0);
		outHandle = conOutHandle;
	}

	if (outHandle != 0 && outHandle != INVALID_HANDLE_VALUE)
	{
		DWORD nrBytes;
		WriteFile(outHandle, buf, (DWORD) strlen(buf), &nrBytes, 0);

		if (conOutHandle != 0 && conOutHandle != INVALID_HANDLE_VALUE)
			CloseHandle(conOutHandle);
	}
	else
	{
		bool interactive = false;
		{
			HWINSTA winsta = GetProcessWindowStation();
			if (winsta != 0)
			{
				USEROBJECTFLAGS uof;
				DWORD lenNeeded = sizeof(uof);
				if (GetUserObjectInformationW(winsta, UOI_FLAGS, &uof, sizeof(uof), &lenNeeded) != 0)
					interactive = ((uof.dwFlags & WSF_VISIBLE) != 0);
			}
		}

		if (interactive)
			MessageBoxA(0, buf, "Ensure check failed", MB_ICONSTOP | MB_OK);

		// If not interactive, might want to write into a log file, but in that
		// case need to build appropriate security descriptor for the file.
	}

	ExitProcess(ENSUREABORT_EXITCODE);
}



// CriticalSection

CriticalSection::CriticalSection()
{
	::InitializeCriticalSection(&m_cs);
}

CriticalSection::~CriticalSection()
{
	::DeleteCriticalSection(&m_cs);
}

void CriticalSection::Lock()
{
	::EnterCriticalSection(&m_cs);
}

void CriticalSection::Unlock()
{
	::LeaveCriticalSection(&m_cs);
}


CriticalSectionLocker::CriticalSectionLocker(CriticalSection& cs)
	: m_cs(cs)
	, m_locked(false)
{
	Lock();
}

void CriticalSectionLocker::Lock() 
{ 
	if (!m_locked)
	{
		m_cs.Lock();
		m_locked = true; 
	}
}

void CriticalSectionLocker::Unlock() 
{ 
	if (m_locked)
	{
		m_cs.Unlock();
		m_locked = false; 
	}
}



// RefPtr

RefCountable::~RefCountable()
{
	// m_refCount must equal INIT_REFCOUNT or DESTROY_REFCOUNT at this point.
	// There are two permissible ways to destroy a RefCountable object:
	// - without AddRef() ever being called;
	// - if AddRef() was called, from Release() when m_refCount reaches zero.
	// At this point, if m_refCount > DESTROY_REFCOUNT && m_refCount != INIT_REFCOUNT,
	// then one of the derived destructors must have caused a RefPtr to the object
	// undergoing destruction to be stored somewhere, or did not properly release a reference.
	// If m_refCount < DESTROY_REFCOUNT, the object was destroyed with references still outstanding.

	EnsureAbort(m_refCount == DESTROY_REFCOUNT || m_refCount == INIT_REFCOUNT);
}

void RefCountable::Release() const
{
	if (m_refCount == 1)
		goto Destroy;

	long newRefCount = InterlockedDecrement(&m_refCount);
	if (newRefCount != 0)
		EnsureAbort(newRefCount > 0 && newRefCount < DESTROY_REFCOUNT - 1);
	else
	{
	Destroy:
		// AddRef() and Release() may be called in the descendant's destructor.
		m_refCount = DESTROY_REFCOUNT;
		OnDestroy();
	}
}



// Exception

Exception::Exception(wchar_t const* msg) 
	: m_wide(L"Error allocating memory for exception message")
	, m_narrow("Error allocating memory for exception message")
{
	m_buffers = new (std::nothrow) Buffers;
	if (m_buffers.Get())
	{
		size_t wideSizeChars = wcslen(msg) + 1;
		m_buffers->m_wide = new (std::nothrow) wchar_t[wideSizeChars];
		if (m_buffers->m_wide)
		{		
			wcscpy_s(m_buffers->m_wide, wideSizeChars, msg);
			m_wide = m_buffers->m_wide;

			int narrowSizeBytes = WideCharToMultiByte(CP_ACP, 0, msg, -1, 0, 0, 0, 0);
			if (!narrowSizeBytes)
				m_narrow = "Error converting exception message to ANSI code page";
			else
			{
				m_buffers->m_narrow = new (std::nothrow) char[(unsigned int)narrowSizeBytes];
				if (m_buffers->m_narrow)
				{
					int writtenChars = WideCharToMultiByte(CP_ACP, 0, msg, -1, m_buffers->m_narrow, narrowSizeBytes, 0, 0);
					if (!writtenChars)
						m_narrow = "Error converting exception message to ANSI code page";
					else
						m_narrow = m_buffers->m_narrow;
				}
			}
		}
	}
}

Exception::Buffers::~Buffers()
{
	if (m_wide)
		delete [] m_wide;

	if (m_narrow)
		delete [] m_narrow;
}



// Data

Data::Data(BSTR bstr, bool owner) 
	: m_owner(owner)
	, m_bstr(bstr)
	, m_writablePtr(0)
	, m_bstrTemp(0)
{ 
	m_ptr = BStrToPtrSize(m_bstr, m_size);
}

Data::Data(unsigned char const* ptr, unsigned int size) 
	: m_owner(false)
	, m_bstr(0)
	, m_writablePtr(0)
	, m_ptr(ptr)
	, m_size(size)
	, m_bstrTemp(0)
{
}

Data::Data(Data const& data) 
	: m_owner(false)
	, m_bstr(0)
	, m_bstrTemp(0)
{
	*this = data;
}

Data::Data(unsigned int size) 
	: m_owner(false)
	, m_bstr(0)
	, m_size(0)
	, m_bstrTemp(0)
{
	Resize(size);
}

Data& Data::operator=(Data const& data)
{
	Resize(0);
	
	if (data.m_bstr)
	{
		m_owner = data.m_owner;
		m_bstr = data.m_bstr;
		m_ptr = data.m_ptr;
		m_size = data.m_size;
	}
	else
	{
		m_owner = true;
		m_bstr = PtrSizeToBStr(data.m_ptr, data.m_size);
		m_ptr = BStrToPtrSize(m_bstr, m_size);
	}

	data.m_owner = false;
	data.m_writablePtr = 0;

	return *this;
}

void Data::Resize(unsigned int size)
{
	if (m_bstrTemp)
	{
		SysFreeString(m_bstrTemp);
		m_bstrTemp = 0;
	}

	if (!size)
	{
		if (m_bstr && m_owner)
			SysFreeString(m_bstr);

		m_bstr = 0;
		m_writablePtr = 0;
		m_ptr = 0;
		m_size = 0;
	}
	else
	{
		if (m_writablePtr && m_size >= size)
			m_size = size;
		else
		{
			m_bstrTemp = PtrSizeToBStr(0, size);
			if (m_size)
				memcpy(m_bstrTemp, m_ptr, std::min(m_size, size));
					
			if (m_bstr && m_owner)
				SysFreeString(m_bstr);

			m_owner = true;
			m_bstr = m_bstrTemp;
			m_bstrTemp = 0;
			m_writablePtr = BStrToPtrSize(m_bstr, m_size);
			m_ptr = m_writablePtr;
		}
	}
}

unsigned char* Data::BStrToPtrSize(BSTR bstr, unsigned int& size)
{
	if (!bstr)
	{
		size = 0;
		return 0;
	}
	else
	{
		size = SysStringByteLen(bstr);
		return (unsigned char*) bstr;
	}
}

BSTR Data::PtrSizeToBStr(unsigned char const* ptr, unsigned int size)
{
	BSTR bstr = 0;
	
	if (size)
	{
		bstr = SysAllocStringByteLen((char const*) ptr, size);
		if (!bstr)
			throw DataAllocException();
	}

	return bstr;
}	



// Describe operation and error codes

std::wstring DescribeSftpErrCode(unsigned int sftpErrCode)
{
	std::wstring desc;
	switch (sftpErrCode)
	{
	case FlowSshC_SftpErrCode_Eof:					desc = L"Eof"; break;
	case FlowSshC_SftpErrCode_NoSuchFile:			desc = L"NoSuchFile"; break;
	case FlowSshC_SftpErrCode_PermissionDenied:		desc = L"PermissionDenied"; break;
	case FlowSshC_SftpErrCode_Failure:				desc = L"Failure"; break;
	case FlowSshC_SftpErrCode_BadMessage:			desc = L"BadMessage"; break;
	case FlowSshC_SftpErrCode_NoConnection:			desc = L"NoConnection"; break;
	case FlowSshC_SftpErrCode_ConnectionLost:		desc = L"ConnectionLost"; break;
	case FlowSshC_SftpErrCode_OpUnsupported:		desc = L"OpUnsupported"; break;
	// SFTP4+
	case FlowSshC_SftpErrCode_InvalidHandle:		desc = L"InvalidHandle"; break;
	case FlowSshC_SftpErrCode_NoSuchPath:			desc = L"NoSuchPath"; break;
	case FlowSshC_SftpErrCode_FileAlreadyExists:	desc = L"FileAlreadyExists"; break;
	case FlowSshC_SftpErrCode_WriteProtect:			desc = L"WriteProtect"; break;
	case FlowSshC_SftpErrCode_NoMedia:				desc = L"NoMedia"; break;
	// SFTP5+
	case FlowSshC_SftpErrCode_NoSpaceOnFilesystem:	desc = L"NoSpaceOnFilesystem"; break;
	case FlowSshC_SftpErrCode_QuotaExceeded:		desc = L"QuotaExceeded"; break;
	case FlowSshC_SftpErrCode_UnknownPrincipal:		desc = L"UnknownPrincipal"; break;
	case FlowSshC_SftpErrCode_LockConflict:			desc = L"LockConflict"; break;
	// SFTP6+
	case FlowSshC_SftpErrCode_DirNotEmpty:			desc = L"DirNotEmpty"; break;
	case FlowSshC_SftpErrCode_NotADirectory:		desc = L"NotADirectory"; break;
	case FlowSshC_SftpErrCode_InvalidFilename:		desc = L"InvalidFilename"; break;
	case FlowSshC_SftpErrCode_LinkLoop:				desc = L"LinkLoop"; break;
	case FlowSshC_SftpErrCode_CannotDelete:			desc = L"CannotDelete"; break;
	case FlowSshC_SftpErrCode_InvalidParameter:		desc = L"InvalidParameter"; break;
	case FlowSshC_SftpErrCode_FileIsADirectory:		desc = L"FileIsADirectory"; break;
	case FlowSshC_SftpErrCode_ByteRangeLockConflict:desc = L"ByteRangeLockConflict"; break;
	case FlowSshC_SftpErrCode_ByteRangeLockRefused:	desc = L"ByteRangeLockRefused"; break;
	case FlowSshC_SftpErrCode_DeletePending:		desc = L"DeletePending"; break;
	case FlowSshC_SftpErrCode_FileCorrupt:			desc = L"FileCorrupt"; break;
	case FlowSshC_SftpErrCode_OwnerInvalid:			desc = L"OwnerInvalid"; break;
	case FlowSshC_SftpErrCode_GroupInvalid:			desc = L"GroupInvalid"; break;
	case FlowSshC_SftpErrCode_NoMatchingByteRangeLock:	desc = L"NoMatchingByteRangeLock"; break;
	default:
		desc = L"SftpErrCodeCode_" + std::to_wstring((unsigned __int64)sftpErrCode);
	}
	return desc;
}

std::wstring DescribeListOp(unsigned int listOp)
{
	std::wstring desc;
	switch (listOp)
	{
	// Codeless errors
	case FlowSshC_ListOp_Exception:			desc = L"Exception"; break;
	// SFTP errors - FlowSshC_SftpErrCode
	case FlowSshC_ListOp_ListDir:			desc = L"ListDir"; break;
	case FlowSshC_ListOp_ChannelNotOpen:	desc = L"ChannelNotOpen"; break;
	default:
		desc = L"ListOp_" + std::to_wstring((unsigned __int64)listOp);
	}
	return desc;
}

std::wstring DescribeTransferOp(unsigned int transferOp)
{
	std::wstring desc;
	switch (transferOp)
	{
	// Codeless errors
	case FlowSshC_TransferOp_Exception:		desc = L"Exception"; break;
	// SFTP errors - FlowSshC_SftpErrCode
	case FlowSshC_TransferOp_MakeRemoteDir:		desc = L"MakeRemoteDir"; break;
	case FlowSshC_TransferOp_OpenRemoteFile:	desc = L"OpenRemoteFile"; break;
	case FlowSshC_TransferOp_ReadRemoteFile:	desc = L"ReadRemoteFile"; break;
	case FlowSshC_TransferOp_WriteRemoteFile:	desc = L"WriteRemoteFile"; break;
	case FlowSshC_TransferOp_ChannelNotOpen:	desc = L"ChannelNotOpen"; break;
	// WinAPI errors - see MSDN
	case FlowSshC_TransferOp_MakeLocalDir:		desc = L"MakeLocalDir"; break;
	case FlowSshC_TransferOp_OpenLocalFile:		desc = L"OpenLocalFile"; break;
	case FlowSshC_TransferOp_ReadLocalFile:		desc = L"ReadLocalFile"; break;
	case FlowSshC_TransferOp_WriteLocalFile:	desc = L"WriteLocalFile"; break;
	// Resume errors - FlowSshC_ResumeErrCode
	case FlowSshC_TransferOp_ResumeLocalFile:	desc = L"ResumeLocalFile"; break;
	case FlowSshC_TransferOp_ResumeRemoteFile:	desc = L"ResumeRemoteFile"; break;
	default:
		desc = L"TransferOp_" + std::to_wstring((unsigned __int64)transferOp);
	}
	return desc;
}

std::wstring DescribeResumeErrCode(unsigned int resumeErrCode)
{
	std::wstring desc;
	switch (resumeErrCode)
	{
	case FlowSshC_ResumeErrCode_TextMode:		desc = L"TextMode"; break;
	case FlowSshC_ResumeErrCode_FileSize:		desc = L"FileSize"; break;
	case FlowSshC_ResumeErrCode_Conversion:		desc = L"Conversion"; break;
	case FlowSshC_ResumeErrCode_Content:		desc = L"Content"; break;
	case FlowSshC_ResumeErrCode_FileSizeError:	desc = L"FileSizeError"; break;
	default:
		desc = L"ResumeErrCode_" + std::to_wstring((unsigned __int64)resumeErrCode);
	}
	return desc;
}



} // namespace FlowSshCpp
