#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netdb.h>

#define IAM_CLIENT 1
#define IAM_SERVER 2
#define INIT_RETRIES 100
#define INIT_RETRY_WAIT 1
#define BUFFER_SIZE 1024
#define SERVER_PORT 4010

char whoami[16];
void server_code(int socketFd);
void client_code(int socketFd);

void usage(){
    printf("socket client|server\n");
}

int main(int argc, char** argv){
    int mode = 0;
    int socketFd;
    
    if (argc == 2){
        if (strcmp(argv[1], "client") == 0){
            mode = IAM_CLIENT;
            strcpy(whoami, "CLIENT");
            printf("running in client mode\n");
        }
        if (strcmp(argv[1], "server") == 0){
            mode = IAM_SERVER;
            strcpy(whoami, "SERVER");
            printf("running in server mode\n");
        }
    }
    if (mode != IAM_CLIENT && mode != IAM_SERVER){
        usage();
        exit(1);
    }

    // create the socket and get back a file descriptor
    socketFd = socket(PF_INET, SOCK_STREAM, 0);
    if (socketFd < 0){
        printf("%s: failure creating socket\n", whoami);
        exit(1);
    }
    
    if (mode == IAM_SERVER){
		server_code(socketFd);
    }
    else if (mode == IAM_CLIENT){
		client_code(socketFd);
    }

	return 0;
}

void server_code(int socketFd){
    int connectedClientFd;

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

    address.sin_family = AF_INET;
    address.sin_addr.s_addr = htonl(INADDR_ANY);
    address.sin_port = htons(SERVER_PORT);

    int retValue, i;
    for (i = 0; i < INIT_RETRIES; ++i){

    	retValue = bind(socketFd,(struct sockaddr *)&address, sizeof(address));
    	if (retValue == 0){
    		break; // success
    	}
    	sleep(INIT_RETRY_WAIT);
    }
    if (retValue < 0){
    	printf("%s: server failed to bind after %d retries\n", whoami, INIT_RETRIES);
    	exit(1);
    }
   
  	retValue = listen(socketFd, SOMAXCONN);
  	if (retValue < 0){ 
  		printf("%s: server failed to listen\n", whoami);
  		exit(1);
  	}
    socklen_t slen = sizeof(address);
    while (1) {
    	printf("%s waiting to accept incoming connections\n", whoami);
        fflush(stdout);
        connectedClientFd = accept(socketFd, (struct sockaddr *)&address, &slen);
        if (connectedClientFd < 0){
        	printf("%s: error on accept socket\n", whoami);
            exit(1);
        }

        int pid = fork();
        if (pid < 0) { 
        	printf("%s: server could not fork\n", whoami);
            exit(1);
        }
        if (pid == 0){ // child process
        	break; // go on to send and receive mode while parent continues to wait for connections 
        }
    }

    printf("%s: client connected\n", whoami);

	char buffer[BUFFER_SIZE];
 	char recvOutMsg[] = "SERVER: received message from client: ";
 	char nlMsg[] = "\n";
	while (1){
		ssize_t bytes = recv(connectedClientFd, buffer, BUFFER_SIZE, 0);
		if (bytes <= 0){
			printf("%s: error from rev, assuming client disconnection\n", whoami);
			return;
		}
		write(1, recvOutMsg, sizeof(recvOutMsg));
		write(1, buffer, bytes);
		write(1, nlMsg, sizeof(nlMsg));
	}
}

void client_code(int socketFd){
    char* server_ip_address = "192.168.0.108";

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

	struct hostent* host_entry = gethostbyname(server_ip_address);
    address.sin_family = AF_INET;
    memcpy((char *)&address.sin_addr,(char *)host_entry->h_addr, host_entry->h_length);
    address.sin_port = htons(SERVER_PORT);

    int retValue, i;
    for (i = 0; i < INIT_RETRIES; ++i){
    	retValue = connect(socketFd,(struct sockaddr *)&address, sizeof(address));
    	if (retValue == 0){
			break;
		}
    	printf("%s: could not connect to server ...\n", whoami);
		sleep(INIT_RETRY_WAIT);
	}
	if (retValue < 0){
    	printf("%s: could not connect to server after %d retries\n", whoami, INIT_RETRIES);
        exit(1);
    }

    printf("%s: connected to server\n", whoami);

	char buffer[BUFFER_SIZE];
	unsigned int bufsize;
	for (i = 0; i < 60; ++i){
		sprintf(buffer, "msg from pid %d, iteration %d, current time %d", getpid(), i, time(NULL));
		bufsize = strlen(buffer);
		int bytes = send(socketFd, buffer, bufsize, 0);
		if (bytes != bufsize){
			printf("%s: error sending msg, assuming server is down\n", whoami);
			exit(1);
		}
		sleep(1);
	}

	printf("%s: client done sending all messages, disconnecting\n", whoami);
}