/*******************************************************************************
* Solution made in class
*******************************************************************************/

#include<stdio.h>
#include<signal.h>                      /* For kill()                         */
#include<stdlib.h>
#include<string.h>
#include<strings.h>

#include<sys/select.h>                  /* For select()                       */
#include<unistd.h>                      /* For select()                       */
#include<sys/select.h>                  /* For select()                       */
#include<sys/time.h>                    /* For select()                       */
#include<sys/socket.h>                  /* For socket(), listen(),            */
                                        /* setsockopt()                       */
#include<sys/types.h>                   /* For bind(), setsockopt()           */
#include<netinet/in.h>                  /* For IPPROTO_TCP                    */
#include<netdb.h>                       /* For gethostbyname()                */

#define PORT_NR 1025
#define STR_SIZE 100

/* Function decleration */
void printSetup();
int client_setup(char *server);
int server_setup();

/*******************************************************************************
* Function name:    INT main(INT argc, CHAR, **argv)
* Parameters:       argc - number of arguments
*                   **argv - array of arguments
* Returns:          0 on success
*                   < 0 on failure
* Description:      Assumes client if one argument is incuded on startup, else
*                   server. Initiates the server/client setup respectively.
*******************************************************************************/
int main(int argc, char **argv) {
  char *machine;
  struct hostent *hostp;

  if (argc == 1) {                      /* Start as a server                  */
    printf("Server\n");
    return server_setup();

  } else if (argc == 2){                /* Start as client                    */
    printf("Client\n");
    machine = argv[1];                  /* Get machine                        */
    int retVal = client_setup(machine);
    return retVal;
  }

  printf("Invalid argument count");
  return -1;
}

/*******************************************************************************
* Function name:    INT server_setup(VOID)
* Returns:          0 on success
*                   -1 if unable to find create socket
*                   -2 if bind error
*                   -3 if listen error
* Description:      Assumes client if one argument is incuded on startup, else
*                   server. Initiates the server/client setup respectively.
*******************************************************************************/
int server_setup() {

  /* Declearations */
  struct sockaddr_in serv_addr, client_addr;

  printSetup();                         /* Print server data                  */

  /* Create socket */
  int request_sd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
  int client_sd;

  /* Init address struct */
  bzero(&serv_addr, sizeof(struct sockaddr_in));
  serv_addr.sin_family = AF_INET;
  serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
  serv_addr.sin_port = htons(PORT_NR);

  /********************** Prepare for incomming connections *******************/
  bind(request_sd, (struct sockaddr *)&serv_addr, sizeof(struct sockaddr_in));

  //"CONNECT" == NEW CLIENT
  char buffer[8];
  buffer[7] = '\0';

  socklen_t len = sizeof(struct sockaddr_in);
  recvfrom(request_sd, buffer, 7, 0, (struct sockaddr*) &client_addr, &len);

  printf("%s\n", buffer);

  if(!strcmp(buffer, "CONNECT"))
    chat(request_sd);                     /* Start chat program               */

  close(request_sd);                    /* Cleanup                            */

  printf("Chat ended\n");

  return 0;
}

/*******************************************************************************
* Function name:    INT client_setup(CHAR *server)
* Parameter:        *server - name of server
* Returns:          0 on success
*                   -1 if unable to get host
*                   -2 if unable to create socket
*                   -3 if error during connect
* Description:      Creates a connection to a server and starts a chat procedure
*******************************************************************************/
int client_setup(char *server) {

  /* Get server details */
  struct hostent *hostp;
  hostp = gethostbyname(server);
  if(hostp == 0) {                      /* Unable to get host                 */
    perror("Hostname error");
    return -1;
  }

  /* Initiate server address structure */
  struct sockaddr_in serv_addr;
  bzero(&serv_addr, sizeof(struct sockaddr_in));
  serv_addr.sin_family = AF_INET;
  memcpy(&serv_addr.sin_addr, hostp->h_addr, hostp->h_length);
  serv_addr.sin_port = htons(PORT_NR);

  /***************************** Connect to server ****************************/
  int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

  char *buf = "CONNECT";
  sendto(sock, buf, 7, 0, (struct sockaddr*)&serv_addr, 
	 sizeof(struct sockaddr_in));

  printf("Connected to server\n");

  /* Start chat program */
  chat(sock);
  close(sock);
  
  printf("Chat ended\n");

  return 0;
}

struct sms {
  int size;
  int checksum;
  char txt[100];
};

/*******************************************************************************
* Function name:    INT char(INT socket)
* Parameter:        socket - socket connection is established on
* Returns:          0 on success
*                   -1 on select error
* Description:      Allows chatting over a socket connection
*******************************************************************************/
int chat(int socket) {
  char rcv_buf[STR_SIZE];
  char inp_buf[STR_SIZE];

  fd_set fds, readfds;
  int numsocks, maxsocks, active_sock_cnt, i;
  struct timeval timeout;

  numsocks = socket + 1;
  maxsocks = numsocks;
  timeout.tv_sec = 0;
  timeout.tv_usec = 0;

  /*************************** Prepare for select *****************************/
  
  FD_ZERO(&fds);
  FD_SET(socket, &fds);
  FD_SET(0, &fds);

  fprintf(stderr, "entering chat on sock %i\n>", socket);

  do {
    readfds = fds;
    active_sock_cnt = select(numsocks, &readfds, NULL, NULL, &timeout);
    if(active_sock_cnt == -1) {         /* Select error */
      perror("Select error");
      return -1;
    }
    
    /************************* Handle socket I/O ******************************/
    bzero(rcv_buf, STR_SIZE);
    if(FD_ISSET(0,&readfds)) {
      fgets(inp_buf, STR_SIZE, stdin);
      write(socket, inp_buf, strlen(inp_buf));
    } else if(FD_ISSET(socket, &readfds)) {
      read(socket, rcv_buf, STR_SIZE - 1);
      fprintf(stderr, "received: %s", rcv_buf);
    }

  } while(strcmp("exit\n", rcv_buf) && strcmp("exit\n", inp_buf));
  
  return 0;
}

/*******************************************************************************
* Function name:    VOID printSetup(VOID)
* Returns:          VOID
* Description:      Prints server data
*******************************************************************************/
void printSetup() {
  fprintf(stderr, "Selected port number: %d\n", PORT_NR       );
  fprintf(stderr, "Name:                 %s\n", getenv("HOST"));
}