#pragma once

#include "FlowSshCppIncludes.h"


namespace FlowSshCpp
{


// EnsureThrow, EnsureAbort
//
// EnsureThrow: Watch out for EnsureFailedException.
//	- In a handler thread it should be handled. Otherwise FlowSshC might exit the process.
//	- If it is passed onto FlowSshNet it should be manually converted to a manged exception;
//	  otherwise a generic description only is available to the user.
//
// EnsureAbort: Is somewhat more generic and direct in this regard. It
// stops the application right away if something is fatally wrong.

#define EnsureThrow(TEST)		((TEST) ? 0 : (EnsureFailure_Throw(#TEST, __FUNCTION__, __LINE__), 0))
#define EnsureAbort(TEST)		((TEST) ? 0 : (EnsureFailure_Abort(#TEST, __FUNCTION__, __LINE__)))


// Handles a EnsureThrow failure by throwing an EnsureFailedException describing the assert.
// Uses dynamic memory allocation, so MUST NOT be called from an operator new or delete.

void EnsureFailure_Throw(char const* test, char const* funcsig, unsigned int line);


// Handles a EnsureFailure_Abort failure as follows:
// - Outputs assertion info to CONOUT if the application has a console.
// - Outputs assertion info with MessageBox() if the application is interactive and does not have a console.
// - Exits the program with ENSUREABORT_EXITCODE.

enum { ENSUREABORT_EXITCODE = 34821 };
void EnsureFailure_Abort(char const* test, char const* funcsig, unsigned int line);



// NoCopy

class NoCopy
{
public:
	NoCopy() {}

private:
	NoCopy(NoCopy const&);
	void operator= (NoCopy const&);
};



// CriticalSection

class CriticalSection : public NoCopy
{
public:
	CriticalSection();
	~CriticalSection();

	void Lock();
	void Unlock();

private:
	CRITICAL_SECTION m_cs;
};


class CriticalSectionLocker : public NoCopy
{
public:
	CriticalSectionLocker(CriticalSection& cs);
	~CriticalSectionLocker() { Unlock(); }

	bool IsLocked() const { return m_locked; }
	void Lock();		// Locks only if unlocked
	void Unlock();		// Unlocks only if locked

private:
	CriticalSection& m_cs;
	bool m_locked;
};

typedef CriticalSectionLocker CsLocker;	// Alias



// RefCountable: a template class that implements thread-safe object management.
// The object will destroy itself when its last holder calls Release().
// Derive from this class to use it.

class RefCountable
{
private:
	enum
	{
		INIT_REFCOUNT = 0x60000000,		// Must not equal DESTROY_REFCOUNT
		DESTROY_REFCOUNT = 0x40000000,	// Must be less than INIT_REFCOUNT, with room to spare
	};

	mutable long volatile m_refCount;

protected:
	virtual void OnDestroy() const { delete this; }

public:
	RefCountable() : m_refCount(INIT_REFCOUNT) {}
	virtual ~RefCountable();

	RefCountable(RefCountable const&) : m_refCount(INIT_REFCOUNT) {}    // Copy construction should not affect ref count in new object
	RefCountable& operator= (RefCountable const&) { return *this; }		// Assignment should not affect ref count in destination object

	void AddRef() const
	{
		if (m_refCount == INIT_REFCOUNT)
			m_refCount = 1;
		else
			InterlockedIncrement(&m_refCount);
	}

	void Release() const;
};



// RefPtr: a smart pointer class that implements reference
// counting on objects derived from RefCountable.

template <class T>
struct RefData
{
	// No destructor. Safe to use as an element in data vector,
	// where it can be copied around as the vector is resized.
	RefData() : m_ptr(0) {}
	
	T* m_ptr;

	void Detach(T* ptr) 
	{ 
		RefData_Detach(m_ptr, ptr);
		m_ptr = ptr; 
	}
	
	template <class TypeOrDescendant>
	void Detach(RefData<TypeOrDescendant> const& x) 
	{ 
		RefData_Detach(m_ptr, x.m_ptr);
		m_ptr = x.m_ptr; 
	}

	static void RefData_Detach(RefCountable* curPtr, RefCountable* newPtr)
	{
		if (newPtr != 0) newPtr->AddRef();
		if (curPtr != 0) curPtr->Release();
	}
};


template <class T>
class RefPtrConst : public RefData<T>
{
public:
	// Constructors
	RefPtrConst() {}
	RefPtrConst(T const* ptr) { Detach((T*)ptr); }
	
	template <class TypeOrDescendant>
	RefPtrConst(RefPtrConst<TypeOrDescendant> const& sp) { Detach(sp); }

	RefPtrConst(RefPtrConst<T> const& sp) { Detach(sp); }

	// Destructor
	~RefPtrConst() { Detach(0); }

