#include "common.h"
#include <time.h>
#include <stdbool.h>
#include <openssl/e_os2.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/wait.h>
#include <unistd.h>
#include <netdb.h>

#define KEYFILE "server.pem"
#define CA_FILE "ca.pem"
#define PASSWORD "password"
#define DHFILE "dh1024.pem"
#define BUFSIZZ 20000

#ifdef DEBUG
#define DEBUG_PRINT(...) do{ fprintf( stdout, __VA_ARGS__ ); } while( false )
#else
#define DEBUG_PRINT(...) do{ } while ( false )
#endif


static int disable_nagle  = 0;
char* middlebox_id = "Middlebox1";
char* handler_command = NULL; //"./mboxHandler.sh";


int tcp_listen(int port)
{
    int sock;
    struct sockaddr_in sin;
    int val=1;

    if((sock=socket(AF_INET,SOCK_STREAM,0))<0)
    {
      printf("[MBOX ERROR] Couldn't make socket\n");
      exit(-1);
    }
    
    memset(&sin,0,sizeof(sin));
    sin.sin_addr.s_addr=INADDR_ANY;
    sin.sin_family=AF_INET;
    sin.sin_port=htons(port);
    setsockopt(sock,SOL_SOCKET,SO_REUSEADDR,
      &val,sizeof(val));

/*    if (disable_nagle == 1)
    	set_nagle(sock, 1); */
    
    if(bind(sock,(struct sockaddr *)&sin,
      sizeof(sin))<0)
    {
      printf("[MBOX] Couldn't bind\n");
      exit(-1);
    }

    listen(sock,5);  

	DEBUG_PRINT("[MBOX] Listening at port: %d\n", port); 

    return(sock);
  }


  // TCP connect function 
int tcp_connect(char *host, int port)
{
	struct hostent *hp;
	struct sockaddr_in addr;
	int sock;

	// Resolve host 
	if(!(hp = gethostbyname(host)))
    {
		printf("[MBOX ERROR] Couldn't resolve host %s\n", host);
        exit(-1);
	}
	DEBUG_PRINT("[MBOX DEBUG] Host %s resolved to %s port %d\n",  host,  hp->h_addr_list[0], port);

	memset(&addr, 0, sizeof(addr));
	addr.sin_addr = *(struct in_addr*)
	hp->h_addr_list[0];
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	if((sock=socket(AF_INET,SOCK_STREAM, IPPROTO_TCP))<0)
    {
		printf("[MBOX ERROR] Couldn't create socket\n");
	}
	DEBUG_PRINT("[MBOX DEBUG] Socket created\n"); 

/*    if (disable_nagle == 1)
    	set_nagle(sock, 1); */

	if(connect(sock,(struct sockaddr *)&addr, sizeof(addr))<0)
    {
		printf("[MBOX ERROR] Couldn't connect socket\n");
        exit(-1);
	}
	DEBUG_PRINT("[MBOX DEBUG] Socket connected\n"); 
	
	return sock;
}

void load_dh_params(SSL_CTX *ctx, char *file)
{
	DH *ret=0;
	BIO *bio;

    if ((bio=BIO_new_file(file,"r")) == NULL)
    {
        printf ("[MBOX ERROR] Couldn't open DH file %s\n", file);
        exit(0);
	}

	ret = PEM_read_bio_DHparams(bio, NULL, NULL, NULL);
	BIO_free(bio);
	if(SSL_CTX_set_tmp_dh(ctx,ret) < 0)
    {
        printf ("[MBOX ERROR] Couldn't set DH parameters\n");
		exit(0);
	}
}

// -------------------------------------------------------------
// Check for SSL_write error (just write at this point) 
// TO DO: check behavior per slice 
// -------------------------------------------------------------
void check_SSL_write_error(SSL *ssl, int r, int request_len)
{
	int errorCode = SSL_get_error(ssl, r);	
	switch(errorCode)
	{
		case SSL_ERROR_NONE:
			if(request_len != r)
			{
				printf("[MBOX ERROR] Incomplete write!");
				exit(-1);
			}
			break;

		default:
		{
			printf("[MBOX ERROR] SSL write error %d", errorCode);
			exit(-1);
		}
	}
}



