/*
 * This is a better server than simple_server.c
 * I'm only going to comment changes from simple_server.c
 * So if there's something that doesn't make any sense in here, maybe look at that one.
 */
#include<sys/socket.h>
#include<sys/types.h>
#include<netinet/in.h>
#include<netinet/ip.h>
#include<netdb.h>
#include<arpa/inet.h>
#include<stdio.h>
#include<unistd.h>
#include<signal.h>
#include<stdlib.h>
#include<pthread.h>

int skt;
void shut_down_server(int signal){
	printf("\nDoing everything we need to do to shut down the server\n");
	close(skt);
	exit(0); 
}

char debug = 1;

/*
 * To use this:
 * 	proxy_server listen_port destination_address destination_port
 */
int main(int argc, char ** argv){
	uint16_t listen_port = 5141; // 5141 is the default
	if(argc > 1){
		listen_port = atoi(argv[1]);
		if(listen_port < 1){
			printf("Invalid port:  %s\n", argv[1]);
			exit(1);
		}
		printf("Will listen on port %d\n", listen_port);
	} else {
		printf("No port specified, defaulting to %d\n", listen_port);
	}	
	struct sigaction sa;
	sa.sa_handler = shut_down_server;
	sigaction(SIGINT, &sa, 0);

	struct sockaddr_in sad;
	sad.sin_port = htons(listen_port);
	sad.sin_addr.s_addr = INADDR_ANY;
	sad.sin_family = AF_INET;

	skt = socket(AF_INET, SOCK_STREAM, 0); // Step 1
	if(skt == -1){
		perror("socket");
		return 1;
	}
	if( bind(skt, (struct sockaddr *)(&sad), sizeof(struct sockaddr_in)) ){ // step 2
		perror("bind");
		return 1;
	}
	if( listen(skt, 5) ){ // step 3
		perror("listen");
		return 1;
	}	
	int client_fd;
	struct sockaddr_in client_address;
	socklen_t address_size = sizeof(struct sockaddr_in);
	while(1){
		client_fd = accept(skt, (struct sockaddr *)(&client_address), &address_size); // step 4

		/* The next bit establishes the connection to the final destination */
		int destination_fd = socket(AF_INET, SOCK_STREAM, 0);
		struct sockaddr_in sad;
		sad.sin_port = htons(atoi(argv[3]));
		sad.sin_family = AF_INET;

		int skt = socket(AF_INET, SOCK_STREAM, 0);

		// do a dns lookup
		struct hostent* entry = gethostbyname(argv[2]);
		if(!entry){
			if(h_errno == HOST_NOT_FOUND){
				printf("This is our own message that says the host wasn't found\n");
			}
			herror("gethostbyname");
			return 1;
		}
		struct in_addr **addr_list = (struct in_addr**)entry->h_addr_list;
		struct in_addr* c_addr = addr_list[0];
		char* ip_string = inet_ntoa(*c_addr);
		sad.sin_addr = *c_addr; // copy the address we found into sad
		// Finally done with DNS!
		printf("Connecting to:  %s\n", ip_string);

		if( connect(destination_fd, (struct sockaddr*)&sad, sizeof(struct sockaddr_in)) ){
			perror("connect");
			return 1;
		}
		void* thread_two(void*){
			char buffer[1024];
			ssize_t len;
			while(1){
				// thread two
				len = read(destination_fd, buffer, 1023);
				if(debug){
					buffer[len] = 0;
					printf("Server sent:  %s\n", buffer);
				}
				if(len <= 0)
					break;
				write(client_fd, buffer, len);
			}
		}
		pthread_t T;
		pthread_create(&T, 0, thread_two, 0);

		/* Should be all connected now! */
		char buffer[1024];
		ssize_t len;
		while(1){
			// thread one
			len = read(client_fd, buffer, 1024);
			if(debug){
				buffer[len] = 0;
				printf("Client sent:  %s\n", buffer);
			}
			if(len <= 0)
				break;
			write(destination_fd, buffer, len);

		}

		close(destination_fd);
		close(client_fd);
	}
	return 0;
}