#pragma once

#include <stdio.h>
#include <errno.h>
#include <sys/types.h>
#include <fcntl.h>

#ifdef _WIN32
#include <winsock2.h>
#include <WS2tcpip.h>
#include <stdlib.h>

#define close(fd) closesocket(fd)
#define ssize_t int
#else
#include <sys/time.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <poll.h>
#endif

#ifndef constrain
#define constrain(x, l, u) (x<l?l:(x>u?u:x))
#endif

class TCPHelperClient
{
#ifdef _WIN32
	SOCKET _fd = -1;
#else
	int _fd = -1;
#endif
	sockaddr_in _remote_addr;

public:
	TCPHelperClient() {}
	TCPHelperClient(int fd)
	{
		_fd = fd;
	}
	TCPHelperClient(const TCPHelperClient &another)
	{
		*this = another;
	}

	/*
	  ~TCPHelperClient() {
	#ifdef _WIN32
	    if (_fd != INVALID_SOCKET) closesocket(_fd);
	    WSACleanup();
	#else
	    if (_fd != INVALID_SOCKET) close(_fd);
	#endif
	  }
	*/
	int available()
	{
#ifdef _WIN32
		WSAPOLLFD pfs;
		pfs.fd = _fd;
		pfs.events = POLLIN;
		pfs.revents = 0;
		int rv = WSAPoll(&pfs, 1, 0);
		if ((rv == -1 && WSAGetLastError() == WSAENETDOWN) || (rv > 0 && (pfs.revents & POLLNVAL) != 0)) {
#ifdef ETCP_ERROR_PRINT
			printf("poll triggered stop, rv=%d: %s\n", rv, strerror(errno));
#endif
			stop(); // Socket problem, close it
		} else if (rv > 0) {
			return (pfs.revents & POLLRDNORM) != 0 || (pfs.revents & POLLHUP) != 0;
		}
#else
		struct pollfd pfs;
		pfs.fd = _fd;
		pfs.events = POLLIN;
		pfs.revents = 0;
		int rv = ::poll(&pfs, 1, 0);
		if (rv == -1 || (rv > 0 && (pfs.revents & POLLNVAL) != 0)) {
#ifdef ETCP_ERROR_PRINT
			printf("poll triggered stop, rv=%d: %s\n", rv, strerror(errno));
#endif
			stop(); // Socket problem, close it
		} else if (rv > 0) {
			return (pfs.revents & POLLIN) != 0 || (pfs.revents & POLLHUP) != 0;
		}
#endif
		return 0;
	}

	bool connect(const uint8_t *address, uint16_t port)
	{
		if (_fd != -1) {
#ifdef ETCP_ERROR_PRINT
			printf("connect triggered stop, _fd=%d\n", _fd);
#endif
			stop();
		}
		_fd = ::socket(AF_INET,SOCK_STREAM,0);
		if (_fd == -1) {
			return false;
		}
		memset(&_remote_addr, 0, sizeof(_remote_addr));
		_remote_addr.sin_family = AF_INET;
		_remote_addr.sin_port = htons(port);
		memcpy(&_remote_addr.sin_addr.s_addr, address, 4);

		// Shorten timeout for connecting
		struct timeval read_timeout;
		read_timeout.tv_sec = 0;
		read_timeout.tv_usec = 2000000;
#ifndef __ZEPHYR__
		setsockopt(_fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&read_timeout, sizeof read_timeout);
		setsockopt(_fd, SOL_SOCKET, SO_SNDTIMEO, (char*)&read_timeout, sizeof read_timeout);
#endif

		bool connected = ::connect(_fd, (struct sockaddr *) &_remote_addr,sizeof(_remote_addr)) == 0;
		if (connected) {
#ifndef __ZEPHYR__
			// Shorten timeouts for reading and writing
			struct timeval read_timeout;
			read_timeout.tv_sec = 0;
			read_timeout.tv_usec = 1000;
			setsockopt(_fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&read_timeout, sizeof read_timeout);
			read_timeout.tv_usec = 2000000;
			setsockopt(_fd, SOL_SOCKET, SO_SNDTIMEO, (char*)&read_timeout, sizeof read_timeout);
#endif
			// Disable Nagles algorith because we are sending small packets and waiting for reply
			set_nodelay(true);
		}
		return connected;
	}

