#include <cerrno></cerrno>

#include <cassert></cassert>

#include <cstring></cstring>

#include <climits></climits>

#include <cctype></cctype>

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

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

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

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

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

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

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

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

#include "tcp.h"

static const int pollTimeoutMs = 100;

static const int connTimeoutMs = 300000;

using namespace Tcp;

ErrorCode Address::resolve(const char *addrString) {

	*this = Address();

	if (!addrString)

		return ERR_TCP_BAD_ADDRESS;

	ErrorCode errorCode = ERR_NONE;

	const char *colon = addrString;

	while(*colon && *colon != ':') ++colon;

	char *hostBuf = nullptr;

	const char *host = nullptr;

	const char *port = nullptr;

	if (*colon) {

		hostBuf = new char[colon - addrString + 1];

		memcpy(hostBuf, addrString, colon - addrString);

		hostBuf[colon - addrString] = 0;

		host = hostBuf;

		port = colon + 1;

	} else {

		bool allDigits = true;

		for(const char *c = addrString; *c && allDigits; ++c)

			if (*c < '0' || *c > '9') allDigits = false;

		(allDigits ? port : host) = addrString;

	}

	if (host) {

		if (4 != sscanf(host, "%hhu.%hhu.%hhu.%hhu", &ip[0], &ip[1], &ip[2], &ip[3])) {

			memset(ip, 0, sizeof(ip));

			hostent *he;

			in_addr **addr_list;

			if ( (he = gethostbyname(host))

			  && (addr_list = (struct in_addr**)he->h_addr_list)

			  && (*addr_list) )

			{

				memcpy(ip, *addr_list, sizeof(ip));

			} else {

				errorCode = ERR_TCP_BAD_ADDRESS;

			}

		}

	}

	if (port) {

		if (1 != sscanf(port, "%hu", &this->port))

			errorCode = ERR_TCP_BAD_ADDRESS;

	}

	if (!host && !port)

		errorCode = ERR_TCP_BAD_ADDRESS;

	if (hostBuf) delete[] hostBuf;

	return errorCode;

}

Connection::Connection():

	status(STATUS_NONE),

	lastError(ERR_NONE),

	sockId(),

	stopping(false)

	{ }

Connection::Connection(SockId sockId, const Address &remoteAddr):

	status(STATUS_OPEN),

	lastError(ERR_NONE),

	sockId(sockId),

	remoteAddr(remoteAddr),

	stopping(false)

	{ }

Connection::~Connection()

	{ close(); }

bool Connection::open(const Address &remoteAddr) {

	close();

	this->remoteAddr = remoteAddr;

	if (!this->remoteAddr.isValid())

		{ close(ERR_TCP_BAD_ADDRESS); return false; }

	sockId = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

	status = STATUS_OPEN;

	struct sockaddr_in addr = {};

	addr.sin_family = AF_INET;

	memcpy(&addr.sin_addr, remoteAddr.ip, sizeof(remoteAddr.ip));

	addr.sin_port = htons(remoteAddr.port);

	if (0 != connect(sockId, (sockaddr*)&addr, sizeof(addr)))

		{ close(ERR_TCP_CONNECTION_FAILED); return false; }

	int tcp_nodelay = 1;

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

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

	return true;

}

void Connection::close(ErrorCode errorCode) {

	if (status == STATUS_OPEN)

		::close(sockId);

	sockId = 0;

	status = errorCode ? STATUS_LOST : STATUS_NONE;

	lastError = errorCode;

	stopping = false;

}

bool Connection::read(void *data, size_t size) {

	if (!data)

		return check() && size == 0;

	struct pollfd pfd = {};

	pfd.fd = sockId;

	pfd.events = POLLIN;

	int remainMs = connTimeoutMs;

	while(check() && size && remainMs >= 0) {

		int pr = poll(&pfd, 1, pollTimeoutMs);

		if (pr < 0)

			{ close(ERR_TCP_CONNECTION_LOST); return false; }

		if (pr > 0) {

			int s = size > INT_MAX ? INT_MAX : (int)size;

			int r = ::recv(sockId, data, s, MSG_DONTWAIT | MSG_NOSIGNAL);

			if (r < 0) {

				if (errno != EAGAIN)

					{ close(ERR_TCP_CONNECTION_LOST); return false; }

			} else {

				size -= r;

				data = (char*)data + r;

			}

		}

		remainMs -= pollTimeoutMs;

	}

	return size == 0;

}

