#include <climits></climits>

#include <cstring></cstring>

#include <fcntl.h></fcntl.h>

#include <unistd.h></unistd.h>

#include <sys types.h=""></sys>

#include <sys socket.h=""></sys>

#include <sys epoll.h=""></sys>

#include <sys eventfd.h=""></sys>

#include <netinet in.h=""></netinet>

#include <netinet tcp.h=""></netinet>

#include <netdb.h></netdb.h>

#include "protocol.h"

Protocol::~Protocol() {

	assert(state == STATE_NONE || state == STATE_CLOSED);

	closeWait();

	while(finishWaiters > 0) std::this_thread::yield();

	if (thread) thread->join();

}

bool Protocol::open() {

	State s = state;

	do {

		if (s != STATE_NONE && s != STATE_FINISHED) return false;

	} while(!state.compare_exchange_weak(s, STATE_OPENING));

	if (s == STATE_FINISHED) while(finishWaiters > 0) std::this_thread::yield();

	if (thread) thread->join();

	thread = new std::thread(&Protocol::threadRun, this);

	return true;

}

void Protocol::wantState(State state) {

	ReadLock lock(stateProtector);

	State s = this->state;

	if (s <= STATE_OPENING || s >= STATE_CLOSING || s >= state) return;

	State sw = stateWanted;

	do {

		if (sw <= STATE_OPENING || sw >= state) return;

	} while(!stateWanted.compare_exchange_weak(sw, state));

	wakeup();

}

void Protocol::closeReq()

	{ wantState(STATE_CLOSE_REQ); }

void Protocol::close()

	{ wantState(STATE_CLOSING); }

void Protocol::closeWait(unsigned long long timeoutUs, bool withReq) {

	stateProtector.lock();

	State s = state;

	if (s <= STATE_OPENING || s >= STATE_FINISHED)

		{ stateProtector.unlock(); return; }

	if (s == STATE_CLOSED) {

		finishWaiters++;

		stateProtector.unlock();

	} else {

		if (withReq && s < STATE_CLOSE_REQ) closeReq();

		closeWaiters++;

		stateProtector.unlock();

		std::mutex fakeMutex;

		{

			std::unique_lock<std::mutex> fakeLock(fakeMutex);</std::mutex>

			if (timeoutUs && s < STATE_CLOSING) {

				std::chrono::steady_clock::time_point t

					= std::chrono::steady_clock::now()

						+ std::chrono::microseconds(timeoutUs);

				while(s < STATE_CLOSED) {

					if (t <= std::chrono::steady_clock::now()) { close(); break; }

					closeCondition.wait_until(fakeLock, t);

					s = state;

				}

			}

			if (s < STATE_CLOSING) close();

			while(s < STATE_CLOSED) {

				closeCondition.wait(fakeLock);

				s = state;

			}

		}

		finishWaiters++;

		closeWaiters--;

	}

	while(s != STATE_FINISHED) { std::this_thread::yield(); s = state; }

	finishWaiters--;

}

void Protocol::wakeup() {

	unsigned long long num = 0;

	::write(eventFd, &num, sizeof(num));

}

void Protocol::updateEvents(Connection &connection) {

	unsigned int events = 0;

	if (connection.readQueue.count()) events |= EPOLLIN;

	if (connection.writeQueue.count()) events |= EPOLLOUT;

	if (connection.events.exchange(events) == events) return;

	struct epoll_event event = {};

	event.data.ptr = &connection;

	event.events = events;

	epoll_ctl(epollFd, EPOLL_CTL_MOD, connection.sockId, &event);

}

void Protocol::initSocket(int sockId) {

	int tcp_nodelay = 1;

	setsockopt(sockId, IPPROTO_TCP, TCP_NODELAY, &tcp_nodelay, sizeof(int));

	fcntl(sockId, F_SETFL, fcntl(sockId, F_GETFL, 0) | O_NONBLOCK);

}