	// Copy operator
	RefPtrConst<T>& operator= (T const* ptr) { Detach((T*)ptr); return *this; }
	
	template <class TypeOrDescendant>
	RefPtrConst<T>& operator= (RefPtrConst<TypeOrDescendant> const& sp) { Detach(sp); return *this; }

	RefPtrConst<T>& operator= (RefPtrConst<T> const& sp) { Detach(sp); return *this; }
	
	// Cast
	template <class U>
	bool IsType() const { return DynamicCast<U>() != 0; }

	template <class U>
	U const* DynamicCast() const { if (!m_ptr) return 0; return dynamic_cast<U const*>(m_ptr); }

	// Reset: since operator= is overloaded to two pointer types, "*this = 0" is ambiguous
	void Reset() { Detach(0); }

	// Accessors
	T const* Get() const { return m_ptr; }
	T const* operator-> () const { return m_ptr; }
	T const& GetRef() const { EnsureThrow(m_ptr != 0); return *m_ptr; }
};


template <class T>
class RefPtr : public RefPtrConst<T>
{
public:
	// Constructors
	RefPtr() {}
	RefPtr(T* ptr) : RefPtrConst<T>(ptr) {}
	
	template <class TypeOrDescendant>
	RefPtr(RefPtr<TypeOrDescendant> const& sp) : RefPtrConst<T>(sp) {}

	RefPtr(RefPtr<T> const& sp) : RefPtrConst<T>(sp) {}

	// Copy operator
	RefPtr<T>& operator= (T* ptr) { Detach(ptr); return *this; }
	
	template <class TypeOrDescendant>
	RefPtr<T>& operator= (RefPtr<TypeOrDescendant> const& sp) { Detach(sp); return *this; }

	RefPtr<T>& operator= (RefPtr<T> const& sp) { Detach(sp); return *this; }
	
	// Cast
	template <class U>
	bool IsType() const { return dynamic_cast<U*>(m_ptr) != 0; }

	template <class U>
	U* DynamicCast() const { return dynamic_cast<U*>(m_ptr); }

	// Accessors
	T* Get() const { return m_ptr; }
	T* operator-> () const { return m_ptr; }
	T& GetRef() const { EnsureThrow(m_ptr != 0); return *m_ptr; }

private:
	RefPtrConst<T>& operator= (T const* ptr);
	RefPtrConst<T>& operator= (RefPtrConst<T> const& sp);
};


template <class T>
inline RefPtr<T> CastAsNonConst(RefPtrConst<T> const& x) { return (T*)x.Get(); }



// AutoReleaseRef

struct AutoReleaseRef
{
	AutoReleaseRef(RefCountable const* refCountable = nullptr) { Set(refCountable); }
	~AutoReleaseRef()
	{
		if (m_refCountable)
			m_refCountable->Release();
		m_refCountable = nullptr;
	}

	void Set(RefCountable const* refCountable) { m_refCountable = refCountable; }

private:
	RefCountable const* m_refCountable;
};



// Exception

class Exception : public std::exception
{
public:
	Exception(wchar_t const* msg);
	virtual ~Exception() {}
	
	wchar_t const* What() const { return m_wide; }
	char const* what() const { return m_narrow; }

private:
	struct Buffers : public RefCountable
	{
		Buffers() : m_wide(0), m_narrow(0) {}
		~Buffers();

		wchar_t* m_wide;
		char* m_narrow;
	};

	RefPtr<Buffers> m_buffers;
	wchar_t const*	m_wide;
	char const*		m_narrow;
};



// EnsureFailedException
//
// Thrown from EnsureFailure_Throw.

class EnsureFailedException : public Exception
{
public:
	EnsureFailedException(wchar_t const* msg) : Exception(msg) {}
};



// Data

class DataAllocException : public Exception  
{ 
public: 
	DataAllocException() : Exception(L"Error allocating data") {}
};


class Data
{
public:
	Data(BSTR bstr, bool owner = true);
	Data(unsigned char const* ptr, unsigned int size);	
	Data(unsigned int size = 0);		// Creates a writable object if size != 0
	Data(Data const& data);				// Both objects become unwritable
	~Data() { Resize(0); }

	Data& operator=(Data const& data);	// Both objects become unwritable

	void Resize(unsigned int size);		// Object will be writable if size != 0

	// Returns NULL if the object is unwritable. Otherwise, calling Resize()
	// or passing the object as a parameter assignment or copy construction
	// will invalidate the pointer returned here.
	unsigned char* GetWritablePtr() { return m_writablePtr; }

	unsigned char const* GetPtr() const { return m_ptr; }
	unsigned int GetSize() const { return m_size; }

private:
	static unsigned char* BStrToPtrSize(BSTR bstr, unsigned int& size);
	static BSTR PtrSizeToBStr(unsigned char const* ptr, unsigned int size);