	// Prepare a socket for non-blocking connect
	bool prepare_connect(const uint8_t *address, uint16_t port)
	{
		if (_fd != -1) {
#ifdef ETCP_ERROR_PRINT
			printf("prepare_connect triggered stop, _fd=%d\n", _fd);
			stop();
#endif
		}
		_fd = ::socket(AF_INET,SOCK_STREAM,0);
		if (_fd == -1) {
			return false;
		}
		memset(&_remote_addr, 0, sizeof(_remote_addr));
		_remote_addr.sin_family = AF_INET;
		_remote_addr.sin_port = htons(port);
		memcpy(&_remote_addr.sin_addr.s_addr, address, 4);
#ifdef _WIN32
		unsigned long ul = 1;
		int flags = ioctlsocket(_fd, FIONBIO, (unsigned long *)&ul); //Set into non blocking mode.
		if (flags != SOCKET_ERROR) ; // Failed to set?
#else
		int flags = fcntl(_fd, F_GETFL, 0);
		if (flags!=-1) {
			fcntl(_fd, F_SETFL, flags | O_NONBLOCK);
		}
#endif
		else {
#ifdef ETCP_ERROR_PRINT
			printf("prepare_connect triggered stop, _fd=%d: %s\n", _fd, strerror(errno));
#endif
		}
		return flags != -1;
	}

	// Try a non-blocking connect, it may succeed at once or after some calls
	// Returns -1 if failure, 0 if waiting to connect, 1 when connected
	int8_t try_connect()
	{
		int8_t connected = 0;
		int status = ::connect(_fd, (struct sockaddr *) &_remote_addr,sizeof(_remote_addr));
		if (status == -1) {
#ifdef _WIN32
			int errno2 = WSAGetLastError();
			if (errno2 == WSAEISCONN); // Already connected, proceed with setting timeouts
			else if (errno2 == WSAEWOULDBLOCK || errno2 == WSAEINPROGRESS || errno2 == WSAEALREADY) {
				return 0;    // Normal, we are waiting for connect to succeed
			}
#else
			if (errno == EISCONN); // Already connected, proceed with setting timeouts
			else if (errno == EINPROGRESS || errno == EALREADY) {
				return 0;    // Normal, we are waiting for connect to succeed
			}
#endif
			else {
#ifdef ETCP_ERROR_PRINT
#ifdef _WIN32
				printf("try_connect triggered stop, _fd=%d: %d\n", _fd, errno2);
#else
				printf("try_connect triggered stop, _fd=%d: %d %s\n", _fd, errno, strerror(errno));
#endif
#endif
				stop();
				return -1; // Some other error
			}
		}

		// Enable blocking operations
#ifdef _WIN32
		unsigned long ul = 0;
		int flags = ioctlsocket(_fd, FIONBIO, (unsigned long *)&ul); // Set into non blocking mode.
		//if (flags != SOCKET_ERROR) ; // Failed to set.
#else
		int flags = fcntl(_fd, F_GETFL, 0);
		if (flags != -1) {
			fcntl(_fd, F_SETFL, flags & ~O_NONBLOCK);
		}
#endif

		// Set timeouts for reading and writing
#ifndef __ZEPHYR__
		struct timeval read_timeout;
		read_timeout.tv_sec = 0;
		read_timeout.tv_usec = 1000;
		setsockopt(_fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&read_timeout, sizeof read_timeout);
		read_timeout.tv_usec = 2000000;
		setsockopt(_fd, SOL_SOCKET, SO_SNDTIMEO, (char*)&read_timeout, sizeof read_timeout);
#endif
		// Disable Nagles algorith because we are sending small packets and waiting for reply
		set_nodelay(true);

		return 1;
	}

	// Enable or disable Nagles algorithm to send unfilled packets immediately or not
	void set_nodelay(bool nodelay)
	{
		int flag = nodelay;
		if (_fd != -1) {
			setsockopt(_fd, IPPROTO_TCP, TCP_NODELAY, (char*)&flag, sizeof flag);
		}
	}

	bool connected()
	{
		return _fd != -1;
	}

	int read(uint8_t *buffer, int buffer_size)
	{
		if (_fd == -1) {
			return -1;
		}
		int r = ::recv(_fd, (char*)buffer, buffer_size, 0); //MSG_DONTWAIT);
		if (r == -1) {
#ifdef _WIN32
			int errno2 = WSAGetLastError();
			if (errno2 == WSAEWOULDBLOCK) {
				return 0;    // Waiting for more
			}
#ifdef ETCP_ERROR_PRINT
			printf("read triggered stop, r=%d: %d\n", r, errno2);
#endif
#else
			if (errno == EAGAIN) {
				return 0;    // Waiting for more
			}
#ifdef ETCP_ERROR_PRINT
			printf("read triggered stop, r=%d: %s\n", r, strerror(errno));
#endif
#endif
			stop();
		}
		return r;
	}

	int write(const uint8_t *buffer, int size)
	{
		if (_fd == -1) {
			return -1;
		}
#if defined(_WIN32) || defined(__ZEPHYR__)
		int w = ::send(_fd, (char*)buffer, size, 0);
#else
		int w = ::send(_fd, (char*)buffer, size, MSG_NOSIGNAL);
#endif
		if (w == -1) {
#ifdef ETCP_ERROR_PRINT
#ifdef _WIN32
			int errno2 = WSAGetLastError();
			printf("write triggered stop, w=%d: %d\n", w, errno2);
#else
			if (errno != EPIPE) {
				printf("write triggered stop, w=%d: %s\n", w, strerror(errno));
			}
#endif
#endif
			stop();
		}
		return w;
	}