bool Connection::write(const void *data, size_t size) {

	if (!data)

		return check() && size == 0;

	struct pollfd pfd = {};

	pfd.fd = sockId;

	pfd.events = POLLOUT;

	int remainMs = connTimeoutMs;

	while(check() && size && remainMs >= 0) {

		int pr = poll(&pfd, 1, pollTimeoutMs);

		if (pr < 0)

			{ close(ERR_TCP_CONNECTION_LOST); return false; }

		if (pr > 0) {

			int s = size > INT_MAX ? INT_MAX : (int)size;

			int r = ::send(sockId, data, s, MSG_DONTWAIT | MSG_NOSIGNAL);

			if (r < 0) {

				if (errno != EAGAIN)

					{ close(ERR_TCP_CONNECTION_LOST); return false; }

			} else {

				size -= r;

				data = (const char*)data + r;

			}

		}

		remainMs -= pollTimeoutMs;

	}

	return size == 0;

}

Handler::~Handler() { }

Server::Server():

	status(STATUS_NONE),

	lastError(ERR_NONE),

	sockId(),

	localAddr(),

	handler(),

	thread(),

	stopping(false)

	{ }

Server::~Server()

	{ stop(); }

void Server::handleConnection(ConnList::iterator iter) {

	handler->handleTcpConnection(*iter->connection);

	{

		std::lock_guard<std::mutex> lock(connectionsMutex);</std::mutex>

		connectionsFinished.push_back(*iter);

		connections.erase(iter);

	}

}

void Server::cleanConnections() {

	while(!connectionsFinished.empty()) {

		ConnDesc &desc = connectionsFinished.back();

		desc.thread->join();

		delete desc.thread;

		delete desc.connection;

		connectionsFinished.pop_back();

	}

}

void Server::listen() {

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

	struct sockaddr_in addr = {};

	Address remoteAddr;

	struct pollfd pfd = {};

	pfd.fd = sockId;

	pfd.events = POLLIN;

	status = STATUS_OPEN;

	while(!stopping) {

		int pr = poll(&pfd, 1, pollTimeoutMs);

		if (pr < 0)

			{ lastError = ERR_TCP_LISTEN_LOST; break; }

		if (pr > 0 && (pr & POLLIN)) {

			socklen_t addrlen = sizeof(addr);

			int sid = ::accept(sockId, (sockaddr*)&addr, &addrlen);

			if (sid < 0) {

				if (errno != EAGAIN)

					{ lastError = ERR_TCP_LISTEN_LOST; break; }

				continue;

			}

			memcpy(remoteAddr.ip, &addr.sin_addr, sizeof(remoteAddr.ip));

			remoteAddr.port = ntohs(addr.sin_port);

			ConnList::iterator iter;

			{

				std::lock_guard<std::mutex> lock(connectionsMutex);</std::mutex>

				cleanConnections();

				iter = connections.emplace(connections.end());

			}

			iter->connection = new Connection(sid, remoteAddr);

			iter->thread = new std::thread(&Server::handleConnection, this, iter);

		} else {

			std::lock_guard<std::mutex> lock(connectionsMutex);</std::mutex>

			cleanConnections();

		}

	}

	status = lastError ? STATUS_LOST : STATUS_CLOSED;

	connectionsMutex.lock();

	for(ConnList::iterator i = connections.begin(); i != connections.end(); ++i)

		i->connection->stop();

	while(!connections.empty()) {

		std::thread *thread = connections.front().thread;

		connectionsMutex.unlock();

		thread->join();

		connectionsMutex.lock();

	}

	cleanConnections();

	connectionsMutex.unlock();

	stopping = false;

}

bool Server::start(const Address &localAddr, Handler *handler) {

	stop();

	this->localAddr = localAddr;

	this->handler = handler;

	if (!this->localAddr.isValidPort())

		{ status = STATUS_LOST; lastError = ERR_TCP_BAD_ADDRESS; return false; }

	if (!this->handler)

		{ status = STATUS_LOST; lastError = ERR_TCP_BAD_HANDLER; return false; }

	sockId = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

	struct sockaddr_in addr = {};

	addr.sin_family = AF_INET;

	memcpy(&addr.sin_addr, localAddr.ip, sizeof(localAddr.ip));

	addr.sin_port = htons(localAddr.port);

	if ( 0 != ::bind(sockId, (sockaddr*)&addr, sizeof(addr))

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

	{

		status = STATUS_LOST;

		lastError = ERR_TCP_LISTEN_FAILED;

		this->handler = nullptr;

		return false;

	}

	int tcp_nodelay = 1;

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

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

	thread = new std::thread(&Server::listen, this);

	return true;

}

void Server::join() {

	if (!thread)

		return;

	thread->join();

	delete thread;

	thread = nullptr;

	handler = nullptr;

}

void Server::stop() {

	if (!thread)

		return;

	stopping = true;

	join();

}