	mutable bool m_owner;
	mutable BSTR m_bstr;
	mutable unsigned char* m_writablePtr;
	unsigned char const* m_ptr;
	unsigned int m_size;
	BSTR m_bstrTemp;
};



// MainBase

class MainBase : public RefCountable, public NoCopy
{};



// HandlerBase

class HandlerBusyException : public Exception
{
public:
	HandlerBusyException() 
		: Exception(L"You cannot pass the same handler object to several requests at once.") {}
};

class ArgumentNullRefPtrException : public Exception
{
public:
	ArgumentNullRefPtrException() : Exception(L"You cannot use a RefPtr(NULL) argument here.") {}
};


template <class H, bool bThrowOnNull>
class HandlerBase : public RefCountable, public NoCopy
{
protected:
	HandlerBase() : m_locked(0), m_isInProgress(false) {}

	// Helpers (static)

	static void S_TryLock(RefPtr<H> obj)
	{
		if (obj.Get()) obj->Base_TryLock();
		else if (bThrowOnNull) throw ArgumentNullRefPtrException();
	}
	static void S_OnStart(RefPtr<H> obj, MainBase* mainObj) { if (obj.Get()) obj->Base_OnStart(mainObj); }
	static void S_OnStartFailed(RefPtr<H> obj) { if (obj.Get()) obj->Base_OnStartFailed(); }

	// Helpers

	void Base_TryLock()
	{
		if (InterlockedIncrement(&m_locked) != 1)
		{ InterlockedDecrement(&m_locked); throw HandlerBusyException(); }
	}

	void Base_OnStart(MainBase* mainObj)
	{
		EnsureThrow(!!mainObj);
		this->AddRef();		 // released in XyHandler::FlowSshC after Base_OnDone
		m_mainObj = mainObj; // released in Base_OnDone after OnDone
		m_isInProgress = true;
		OnStart();
	}

	void Base_OnStartFailed()
	{
		if (m_isInProgress)
		{
			AutoReleaseRef releaseRef(this);
			m_mainObj.Reset();
			m_isInProgress = false;	InterlockedDecrement(&m_locked);
		}
	}
	
	void Base_OnDone()
	{
		EnsureAbort(m_isInProgress);
		RefPtrConst<MainBase> releaseRef = m_mainObj; m_mainObj.Reset();
		m_isInProgress = false;	InterlockedDecrement(&m_locked);
		OnDone();
	}

	virtual void OnStart() = 0;
	virtual void OnDone() = 0;
	bool IsDone() const { return InterlockedCompareExchange(&m_locked, 0, 0) == 0; }

private:
	mutable volatile LONG m_locked;
	bool m_isInProgress;
	RefPtrConst<MainBase> m_mainObj; // Keep mainObj alive while a handler on it is active
									 // (so we don't miss the OnDoneBase call)
};



// DoneEvent

class CreateEventException : public Exception
{ 
public: 
	CreateEventException() : Exception(L"CreateEvent failed.") {}
};


template <class T>
class DoneEvent : virtual public T
{
protected:
	// This object must be created on the heap.
	virtual ~DoneEvent() { CloseHandle(m_event); }

public:
	DoneEvent() : m_event(CreateEvent(0, true, true, 0))
		{ if (!m_event) throw CreateEventException(); }
	
	HANDLE GetDoneEvent() const { return m_event; }
	bool WaitDone(DWORD timeout = INFINITE) const // Returns true if done.
		{ return WaitForSingleObject(m_event, timeout) == WAIT_OBJECT_0; }
	bool IsDone() const { return WaitDone(0); } 

protected:
	void OnStart() { ResetEvent(m_event); T::OnStart(); }
	void OnDone()  { SetEvent(m_event); T::OnDone(); }

private:
	HANDLE m_event;
};



// DoneMsg

template <class T>
class DoneMsg : virtual public T
{
protected:
	// This object must be created on the heap.
	virtual ~DoneMsg() {}

public:
	DoneMsg(HWND wnd, unsigned int msg, unsigned int requestId=0) : m_wnd(wnd), m_msg(msg), m_requestId(requestId) {}	

protected:
	void OnDone()
	{
		if (IsWindow(m_wnd)) SendMessageW(m_wnd, m_msg, WPARAM(this), LPARAM(m_requestId));
		T::OnDone(); 
	}

	HWND m_wnd;
	unsigned int m_msg;
	unsigned int m_requestId;
};



// Describe operation and error codes

std::wstring DescribeSftpErrCode(unsigned int sftpErrCode);
std::wstring DescribeListOp(unsigned int listOp);
std::wstring DescribeTransferOp(unsigned int transferOp);
std::wstring DescribeResumeErrCode(unsigned int resumeErrCode);



} // namespace FlowSshCpp