	void flush() { }

	void stop()
	{
		if (_fd != -1) {
#ifndef _WIN32
			::shutdown(_fd, SHUT_RDWR);
#endif
			::close(_fd);
			_fd = -1;
		}
	}

	void operator=(const TCPHelperClient &another)
	{
		_fd = another._fd;
	}

	operator bool()
	{
		return connected();
	}
	bool operator==(const bool value)
	{
		return bool() == value;
	}
	bool operator!=(const bool value)
	{
		return bool() != value;
	}
	bool operator==(const TCPHelperClient& rhs)
	{
		return _fd == rhs._fd && _fd != -1 && rhs._fd != -1;
	}
	bool operator!=(const TCPHelperClient& rhs)
	{
		return !this->operator==(rhs);
	}
	uint8_t getSocketNumber()
	{
		return _fd;
	}

	int print(const char *msg)
	{
		return write((const uint8_t*) msg, strlen(msg));
	}
};


class TCPHelperServer
{
	uint16_t _port = 0;
	int _fd = -1;
	sockaddr_in _localaddr, _remote_sender_addr;

public:
	TCPHelperServer(uint16_t listening_port)
	{
		_port = listening_port;
	}
	~TCPHelperServer()
	{
#ifdef ETCP_DEBUG_PRINT
		printf("destructor triggered stop, _fd=%d\n", _fd);
#endif
		stop();
	}

	TCPHelperClient available()
	{
		socklen_t len = sizeof(_remote_sender_addr);
		memset(&_remote_sender_addr, 0, len);
		int connected_fd = _fd == -1 ? -1 : ::accept(_fd, (struct sockaddr *) &_remote_sender_addr, &len);
		if (connected_fd != -1) {
#ifndef __ZEPHYR__
			// Shorten timeout for reading and writing
			struct timeval read_timeout;
			read_timeout.tv_sec = 0;
			read_timeout.tv_usec = 1000;
			setsockopt(connected_fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&read_timeout, sizeof read_timeout);
			read_timeout.tv_usec = 2000000;
			setsockopt(connected_fd, SOL_SOCKET, SO_SNDTIMEO, (char*)&read_timeout, sizeof read_timeout);
#endif
			// Disable Nagles algorith because we are sending small packets and waiting for reply
			int flag = 1;
			setsockopt(connected_fd, IPPROTO_TCP, TCP_NODELAY, (char*)&flag, sizeof flag);
		}
		return TCPHelperClient(connected_fd);
	}

	bool begin()
	{
		// Close if open after previous init attempt
		if (_fd != -1) {
			close(_fd);
			_fd = -1;
		}

		// Prepare socket
		_fd=socket(AF_INET,SOCK_STREAM,0);
		if (_fd==-1) {
#ifdef ETCP_ERROR_PRINT
			printf("failure creating socket: %s\n", strerror(errno));
#endif
			return false;
		}

		int yes = 1;
		setsockopt(_fd, SOL_SOCKET, SO_REUSEADDR, (char*)&yes, sizeof(yes));

		// Bind to specific local port
		memset(&_localaddr, 0, sizeof(_localaddr));
		_localaddr.sin_family = AF_INET;
		_localaddr.sin_port = htons(_port);
		_localaddr.sin_addr.s_addr = htonl(INADDR_ANY);
		if (bind(_fd,(struct sockaddr *) &_localaddr,sizeof(_localaddr))==-1) {
#ifdef ETCP_ERROR_PRINT
			printf("bind failed: %s\n", strerror(errno));
#endif
			return false;
		}

		// Start listening
		if (listen(_fd, 5) == -1) {
#ifdef ETCP_ERROR_PRINT
			printf("listen failed: %s\n", strerror(errno));
#endif
			return false;
		}

		// Set the socket to nonblocking so that accept will not block
#ifdef _WIN32
		unsigned long ul = 1;
		int flags = ioctlsocket(_fd, FIONBIO, (unsigned long *)&ul); //Set into non blocking mode.
		if (flags != NO_ERROR) { // Failed to set.
#ifdef ETCP_ERROR_PRINT
			printf("ioctlsocket FIONBIO failed.");
#endif
		}
#else
		int flags = fcntl(_fd, F_GETFL, 0);
		if (flags != -1) {
			fcntl(_fd, F_SETFL, flags | O_NONBLOCK);
		}
#endif
		memset(&_remote_sender_addr, 0, sizeof(_remote_sender_addr));
		return true;
	}

	void stop()
	{
		if (_fd != -1) {
			::close(_fd);
			_fd = -1;
		}
	}
};

#undef close
