#pragma once

// standard locker for multi thread.

#if defined(WIN32) || defined(WIN64)
  // windows mutex.
  #define WIN32_LEAN_AND_MEAN
  #include <windows.h>
  #include <assert.h>
#ifndef GNAPI
  #define GNAPI __stdcall
#endif

#define mutex_t           CRITICAL_SECTION
#define mutex_init(m)     InitializeCriticalSection(m)
#define mutex_acquire(m)  EnterCriticalSection(m)
#define mutex_release(m)  LeaveCriticalSection(m)
#define mutex_destroy(m)  DeleteCriticalSection(m)

#else
// Linux mutex.
#include <pthread.h>
#define GNAPI
#define mutex_t           pthread_mutex_t
#define mutex_init(m)     pthread_mutex_init(m, nullptr)
#define mutex_acquire(m)  pthread_mutex_lock(m)
#define mutex_release(m)  pthread_mutex_unlock(m)
#define mutex_destroy(m)  pthread_mutex_destroy(m)

#endif

namespace KCP {
	// interface of mutex.
	class IGNMutex {
	public:
		virtual ~IGNMutex() {}
		virtual void GNAPI Lock() = 0;
		virtual void GNAPI Unlock() = 0;
		virtual void GNAPI Release() = 0;
	};

	// locker
	class CLocker : public IGNMutex {
	public:
		CLocker() {
			mutex_init(&m_oCS);
		}
		virtual ~CLocker() {
			mutex_destroy(&m_oCS);
		}

		virtual void GNAPI Lock() {
			mutex_acquire(&m_oCS);
		}

		virtual void GNAPI Unlock() {
			mutex_release(&m_oCS);
		}

		virtual void GNAPI Release() {
		}

	protected:
		mutex_t m_oCS;
	};

	// ????
	class CNonMutex : public IGNMutex {
	public:
		CNonMutex() {}
		virtual ~CNonMutex() {}

	public:
		virtual void GNAPI Lock() {}
		virtual void GNAPI Unlock() {}
		virtual void GNAPI Release() {}
	};

	// ?????
	template <typename MT = CNonMutex>
	class CAutoLock {
	public:
		CAutoLock(MT& oLocker) : m_oLocker(oLocker) {
			m_oLocker.Lock();
		}
		~CAutoLock() {
			m_oLocker.Unlock();
		}

	protected:
		MT& m_oLocker;
	};

	// rwlock
#if defined(WIN32) || defined(WIN64)
#define RWLOCK_IDLE 0 
#define RWLOCK_R 0x01 
#define RWLOCK_W 0x02

	class CRWLock {
	private:
		int              _st;           
		int              _rlockCount;   
		int              _rwaitingCount; 
		int              _wwaitingCount; 
		HANDLE           _ev;         
		CRITICAL_SECTION _stLock;

	public:
		CRWLock()
			: _rlockCount(0),
			_st(RWLOCK_IDLE),
			_rwaitingCount(0),
			_wwaitingCount(0) {
			InitializeCriticalSection(&_stLock);
			_ev = CreateEvent(nullptr, TRUE, FALSE, nullptr);
			assert(_ev != INVALID_HANDLE_VALUE);
		}
		~CRWLock() {
			DeleteCriticalSection(&_stLock);
			CloseHandle(_ev);
		}

	public:
		void rlock() {
			bool isWaitReturn = false;
			while (1) {
				EnterCriticalSection(&_stLock);
				if (isWaitReturn) {
					--_rwaitingCount;
				}
				if (_st == RWLOCK_IDLE) {
					_st = RWLOCK_R;
					_rlockCount++;
					//ReleaseMutex(_stLock);
					LeaveCriticalSection(&_stLock);
					break;
				}
				else if (_st == RWLOCK_R) {
					if (_wwaitingCount > 0) {
						++_rwaitingCount;
						ResetEvent(_ev);
						LeaveCriticalSection(&_stLock);

						WaitForSingleObject(_ev, INFINITE);

						isWaitReturn = true;
					}
					else {
						++_rlockCount;
						LeaveCriticalSection(&_stLock);
						break;
					}
				}
				else if (_st == RWLOCK_W) {
					++_rwaitingCount;
					ResetEvent(_ev);
					//SignalObjectAndWait(_stLock, _ev, INFINITE, FALSE);
					LeaveCriticalSection(&_stLock);
					WaitForSingleObject(_ev, INFINITE);
					isWaitReturn = true;
				}
				else {
					assert(0);
					break;
				}
			}
		}
		void wlock() {
			bool isWaitReturn = false;
			while (1) {
				//WaitForSingleObject(_stLock, INFINITE);
				EnterCriticalSection(&_stLock);
				if (isWaitReturn) --_wwaitingCount;
				if (_st == RWLOCK_IDLE) {
					_st = RWLOCK_W;
					LeaveCriticalSection(&_stLock);
					break;
				}
				else {
					++_wwaitingCount;
					ResetEvent(_ev);
					LeaveCriticalSection(&_stLock);
					WaitForSingleObject(_ev, INFINITE);
					isWaitReturn = true;
				}
			}
		}
		void unlock() {
			// WaitForSingleObject(_stLock, INFINITE);
			EnterCriticalSection(&_stLock);
			if (_rlockCount > 0) {
				--_rlockCount;
				if (0 == _rlockCount) {
					_st = RWLOCK_IDLE;
					if (_wwaitingCount > 0 || _rwaitingCount > 0) {
						SetEvent(_ev);
					}
					else {
						/* ???? */
					}
				}
				else {
					/* ???ж??? */
				}
			}
			else {
				_st = RWLOCK_IDLE;
				if (_wwaitingCount > 0 || _rwaitingCount > 0) {
					SetEvent(_ev);
				} else {
					/* ???ж??? */
				}
			}
			//ReleaseMutex(_stLock);
			LeaveCriticalSection(&_stLock);
		}
	};

#else
	class CRWLock { // Linux rwlock.
	protected:
		int m_rd_count;
		int m_wr_count;

	public:
		CRWLock() :m_rd_count(0), m_wr_count(0) {
			::pthread_rwlock_init(&m_rwlock, nullptr);
		}
		~CRWLock() {
			::pthread_rwlock_destroy(&m_rwlock);
		}

		inline void rdlock() {
			::pthread_rwlock_rdlock(&m_rwlock);
			++m_rd_count;
		};

		inline void wrlock() {
			::pthread_rwlock_wrlock(&m_rwlock);
			++m_wr_count;
			++m_rd_count;
		}

		inline void unlock() {
			::pthread_rwlock_unlock(&m_rwlock);
			--m_rd_count;
		}
	private:
		pthread_rwlock_t m_rwlock;		/**< ????д?? */
	};
#endif

	// common mutex module
	IGNMutex* GNAPI GetGNMutex();
}