SSL* create_SSL_connection(char *address, char* method)
{
    DEBUG_PRINT("[MBOX DEBUG] Creating new spp conection to: %s\n", address);

	SSL_CTX *ctx;							// SSL context
	SSL *new_ssl;								// SSL context
	BIO *sbio;								// ?
	int sock;								// socket
	char* ipv4 = strtok(strdup(address), ":");	// ip
    int port = atoi(strtok(NULL, ":"));	// port 

    DEBUG_PRINT("[MBOX DEBUG] Connecting to next hop: %s %d\n", ipv4, port);

	int status = COMMON_InitializeSSLCtx(&ctx, KEYFILE, PASSWORD, CA_FILE, ID_MIDDLEBOX_MIN );
	new_ssl = SSL_new(ctx);
	sock = tcp_connect(ipv4, port);
	sbio = BIO_new_socket(sock, BIO_NOCLOSE);
    SSL_set_bio(new_ssl, sbio, sbio);

    return new_ssl;
}

/*
	Just creates a new SSL instance but it does not connect !
*/
SSL* SPP_Callback(SSL *ssl, char *address)
{
	return create_SSL_connection(address, "middlebox");
}



/*
This will terminate but NOT destroy a socket
*/
int shut_down_connections(SSL* ssl)
{
	int r;
	int socket;

	// get socket 
	socket =  SSL_get_fd(ssl);
	DEBUG_PRINT("[MBOX DEBUG] Shutting down  connection!\n");

	r = SSL_shutdown(ssl);
	
	if (r == 0)
	{
		DEBUG_PRINT("[MBOX DEBUG] r=0 trying again to shutdown\n");
		shutdown(socket, 1);
		r = SSL_shutdown(ssl);
	}
	
	// Verify that all went good 
	switch(r)
    {  
		case 1:
			DEBUG_PRINT("[MBOX DEBUG] Succesfully shut down connection!\n");
       		break; // Success
		case 0:
		case -1:
		default: // Error 
			printf("[MBOX ERROR] Connection shutdown failed\n");
	}

	// FIXME -- this should not be need it but ...
	shutdown(SSL_get_fd(ssl), SHUT_WR);

	// close here
	close(socket);
   
	// All good  
    return 0;
}

cJSON* create_slice_record(char* direction, char* data, int dataLen, SPP_SLICE* slice)
{
	cJSON* sliceRecord = cJSON_CreateObject();

	cJSON_AddStringToObject(sliceRecord, "direction", direction);
	cJSON_AddStringToObject(sliceRecord, "middleboxid", middlebox_id);
	cJSON_AddNumberToObject(sliceRecord, "slice", slice->slice_id);
	cJSON_AddStringToObject(sliceRecord, "slicePurpose", slice->purpose);
	cJSON_AddNumberToObject(sliceRecord, "readAccess", slice->read_access);	
	cJSON_AddNumberToObject(sliceRecord, "writeAccess", slice->write_access);

	if (slice->read_access)
	{
		char* dataStr = malloc(dataLen + 1);
		memcpy(dataStr, data, dataLen);
		dataStr[dataLen] = 0;
		cJSON_AddStringToObject(sliceRecord, "data", dataStr);
		free(dataStr);
	}
	else
	{
		char* printableStr = NULL;
		COMMON_Base64Encode(data, dataLen, &printableStr);
		cJSON_AddStringToObject(sliceRecord, "data", printableStr);
	}

	return sliceRecord;
}

cJSON* call_slice_handler (char* inputJsonFile)
{
	if (!handler_command) return NULL;

    DEBUG_PRINT("[MBOX DEBUG] Calling response handler with file %s\n", inputJsonFile);

    char commandBuf[256];
    sprintf(commandBuf, "%s %s", handler_command, inputJsonFile);

    int bufLen = 0;
    char* response = COMMON_CallExternalProcess(commandBuf, &bufLen);

    DEBUG_PRINT ("[MBOX DEBUG] Received %lu bytes from handler (%s)\n", strlen(response), response);
    cJSON* returnJson = cJSON_Parse(response);
    if (returnJson) return returnJson;

    printf ("[MBOX ERROR] Could not parse json from command handler %s  - received { %s }\n", handler_command, response);
    return NULL;	
}