void Protocol::threadConnQueue() {

	while(connQueue.peek()) {

		Connection &connection = connQueue.pop();

		connection.enqueued = false;

		if (connection.stateLocal == Connection::STATE_NONE) {

			int sockId = connection.sockId;

			bool connected = false;

			if (sockId >= 0) {

				initSocket(sockId);

			} else

			if (connection.addressSize >= sizeof(sa_family_t)) {

				struct sockaddr *addr = (struct sockaddr*)connection.address;

				sockId = ::socket(addr->sa_family, SOCK_STREAM, 0);

				if (sockId >= 0) {

					initSocket(sockId);

					int res = ::connect(sockId, addr, (int)connection.addressSize);

					if (res == 0) {

						connected = true;

					} else

					if (res != EINPROGRESS) {

						::close(sockId);

						sockId = -1;

					}

				}

			}

			if (sockId >= 0) {

				connection.protocol = this;

				connection.prev = connLast;

				assert(!connection.next);

				(connection.prev ? connection.prev->next : connFirst) = &connection;

				if (connection.server) {

					connection.srvPrev = connection.server->connLast;

					assert(!connection.srvNext);

					(connection.srvPrev ? connection.srvPrev->srvNext : connection.server->connFirst) = &connection;

				}

				connection.sockId = sockId;

				struct epoll_event event = {};

				event.data.ptr = &connection;

				if (connected) {

					epoll_ctl(epollFd, EPOLL_CTL_ADD, sockId, &event);

					connection.onOpen(connection.address, connection.addressSize);

					connection.state = connection.stateLocal = Connection::STATE_OPEN;

					connection.stateProtector.wait();

				} else {

					event.events = EPOLLOUT;

					epoll_ctl(epollFd, EPOLL_CTL_ADD, sockId, &event);

					connection.state = connection.stateLocal = Connection::STATE_CONNECTING;

					connection.stateProtector.wait();

				}

			} else {

				connection.sockId = -1;

				connection.server = nullptr;

				connection.onOpeningError();

				connection.state = Connection::STATE_CLOSED;

				while(connection.closeWaiters > 0)

					{ connection.closeCondition.notify_all(); std::this_thread::yield(); }

				connection.state = Connection::STATE_FINISHED;

			}

		} else {

			Connection::State stateWanted = connection.stateWanted;

			if ( connection.stateLocal >= Connection::STATE_CONNECTING

			  && connection.stateLocal <= Connection::STATE_OPEN

			  && stateWanted == Connection::STATE_CLOSE_REQ )

			{

				connection.state = connection.stateLocal = Connection::STATE_CLOSE_REQ;

				connection.stateProtector.wait();

				connection.onCloseReqested();

			} else

			if ( connection.stateLocal >= Connection::STATE_CONNECTING

			  && connection.stateLocal <= Connection::STATE_CLOSE_REQ

			  && stateWanted == Connection::STATE_CLOSING )

			{

				connection.state = connection.stateLocal = Connection::STATE_CLOSING;

				connection.stateProtector.wait();

				while(Connection::Task *task = connection.readQueue.peek()) {

					connection.readQueue.pop();

					if (!connection.onReadReady(*task, true))

						assert(false);

				}

				while(Connection::Task *task = connection.writeQueue.peek()) {

					connection.writeQueue.pop();

					if (!connection.onWriteReady(*task, true))

						assert(false);

				}

				connection.onClose(false);

			}

			if (connection.stateLocal == Connection::STATE_CLOSING && !connection.enqueued) {

				if (connection.server) {

					(connection.srvPrev ? connection.srvPrev->srvNext : connection.server->connFirst) = connection.srvNext;

					(connection.srvNext ? connection.srvNext->srvPrev : connection.server->connLast) = connection.srvPrev;

					connection.server = nullptr;

				}

				(connection.prev ? connection.prev->next : connFirst) = connection.next;

				(connection.next ? connection.next->prev : connLast) = connection.prev;

				connection.sockId = -1;

				connection.protocol = nullptr;

				memset(connection.address, 0, sizeof(connection.address));

				connection.addressSize = 0;

				connection.stateWanted = Connection::STATE_NONE;

				connection.stateLocal = Connection::STATE_NONE;

				connection.state = Connection::STATE_CLOSED;

				while(connection.closeWaiters > 0)

					{ connection.closeCondition.notify_all(); std::this_thread::yield(); }

				connection.state = Connection::STATE_FINISHED;

			}

		}

	}

}

