#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#define PORT 45894

#ifdef WIN32
#include <winsock2.h>
#include <windows.h>

#define ERRORME(x) == (x)

#else
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>

#define ERRORME(x) < 0

#endif //WIN32

struct connection {
	#ifdef WIN32
	SOCKET sockfd;
	#else
	int sockfd;
	#endif //WIN32
	int run, n;
	struct sockaddr_in partner;
};

void start(int type);
void *userIn(void *arg);
void userRecv(void *arg);

int main(int argc, char **argv) {
	#ifdef WIN32
	WSADATA wsa;
	if (WSAStartup(MAKEWORD(2, 0), &wsa) < 0) {
		printf("Error starting wsa: %d\n", WSAGetLastError());
		return -1;
	}
	#endif //WIN32

	char buf[256];
	int type;
	printf("Enter type! 1 for server, 2 for client!");
	fgets(buf, 256, stdin);
	sscanf(buf, "%d", &type);

	start(type);

	#ifdef WIN32
	WSACleanup();
	#endif //WIN32

	return 0;
}

void start(int type) {
	if (type == 1) {
		struct connection s;

		#ifdef WIN32
		SOCKET sockfd, newsockfd;
		int clisize;
		#else
		int sockfd, newsockfd;
		socklen_t clisize;
		#endif //WIN32

		struct sockaddr_in server, client;
		memset(&server, 0, sizeof(server));

		server.sin_family = AF_INET;
		server.sin_port = htons(PORT);
		server.sin_addr.s_addr = INADDR_ANY;

		if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) ERRORME(INVALID_SOCKET)) {
			#ifdef WIN32
			closesocket(sockfd);
			printf("Error creating socket: %d\n", WSAGetLastError());
			#else
			close(sockfd);
			printf("Error creating socket.\n");
			#endif //WIN32

			return;
		}

		if (bind(sockfd, (struct sockaddr *)&server, sizeof(server)) ERRORME(SOCKET_ERROR)) {
			#ifdef WIN32
			printf("Error binding socket: %d\n", WSAGetLastError());
			closesocket(sockfd);
			#else
			perror("Error binding socket.\n");
			close(sockfd);
			#endif //WIN32

			return;
		}

		listen(sockfd, 0);

		clisize = sizeof(client);
		if ((newsockfd = accept(sockfd, (struct sockaddr *)&client, &clisize)) ERRORME(INVALID_SOCKET)) {
			#ifdef WIN32
			printf("Error accepting connection: %d\n", WSAGetLastError());
			closesocket(sockfd);
			#else
			printf("Error accepting connection.\n");
			close(sockfd);
			#endif //WIN32

			return;
		}

		send(newsockfd, "Connected.\n", 12, 0);

		s.sockfd = newsockfd;
		s.run = 1;
		s.partner = client;

		pthread_t input;
		pthread_create(&input, NULL, userIn, &s);

		userRecv(&s);

		#ifdef WIN32
		closesocket(newsockfd);
		closesocket(sockfd);
		#else
		close(newsockfd);
		close(sockfd);
		#endif //WIN32

		return;
	} else if (type == 2) {
		struct connection c;

		#ifdef WIN32
		SOCKET sockfd;
		#else
		int sockfd;
		#endif //WIN32

		struct sockaddr_in server;
		memset(&server, 0, sizeof(server));

		server.sin_port = htons(PORT);
		server.sin_family = AF_INET;

		printf("Enter IP to connect to: ");
		char in[256];
		char ip[15];
		fgets(in, 256, stdin);
		sscanf(in, "%s\n", ip);

		server.sin_addr.s_addr = inet_addr(ip);

		if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) ERRORME(INVALID_SOCKET)) {
			#ifdef WIN32
			printf("Error creating socket: %d\n", WSAGetLastError());
			#else
			printf("Error creating socket.\n");
			#endif //WIN32

			return;
		}

		if (connect(sockfd, (struct sockaddr *)&server, sizeof(server)) < 0) {
			#ifdef WIN32
			printf("Error conencting: %d\n", WSAGetLastError());
			#else
			printf("Error connecting.\n");
			#endif //WIN32

			return;
		}

		send(sockfd, "Connected.\n", 12, 0);

		c.sockfd = sockfd;
		c.run = 1;
		c.partner = server;

		pthread_t input;
		pthread_create(&input, NULL, userIn, &c);

		userRecv(&c);

		#ifdef WIN32
		closesocket(sockfd);
		#else
		close(sockfd);
		#endif //WIN32

		return;
	} else {
		printf("Error! Invalid type!");
		return;
	}
}

void *userIn(void *arg) {
	struct connection *me = (struct connection *)arg;
	char buffer[512];
	char test[7];
	while(me->run) {
		memset(buffer, 0, 512);
		memset(test, 0, 6);

		fgets(buffer, 512, stdin);

		memcpy(test, &buffer, 6);
		if(strlen(buffer) > 6) {
			test[6] = '\0';
			if (strcmp(test, "~QUIT~") == 0) {
				me->run = 0;
				send(me->sockfd, "~QUIT~", 7, 0);
			} else
				send(me->sockfd, buffer, strlen(buffer), 0);
		} else
			send(me->sockfd, buffer, strlen(buffer), 0);
	}
}

void userRecv(void *arg) {
	struct connection *me = (struct connection *)arg;
	char buffer[256];
	while(me->run) {
		if ((me->n = recv(me->sockfd, (char *)buffer, 255, 0)) < 0) {
			#ifdef WIN32
			printf("Error at recv: %d\n", WSAGetLastError());
			#else
			printf("Error at recv\n");
			#endif //WIN32

			me->run = 0;
			return;
		}

		buffer[(me->n)] = '\0';
		if (strcmp(buffer,"~QUIT~") == 0) {
			me->run = 0;
			return;
		}
		printf("Buddy: %s", buffer);
	}
}