/*
Handles data from previous hop and forwards them to the next one. 
TODO: in theory the two handlers can be merged in a single function...
*/
int handle_previous_hop_data(SSL* prev_ssl, SSL* next_ssl)
{  
    int r,w; 
	long status; 
	char buf[BUFSIZZ];
	SPP_SLICE *slice;       
	SPP_CTX *ctx;   

	// Read HTTP GET (assuming a single read is enough)
	while(1)
    {
		DEBUG_PRINT("[MBOX DEBUG PREVHOP] Waiting to read data from previous hop\n");

        r = SPP_read_record(prev_ssl, buf, BUFSIZZ, &slice, &ctx);
        DEBUG_PRINT("[MBOX DEBUG PREVHOP] SPP_read_record returned\n");

		status = SSL_get_error(prev_ssl, r);
		if (status ==  SSL_ERROR_ZERO_RETURN || status != SSL_ERROR_NONE)
        {
			printf ("[MBOX ERROR PREVHOP] Error code %ld\n", status);
            ERR_print_errors_fp(stdout);
			char tempBuf[100]; 
			ERR_error_string(status, tempBuf); 
			printf("[MBOX ERROR PREVHOP] Connection with previous hop closed, exiting previous hop handler (error %s)\n", tempBuf);
			break;
		}
		else if  (status != SSL_ERROR_NONE)
        {
			printf("[MBOX ERROR PREVHOP] SSL read error code %lu", status);
            exit(-1);
		}

		DEBUG_PRINT("[MBOX DEBUG PREVHOP] Data received (from previous hop) (length %d, slice %d: %s):\n", r, slice->slice_id, slice->purpose); 

		cJSON* sliceRecord = create_slice_record("ClientToServer", buf, r, slice);
		char* printStr = cJSON_Print(sliceRecord);
		printf("[MBOX PREVHOP] Received data (%d bytes)\n", r);
		printf("%s\n", printStr);
		free(printStr);

		char* jsonFileName = COMMON_WriteJSONFile(sliceRecord, middlebox_id);

		cJSON* modifiedSlice = NULL;
		if (slice->write_access) modifiedSlice = call_slice_handler(jsonFileName);
		if (modifiedSlice)
		{
			char* modifiedBuf = cJSON_GetObjectItem(modifiedSlice, "data")->valuestring;
			int newSliceID = cJSON_GetObjectItem(modifiedSlice, "slice")->valueint;

			DEBUG_PRINT("[MBOX DEBUG PREVHOP] Forwarding record to next hop modifed by command handler\n");
			w = SPP_forward_record(next_ssl, modifiedBuf, strlen(modifiedBuf) , SPP_get_slice_by_id(next_ssl, newSliceID), ctx, 1);
			check_SSL_write_error(next_ssl, w, strlen(modifiedBuf)); 
		}
		else
		{
			DEBUG_PRINT("[MBOX DEBUG PREVHOP] Forwarding record to next hop unmodified\n");
			w = SPP_forward_record(next_ssl, buf,r , slice, ctx, 0);
			check_SSL_write_error(next_ssl, w, r); 
		}

		free (jsonFileName);
		cJSON_Delete(sliceRecord);


		//this is probably not necessary...
		if(r == 0 )
        {
			break; 
		}
	}

	// NEW SHUT DOWN 
	DEBUG_PRINT("[MBOX DEBUG PREVHOP] Shutting down next hop\n"); 
	shut_down_connections(next_ssl);
	DEBUG_PRINT("[MBOX DEBUG PREVHOP] Next hop shut down\n"); 

    return(0);
}

int handle_next_hop_data(SSL* prev_ssl, SSL* next_ssl)
{  
    int r,w,status; 
	char buf[BUFSIZZ];
	SPP_SLICE *slice;       
	SPP_CTX *ctx; 

	// Read HTTP GET (assuming a single read is enough)
	while(1)
    {
		DEBUG_PRINT("[MBOX DEBUG NEXTHOP] Waiting to read data from next hop\n");

        r = SPP_read_record(next_ssl, buf, BUFSIZZ, &slice, &ctx);

		status = SSL_get_error(next_ssl, r);
		if (status ==  SSL_ERROR_ZERO_RETURN || status != SSL_ERROR_NONE)
        {
			printf("[MBOX ERROR NEXTHOP] Connection with next hop closed, exiting next hop handler and also closing previous hop connection\n");
			break;
		}
		else if  (status != SSL_ERROR_NONE)
        {
			printf("[MBOX ERROR NEXTHOP] SSL read error code %d", status);
            exit(-1);
		}

		DEBUG_PRINT("[MBOX DEBUG NEXTHOP] Data received (from previous hop) (length %d bytes, slice %d: %s):\n", r, slice->slice_id, slice->purpose); 
		
		cJSON* sliceRecord = create_slice_record("ServerToClient", buf, r, slice);
		char* printStr = cJSON_Print(sliceRecord);
		printf("[MBOX NEXTHOP] Received data (%d bytes)\n", r);
		printf("%s\n", printStr);
		COMMON_WriteJSONFile(sliceRecord, middlebox_id);
		free(printStr);

        w = SPP_forward_record(prev_ssl, buf,r , slice, ctx, 0);
        check_SSL_write_error(prev_ssl, w, r); 

		// In theory this is not necessary
		if(r == 0 )
        {
			break; 
		}
	}
	
	DEBUG_PRINT("[MBOX DEBUG NEXTHOP] Triggering connection with previous hop to close too\n");

    return(0);
}