void Protocol::threadConnEvents(Connection &connection, unsigned int events) {

	bool needUpdate = false;

	if (connection.stateLocal == Connection::STATE_CONNECTING) {

		if (events & EPOLLHUP) {

			connection.onOpeningError();

			if (connection.server) {

				(connection.srvPrev ? connection.srvPrev->srvNext : connection.server->connFirst) = connection.srvNext;

				(connection.srvNext ? connection.srvNext->srvPrev : connection.server->connLast) = connection.srvPrev;

				connection.server = nullptr;

			}

			(connection.prev ? connection.prev->next : connFirst) = connection.next;

			(connection.next ? connection.next->prev : connLast) = connection.prev;

			connection.sockId = -1;

			connection.protocol = nullptr;

			memset(connection.address, 0, sizeof(connection.address));

			connection.addressSize = 0;

			connection.stateWanted = Connection::STATE_NONE;

			connection.stateLocal = Connection::STATE_NONE;

			connection.state = Connection::STATE_CLOSED;

			while(connection.closeWaiters > 0)

				{ connection.closeCondition.notify_all(); std::this_thread::yield(); }

			connection.state = Connection::STATE_FINISHED;

		}

		if (events & EPOLLOUT) {

			connection.onOpen(connection.address, connection.addressSize);

			connection.state = connection.stateLocal = Connection::STATE_OPEN;

			connection.stateProtector.wait();

			needUpdate = true;

		}

	}

	if ( connection.stateLocal < Connection::STATE_OPEN

	  || connection.stateLocal >= Connection::STATE_CLOSING )

		return;

	if (events & EPOLLIN) {

		bool ready = true;

		int pops = 0;

		while(ready) {

			Connection::Task *task = connection.readQueue.peek();

			if (!task) {

				if (pops) needUpdate = true;

				break;

			}

			size_t completion = 0;

			while(task->size > task->completion) {

				size_t size = task->size - task->completion;

				if (size > INT_MAX) size = INT_MAX;

				int res = ::recv(

					connection.sockId,

					(unsigned char*)task->data + task->completion,

					(int)size, MSG_DONTWAIT );

				if (res == EAGAIN) {

					ready = false;

					break;

				} else

				if (res <= 0) {

					connection.close(true);

					return;

				} else {

					completion += size;

					task->completion += size;

				}

			}

			if (task->size <= task->completion || (completion && task->watch)) {

				connection.readQueue.pop();

				if (!connection.onReadReady(*task, false))

					connection.readQueue.unpop(*task); else ++pops;

			}

		}

	}

	if (events & EPOLLHUP) {

		connection.close((bool)(events & EPOLLERR));

		return;

	}

	if (events & EPOLLOUT) {

		bool ready = true;

		int pops = 0;

		while(ready) {

			Connection::Task *task = connection.writeQueue.peek();

			if (!task) {

				if (pops) needUpdate = true;

				break;

			}

			size_t completion = 0;

			while(task->size > task->completion) {

				size_t size = task->size - task->completion;

				if (size > INT_MAX) size = INT_MAX;

				int res = ::send(

					connection.sockId,

					(const unsigned char*)task->data + task->completion,

					(int)size, MSG_DONTWAIT );

				if (res == EAGAIN) {

					ready = false;

					break;

				} else

				if (res <= 0) {

					connection.close(true);

					return;

				} else {

					completion += size;

					task->completion += size;

				}

			}

			if (task->size <= task->completion || (completion && task->watch)) {

				connection.writeQueue.pop();

				if (!connection.onWriteReady(*task, false))

					connection.writeQueue.unpop(*task); else ++pops;

			}

		}

	}

	if (needUpdate) updateEvents(connection);

}

void Protocol::threadSrvQueue() {

	while(srvQueue.peek()) {

		Server &server = srvQueue.pop();

		server.enqueued = false;

		if (server.stateLocal == Server::STATE_NONE) {

			int sockId = -1;

			if (server.addressSize >= sizeof(sa_family_t)) {

				struct sockaddr *addr = (struct sockaddr*)server.address;

				sockId = ::socket(addr->sa_family, SOCK_STREAM, 0);

				if (sockId >= 0) {

					initSocket(sockId);

					if ( 0 != ::bind(sockId, addr, (int)server.addressSize)

					  || 0 != ::listen(sockId, 128) )

					{

						::close(sockId);

						sockId = -1;

					}

				}

			}

			if (sockId >= 0) {

				server.protocol = this;

				server.prev = srvLast;

				assert(!server.next);

				(server.prev ? server.prev->next : srvFirst) = &server;

				server.sockId = sockId;

				struct epoll_event event = {};

				event.data.ptr = &server;

				event.events = EPOLLIN;

				epoll_ctl(epollFd, EPOLL_CTL_ADD, sockId, &event);

				server.onOpen();

				server.state = server.stateLocal = Server::STATE_OPEN;

				server.stateProtector.wait();

			} else {

				server.onOpeningError();

				server.state = Server::STATE_CLOSED;

				while(server.closeWaiters > 0)

					{ server.closeCondition.notify_all(); std::this_thread::yield(); }

				server.state = Server::STATE_FINISHED;

			}

		} else {

			Server::State stateWanted = server.stateWanted;

			if ( server.stateLocal == Server::STATE_OPEN

			  && stateWanted == Server::STATE_CLOSE_REQ )

			{

				for(Connection *conn = server.connFirst; conn; conn = conn->srvNext)

					conn->closeReq();

				struct epoll_event event = {};

				epoll_ctl(epollFd, EPOLL_CTL_DEL, server.sockId, &event);

				server.state = server.stateLocal = Server::STATE_CLOSE_REQ;

				server.stateProtector.wait();

				server.onCloseReqested();

			} else

			if ( ( server.stateLocal == Server::STATE_OPEN

				|| server.stateLocal == Server::STATE_CLOSE_REQ )

			  && stateWanted == Server::STATE_CLOSING_CONNECTIONS )

			{

				if (server.stateLocal != Server::STATE_CLOSE_REQ) {

					struct epoll_event event = {};

					epoll_ctl(epollFd, EPOLL_CTL_DEL, server.sockId, &event);

				}

				if (server.connFirst) {

					for(Connection *conn = server.connFirst; conn; conn = conn->srvNext)

						conn->close(true);

					server.state = server.stateLocal = Server::STATE_CLOSING_CONNECTIONS;

				} else {

					server.state = server.stateLocal = Server::STATE_CLOSING;

					server.stateProtector.wait();

					server.onClose(false);

				}

			}

			if (server.stateLocal == Server::STATE_CLOSING && !server.enqueued) {

				server.state = Server::STATE_CLOSED;

				(server.prev ? server.prev->next : srvFirst) = server.next;

				(server.next ? server.next->prev : srvLast) = server.prev;

				server.protocol = nullptr;

				memset(server.address, 0, sizeof(server.address));

				server.addressSize = 0;

				server.stateWanted = Server::STATE_NONE;

				server.stateLocal = Server::STATE_NONE;

				while(server.closeWaiters > 0)

					{ server.closeCondition.notify_all(); std::this_thread::yield(); }

				server.state = Server::STATE_FINISHED;

			}

		}

	}

}

