#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define MAXHOSTNAMELEN 128


int create_and_bind_socket(port)
{
    int fd;
    struct sockaddr_in myAddress;

    /* Create a UDP socket - SOCK_DGRAM means UDP */
    if((fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
	perror("Error creating socket");
	exit(1);
    }

    /* Bind our socket to a particular address (port number for any 
    ** IP address of this machine 
    */
    myAddress.sin_family = AF_INET;
    myAddress.sin_port = htons(port);
    myAddress.sin_addr.s_addr = INADDR_ANY;
    
    /* Have to cast our Internet socket address type to a 
    ** Generic socket type (pointers to).
    */
    if(bind(fd, (struct sockaddr *)&myAddress, 
	    sizeof(struct sockaddr_in)) < 0) {
	perror("Error binding socket to port");
	exit(1);
    }
    return fd;
    /* Now have a UDP socket, bound to a particular port number */
}


void communicate(fd, destPort)
{
    fd_set readSet;
    int numBytes;
    char buffer[1024];
    struct sockaddr_in theirAddress, destAddress;
    int theirAddrSize;
    char hostname[MAXHOSTNAMELEN];

    /* Destination address for any datagrams sent is constant - set it
    ** up now and we'll use it multiple times. Dest address is the 
    ** given port number on localhost.
    */
    destAddress.sin_family = AF_INET;
    destAddress.sin_port = htons(destPort);
    destAddress.sin_addr.s_addr = inet_addr("127.0.0.1");

    FD_ZERO(&readSet);

    while(1) {
	/* Initialise our readSet of fd's before every call to select(). Two
	** fds of interest are stdin (0) and the UDP socket we've bound to a 
	** port.
	*/
	FD_SET(fd, &readSet);
	FD_SET(STDIN_FILENO, &readSet);

	/* Block until data is available (or EOF) on one (or both) of 
	** the fds of interest 
	*/
	if(select(fd + 1, &readSet, NULL, NULL, NULL) < 0) {
	    perror("Select failed");
	    exit(1);
	}
	/* Now know that data is available (or EOF) on one or both
	** of the fds
	*/

	if(FD_ISSET(STDIN_FILENO, &readSet)) {
	    /* Data is available on standard input (or EOF).
	    ** read will not block.
	    */
	    numBytes = read(STDIN_FILENO, buffer, 1024);
	    if(numBytes > 0) {
		/* Got some data from standard input - send it off to
		** the destination address we created earlier */
		if(sendto(fd, buffer, numBytes, 0, 
			(struct sockaddr *)&destAddress, 
			sizeof(struct sockaddr_in)) < 0) {
		    perror("sendto failed");
		    exit(1);
		}
	    } else if(numBytes == 0) {
		/* Stdin closed - exit */
		close(fd);
		exit(1);
	    } else {
		/* error happened */
		perror("Error reading from standard input");
		exit(1);
	    }
	}
	if(FD_ISSET(fd, &readSet)) {
	    /* Data was available on the UDP socket */
	    theirAddrSize = sizeof(struct sockaddr_in);
	    /* Read the data - and get the address it was sent from */
	    numBytes = recvfrom(fd, buffer, 1024, 0, 
		    (struct sockaddr*)&theirAddress, 
		    &theirAddrSize);
	    if(numBytes > 0) {
		/* Got some data - print out where it came from and what it was */
		getnameinfo((struct sockaddr*)&theirAddress, theirAddrSize, hostname,
			MAXHOSTNAMELEN, NULL, 0, 0);
		printf("From %s:%d: ", hostname,
			  ntohs(theirAddress.sin_port));
		fflush(stdout);
		write(STDOUT_FILENO, buffer, numBytes);
	    } else if(numBytes == 0) {
		close(fd);
		exit(1);
	    } else {
		perror("Error in recvfrom");
		exit(1);
	    }
	}
    }
}

int main(int argc, char* argv[])
{
    int myPort, destPort, fd;

    if(argc != 3) {
	printf("Usage: %s my-port dest-port\n", argv[0]);
	exit(1);
    }
    myPort = atoi(argv[1]);
    destPort = atoi(argv[2]);
    /* Get ready to accept UDP datagrams on the given port */
    fd = create_and_bind_socket(myPort);
    /* Implement the chat functionality */
    communicate(fd, destPort);
    return 0;
}