int handle_data(SSL* prev_ssl, SSL* next_ssl)
{
	pid_t next_handler;
	int status;

	DEBUG_PRINT("[MBOX DEBUG] Initializing data handlers\n");

	next_handler = fork();
	if(next_handler == 0) 
	{
		//child process
		// handle traffic from client
		handle_previous_hop_data(prev_ssl, next_ssl);
		DEBUG_PRINT("[MBOX DEBUG] Exiting previous hop handler\n");
		//TODO: wait for child before returning
		exit(0);
	}
	else 
	{
		//parent 
		handle_next_hop_data(prev_ssl, next_ssl);
		DEBUG_PRINT("[MBOX DEBUG] Exiting next hop handler\n");
	}

	DEBUG_PRINT("[MBOX DEBUG] Waiting previous hop handler to quit before quiting data handler\n");

	wait(&status);

	//CLEAN UP
	close(SSL_get_fd(next_ssl));
	close(SSL_get_fd(prev_ssl));

	//CLEAN UP
	SSL_free(next_ssl);
	//SSL_free(prev_ssl);

	DEBUG_PRINT("[MBOX DEBUG] Exiting data handler!\n");

	return 0;
}




int main(int argc, char **argv)
{
	int sock, s;
	BIO *sbio;
	SSL_CTX *ctx;
	SSL *ssl;
	int r;
	pid_t pid;
	int port = 8423;    
	char *prxy_address = "127.0.0.1:8423";
	SSL* ssl_next = NULL;
	char buf[BUFSIZZ];
	int ret;    

    if (argc < 4)
    {
		printf ("Usage: middlebox port proxy-address midlebox-id [command handler...]\n");
		exit(0);
    }
	port = atoi(argv[1]);
	prxy_address = argv[2];
	middlebox_id = argv[3];
    DEBUG_PRINT("[MBOX DEBUG] port=%d prxy_address=%s id=%s\n", port, prxy_address, middlebox_id);  


	if (argc > 4)
	{
		char tempBuf[1024];
		memset(tempBuf, 0, 1024);
		int written = 0;
		for (int i = 4; i < argc; i++)
		{
			written += snprintf(tempBuf + written, 1024 - written, "%s ", argv[i]);
		}
		handler_command = strdup(tempBuf);
		DEBUG_PRINT("[MOX DEBUG] Handler command = %s\n", handler_command);
	}
	else
	{
		DEBUG_PRINT("[MBOX DEBUG] Handler command not set\n");
	}

	// TO-DO: Actually check the return value of this e.g. if the files don't exist
	int status = COMMON_InitializeSSLCtx(&ctx, KEYFILE, PASSWORD, CA_FILE, ID_MIDDLEBOX_MIN );
	load_dh_params(ctx,DHFILE);

	// Socket in listen state
	sock = tcp_listen(port);

	// Wait for client request 
	int nConn = 0;

	while(1)
    {
		// keep track of number of connections
		nConn++; 

		if((s = accept(sock, 0, 0)) < 0)
        {
			printf("[MBOX ERROR] Problem socket accept\n");
		}

		signal(SIGCHLD, SIG_IGN);
		// fork a new proces 
		if((pid = fork()))
        {
			close(s);
		} 
        else 
        {
            sbio = BIO_new_socket(s, BIO_NOCLOSE);
            ssl = SSL_new(ctx);
            SSL_set_bio(ssl, sbio, sbio);

            if ((r = SPP_proxy(ssl, prxy_address, SPP_Callback, &ssl_next)) <= 0) 
            {
                printf ("[MBOX ERROR] SPP proxy error");
            } 
            else 
            {
                DEBUG_PRINT("[MBOX DEBUG] SPP proxy OK\n"); 
            }

            // Set server URL
            //char* targetServercopy = "127.0.0.1:4433";
            //ssl->spp_server_address = targetServercopy;

			handle_data(ssl, ssl_next);

			// Close socket
    		close(s);
    	
			// Exit child thread
			exit(0);
		}
	}
	
	// Clean context
	COMMON_DestroyCtx(ctx);
	
	// Exit
	exit(0);    
}