void Protocol::threadSrvEvents(Server &server, unsigned int events) {

	if (events & EPOLLHUP) {

		if ( server.stateLocal >= Server::STATE_OPEN

		  && server.stateLocal < Server::STATE_CLOSING )

			server.close((bool)(events & EPOLLERR));

		return;

	}

	if (events & EPOLLIN) {

		if ( server.stateLocal >= Server::STATE_OPEN

		  && server.stateLocal < Server::STATE_CLOSE_REQ)

		{

			unsigned char address[MAX_ADDR_SIZE] = {};

			socklen_t len = sizeof(address);

			int res = ::accept(server.sockId, (struct sockaddr*)address, &len);

			assert(len < sizeof(address));

			if (res >= 0) {

				if (Connection *conn = server.onConnect(address, len))

					if (!conn->open(*this, address, len, res))

						server.onDisconnect(conn, true);

			} else

			if (res != EAGAIN) {

				server.close(true);

				return;

			}

		}

	}

}

void Protocol::threadRun() {

	epollFd = ::epoll_create(32);

	eventFd = ::eventfd(0, EFD_NONBLOCK);

	struct epoll_event event = {};

	event.events = EPOLLIN;

	epoll_ctl(epollFd, EPOLL_CTL_ADD, eventFd, &event);

	const int maxCount = 16;

	struct epoll_event events[maxCount] = {};

	int count = 0;

	State stateLocal = STATE_OPEN;

	state = stateLocal;

	while(true) {

		for(int i = 0; i < count; ++i) {

			Root *pointer = (Root*)events[i].data.ptr;

			if (!pointer) {

				unsigned long long num = 0;

				while(::read(eventFd, &num, sizeof(num)) > 0 && !num);

			} else

			if (Connection *connection = dynamic_cast<connection*>(pointer)) {</connection*>

				threadConnEvents(*connection, events[i].events);

			} else

			if (Server *server = dynamic_cast<server*>(pointer)) {</server*>

				threadSrvEvents(*server, events[i].events);

			}

		}

		State sw = stateWanted;

		if (stateLocal == STATE_OPEN && sw == STATE_CLOSE_REQ) {

			state = stateLocal = STATE_CLOSE_REQ;

			stateProtector.wait();

			for(Server *server = srvFirst; server; server = server->next)

				server->closeReq();

			for(Connection *conn = connFirst; conn; conn = conn->next)

				if (!conn->server)

					conn->closeReq();

		} else

		if (stateLocal >= STATE_OPEN && stateLocal <= STATE_CLOSE_REQ && sw == STATE_CLOSING) {

			state = STATE_CLOSING;

			stateProtector.wait();

			for(Server *server = srvFirst; server; server = server->next)

				server->close(true);

			for(Connection *conn = connFirst; conn; conn = conn->next)

				if (!conn->server)

					conn->close(true);

		}

		while(connQueue.peek() || srvQueue.peek()) {

			threadConnQueue();

			threadSrvQueue();

		}

		if (stateLocal == STATE_CLOSING) {

			assert(!srvFirst && !connFirst);

			break;

		}

		count = epoll_wait(epollFd, events, maxCount, 10000);

	}

	while(closeWaiters > 0)

		{ closeCondition.notify_all(); std::this_thread::yield(); }

	::close(epollFd);

	::close(eventFd);

	state = STATE_FINISHED;

}