#include <stdio.h>
#include <stdbool.h>
#include <winsock.h>

#define ERRMSG_BUFFER_SIZE  100  // Size of error message buffer
#define DAYTIME_PORT        13   // Port of daytime service (RFC 867)
#define BUFFER_SIZE         100  // Size of buffer for daytime string

/* Name:         getErrorMsg
 * Desription:   Returns a human readable error message for codes from
 *               WSAGetLastError or GetLastError
 * Parameters:   code - error code to return message for
 * Return code:  Pointer to a static buffer containing the message
 * Notes:        Not thread-safe because of static buffer
 */
const char *getErrorMsg(int code)
{
	static char buffer[ERRMSG_BUFFER_SIZE];

	FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, code, 0, buffer, sizeof(buffer), NULL);

	return buffer;
}

/* Name:         initWinsock
 * Desription:   Initialized Winsock subsystem
 * Parameters:   none
 * Return code:  true if successful
 *               false if an error occurred
 */
bool initWinsock(void)
{
	bool failed = false;

	WSADATA wsa;
	int startupResult = WSAStartup(MAKEWORD(1, 1), &wsa);
	if (startupResult != 0)
	{
		fprintf(stderr, "Couldn't start Winsock: %s\n", getErrorMsg(startupResult));
		failed = true;
	}

	return !failed;
}

/* Name:         deinitWinsock
 * Desription:   Deinitialized Winsock subsystem
 * Parameters:   none
 * Return code:  none
 */
void deinitWinsock(void)
{
	WSACleanup();
}

/* Name:         getCmdLineArgs
 * Description:  Gets the expected arguments from command line
 * Parameters:   argc - number of arguments specified
 *               argv - array with arguments as strings
 *               server - hostname of server to connect to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool getCmdLineArgs(int argc, char *argv[], const char **server)
{
	bool result = false;

	if (argc == 2)
	{
		*server = argv[1];
		result = true;
	}
	else
	{
		fprintf(stderr, "usage: %s <server>\n", argv[0]);
	}

	return result;
}

/* Name:         resolveHostName
 * Description:  Resolves a hostname into an IP address
 * Parameters:   name - hostname to resolve
 *               addr - buffer to store IP address in
 * Return code:  true if successful
 *               false if an error occurred
 */
bool resolveHostName(const char *name, struct in_addr *addr)
{
	bool result = false;

	printf("Resolving hostname \"%s\"...", name);
	struct hostent *host = gethostbyname(name);
	if (host)
	{
		printf(" OK.\n");
		*addr = *(struct in_addr*)host->h_addr;
		result = true;
	}
	else
	{
		printf(" FAILED.\n");
		fprintf(stderr, "Couldn't resolve hostname: %s\n", getErrorMsg(WSAGetLastError()));
	}

	return result;
}

/* Name:         createTcpSocket
 * Description:  Creates a new, unconnected TCP socket
 * Parameters:   sock - buffer to receive socket
 * Return code:  true if successful
 *               false if an error occurred
 */
bool createTcpSocket(SOCKET *sock)
{
	bool failed = false;

	*sock = socket(AF_INET, SOCK_STREAM, 0);
	if (*sock == INVALID_SOCKET)
	{
		fprintf(stderr, "Couldn't create socket: %s\n", getErrorMsg(WSAGetLastError()));
		failed = true;
	}

	return !failed;
}

/* Name:         connectTcpSocket
 * Description:  Connects a TCP socket to a destination end point
 * Parameters:   sock - socket to connect
 *               addr - destination address to connect to
 *               port - destination port to connect to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool connectTcpSocket(SOCKET sock, struct in_addr addr, unsigned short port)
{
	bool result = false;

	printf("Connecting to %s:%u...", inet_ntoa(addr), port);
	fflush(stdout);

	struct sockaddr_in saddr;
	saddr.sin_family = AF_INET;
	saddr.sin_port = htons(port);
	saddr.sin_addr = addr;
	if (connect(sock, (struct sockaddr*) &saddr, sizeof(saddr)) == 0)
	{
		printf(" OK.\n");
		result = true;
	}
	else
	{
		printf(" FAILED.\n");
		fprintf(stderr, "Couldn't connect to host: %s\n", getErrorMsg(WSAGetLastError()));
	}

	return result;
}

/* Name:         receiveString
 * Description:  Receives and prints a string from a connected socket
 * Parameters:   sock - socket to read from
 * Return code:  true if successful
 *               false if an error occurred
 */
bool receiveString(SOCKET sock)
{
	int total = 0;
	bool done = false;
	bool failed = false;
	char buffer[BUFFER_SIZE];

	while (!done && !failed)
	{
		int bytes = recv(sock, buffer + total, sizeof(buffer) - total - 1, 0);
		if (bytes > 0)
		{
			total += bytes;
			buffer[total] = '\0';
		}
		else if (bytes == 0)
		{
			done = true;
		}
		else if (bytes == SOCKET_ERROR)
		{
			fprintf(stderr, "Couldn't read from socket: %s\n", getErrorMsg(WSAGetLastError()));
			failed = true;
		}
	}

	if (!failed)
	{
		printf("Received %i bytes: <%s>\n", total, buffer);
	}

	return !failed;
}

/* Name:         safeCloseSocket
 * Description:  Checks if a socket is valid, closes it. The socket is
 *               set to an invalid value to avoid double free errors.
 * Parameters:   sock - socket to communicate with
 * Return code:  none
 */
void safeCloseSocket(SOCKET *sock)
{
	if (*sock != INVALID_SOCKET)
	{
		closesocket(*sock);
		*sock = INVALID_SOCKET;
	}
}

int main(int argc, char *argv[])
{
	const char *hostname;
	bool good = getCmdLineArgs(argc, argv, &hostname);

	bool startupSuccessful = false;
	if (good)
		good = startupSuccessful = initWinsock();

	struct in_addr addr;
	if (good)
		good = resolveHostName(hostname, &addr);

	SOCKET sock = INVALID_SOCKET;
	if (good)
		good = createTcpSocket(&sock);

	if (good)
		good = connectTcpSocket(sock, addr, DAYTIME_PORT);

	if (good)
		good = receiveString(sock);

	safeCloseSocket(&sock);

	if (startupSuccessful)
		deinitWinsock();

	return good ? EXIT_SUCCESS : EXIT_FAILURE;
}