/* Server side - UDP Code				    */
/* By Connor McKee							*/

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

#include "gethostbyname.h"
#include "networks.h"
#include "safeUtil.h"
#include "cpe464.h"
#include "pollLib.h"
#include "pdu.h"
#include "window.h"

#define MAXBUF 80
#define DEBUG_FLAG 1

// server states
typedef enum STATE {
	INIT, DONE, FILENAME, SEND_DATA, WAIT_ACK, TIMEOUT_ACK, WAIT_EOF_ACK, TIMEOUT_EOF_ACK
} STATE;

// Information passed on from the client in setup/filename packet
typedef struct {
	uint32_t windowSize;
	uint16_t bufferSize;
	char fileName[MAX_FNAME_LEN+1];
	FILE *dataFile;
} ClientInfo;

void processServer(int serverSocket, double errRate);
void processClient(uint8_t *pkt, int32_t recvLen, ConnectionData *client);
int checkArgs(int argc, char *argv[]);
double getRate(char * argv[]);
int safeRecvBuf(int serverSocket, uint8_t *pkt, int len, ConnectionData *client, uint8_t *flag, uint32_t *seqNum);
STATE processFilenameEX(ConnectionData *client, uint8_t *pkt, int32_t recvLen, ClientInfo *meta);
void sendFilenameResponse (ConnectionData *client, uint8_t response);
STATE sendData(ConnectionData *client, ClientInfo metadata, Window *window, uint32_t *seqNum);
void handleDataSending(ConnectionData *client, Window *window, uint32_t *seqNum, uint8_t *dataPayload, uint32_t bytesRead);
int processWindowClosed(ConnectionData *client, Window *window, int *retryCount);
STATE processWaitEOFACK(ConnectionData *client, ClientInfo metadata, Window *window, uint32_t *seqNum);
bool processRRandSREJOpen(ConnectionData *client, Window *window);
bool processRRandSREJ(ConnectionData *client, Window *window);
void processRR(Window *window, uint8_t *packet);
void processSREJ(ConnectionData *client, Window *window, uint8_t *packet);
void sendLowestPacket(ConnectionData *client, Window *window);
void sendEOF(ConnectionData *client, uint32_t *seqNum);
void cleanExit(Window *window, ConnectionData *client);


int main ( int argc, char *argv[]  )
{ 
	int socketNum = 0;				
	int portNumber = 0;

	setupPollSet();

	portNumber = checkArgs(argc, argv);

	sendErr_init(getRate(argv), DROP_ON, FLIP_ON, DEBUG_ON, RSEED_ON);
		
	socketNum = udpServerSetup(portNumber);

	processServer(socketNum, getRate(argv));

	close(socketNum);
	
	return 0;
}


void processServer(int serverSocket, double errRate) {
	pid_t pid = 0;
	int currentSocket = 0;
	int clientSocket = 0;
	int clientStatus = 0;
	int recvLen = 0;
	uint8_t packet[MAX_BUFFER] = {0};
	ConnectionData client;
	uint8_t flag = 0;
	uint32_t seqNum = 0;

	addToPollSet(serverSocket);

	while (1) {
		// block to wait for a new client to connect
		if ((currentSocket = pollCall(POLL_WAIT_FOREVER)) != -1) {
			// Initial receive to gather client info
			recvLen = safeRecvBuf(serverSocket, packet, MAX_BUFFER, &client, &flag, &seqNum);

			// check for CRC Error, if error, continue to next response, don't fork
			if(!validChecksum(packet, recvLen)) {
				continue;
			}

			// Check for fork errors, new process started otherwise
			if ((pid = fork()) < 0) {
				perror("Fork failed");
				exit(-1);
			}

			// Child, client process
			if (pid == 0) {
				// Close the server socket and remove from pollSet
				close(serverSocket);
				removeFromPollSet(serverSocket);

				// Set up the socket for the child client process and add it to the pollSet
				clientSocket = socket(AF_INET6, SOCK_DGRAM, 0);
				client.socket = clientSocket;
				addToPollSet(clientSocket);

				// received data from client, process on the new client socket
				processClient(packet, recvLen, &client);
				
				//done processing
				exit(0);
			}
			
			// For single client processing
			// received data from client, process
			//client.socket = serverSocket;
			//processClient(serverSocket, packet, recvLen, &client);
			//done processing
			//exit(0);

		} else {
			printf("Poll timed out waiting for client\n");

			// Wait for dead children
			while (waitpid(-1, &clientStatus, WNOHANG) > 0) { }
		}
	}

}


void processClient(uint8_t *packet, int32_t recvLen, ConnectionData *client)
{
	STATE state = INIT;
	uint32_t seqNum = 0;
	ClientInfo metadata;
	Window window;

	while (state != DONE) {
		switch(state) {
			case INIT:
				state = FILENAME;
				break;

			case FILENAME:
				state = processFilenameEX(client, packet, recvLen, &metadata);
				break;

			case SEND_DATA:
				// setup the window
				window = createWindow(metadata.windowSize);
				state = sendData(client, metadata, &window, &seqNum);
				break;

			case WAIT_EOF_ACK:
				state = processWaitEOFACK(client, metadata, &window, &seqNum);
				return;
				break;

			case DONE:
				cleanExit(&window, client);
				break;

			default:
				printf("got to default in state machine\n");
				state = DONE;
				break;
		}
	}

	// Clean up the window table, exited the state machine
	free(window.table);
}

STATE processWaitEOFACK(ConnectionData *client, ClientInfo metadata, Window *window, uint32_t *seqNum) {
	uint8_t packet[MAX_BUFFER] = {0};
	STATE nextState = DONE;
	int retryCount = 0;
	int recvLen = 0;
	uint32_t eofSeqNum = *seqNum;

	// send out the EOF initially
	sendEOF(client, seqNum);

	// continuously loop to poll and get client packets. No more sending of data now.
	// could be getting RRs, SREJS, EOF_ACK (Finally done! clean up after the client)
	// else, poll expires, resend EOF, increment counter, poll again
	while (1) {
		if (pollCall(POLL_WAIT_1_SEC) != -1) {
			// Data to process
			recvLen = safeRecvfrom(client->socket, packet, MAX_BUFFER, 0, (struct sockaddr *) &((*client).address), &((*client).len));

			// Verify CRC
			if (!validChecksum(packet, recvLen)) {
				// Bad packet, continue
				continue;
			}

			// Check flags
			if (packet[FLAG_LOC] == RR_FLAG) {
				processRR(window, packet);
				retryCount = 0;

			} else if (packet[FLAG_LOC] == SREJ_FLAG) {

				// if the SREJ is for EOF, don't process, send lowest packet until fixed
				if (getOffsetSeqNumber(packet) == eofSeqNum) {
					sendLowestPacket(client, window);
				} else {
					processSREJ(client, window, packet);
					retryCount = 0;
				}

			
			} else if (packet[FLAG_LOC] == EOF_ACK_FLAG) {
				printf("Received Client EOF ACK. Cleaning up and Exiting.\n");
				nextState = DONE;
				break;
			
			} else {
				// Flag was something else, garbage? Send EOF again and increment
				sendEOF(client, seqNum);
				retryCount++;
			}

		} else {
			// timeout, no data from client. Send EOF again, increment, continue.
			sendEOF(client, seqNum);
			retryCount++;
		}

		if (retryCount > 9) {
			// Poll timed out ten times. Client unresponsive
			printf("Client EOF ACK lost. Cleaning up.\n");
			nextState = DONE;
			break;
		}

	}

	return nextState;
}


// reads from the data file, places information into the buffer and sets the next state based on
// the sizes set in ClientInfo
STATE sendData(ConnectionData *client, ClientInfo metadata, Window *window, uint32_t *seqNum) {
	uint8_t dataPayload[MAX_BUFFER] = {0};
	uint32_t bytesRead = 0;
	int retryCount = 0;
	bool eofFound = false;
	STATE nextState = DONE;

	// loop until EOF read from file
	while(!eofFound) {
		while (isWindowOpen(*window)) {
			// Read the specified number of bytes (1 byte chunks) from file and fill a payload
			bytesRead = fread(dataPayload, sizeof(uint8_t), metadata.bufferSize, metadata.dataFile);

			// check for EOF after reading
			if (bytesRead == 0) {
				if (feof(metadata.dataFile)) {
					// EOF! Need to send EOF packet and move to EOF ack state
					eofFound = true;
				}

				if (ferror(metadata.dataFile)) {
					printf("ERROR WHEN READING FROM FILE: %s\n", metadata.fileName);
					nextState = DONE;
				}
				// if 0 bytes were read, either EOF was found and set or an error occurred while reading. 
				// Break from loop and double check for EOF
				break;
			}

			// EOF not found, send data
			handleDataSending(client, window, seqNum, dataPayload, bytesRead);

			// Handle and process any RR/SREJs
			if (!processRRandSREJOpen(client, window)) {
				// discard the packet and continue looking for more if checksum failed
				continue;
			}
		}

		// window closed
		while(!isWindowOpen(*window)) {
			if(processWindowClosed(client, window, &retryCount) < 0) {
				// Tried to break deadlock 10 times, time to exit
				return DONE;
			}
		}

		// reset the retry counter to 0 here
		retryCount = 0;
	}

	// if EOF, go to EOF sending state here
	if (eofFound) {
		nextState = WAIT_EOF_ACK;
	}

	return nextState;
}

// Main processing when the window is closed. Tries to process SREJs/RRs, if none available, sends back the lowest packet.
// Returns -1 on recovery failure
int processWindowClosed(ConnectionData *client, Window *window, int *retryCount) {
	//if count is 10, we're done trying to recover.
	if ((*retryCount) > 9) {
		// Error'd out, go to DONE state
		printf("Deadlock recovery failed, aborting transfer\n");
		//cleanExit(window, client);
		return -1;
	}

	// try to poll at one second for a max of ten times to try and recv RR/SREJs, 
	// else enter recovery and send lowest packet
	if (pollCall(POLL_WAIT_1_SEC) != -1) {
		// there's a packet waiting to be processed, hopefully this will allow 
		processRRandSREJ(client, window);
		// Reset count, recv'd some type of data from client whether valid or not
		(*retryCount) = 0;

	} else {
		// entered recovery, send the packet pointed to by lower
		sendLowestPacket(client, window);
		(*retryCount)++;
	}

	return 1;
}

// Sends an EOF packet to a client
void sendEOF(ConnectionData *client, uint32_t *seqNum) {
	uint8_t *PDU;
	uint8_t payload = 0;
	PDU = createPDU(*seqNum, EOF_FLAG, &payload, 0);
	safeSendto(client->socket, PDU, HEADER_LEN, 0, (struct sockaddr *) &(client->address), client->len);
}

// send the lowest (RR) packet to client
void sendLowestPacket(ConnectionData *client, Window *window) {
	windowEntry we = getLowestEntry(*window);

	// TODO, MAKE A FUNCTION THAT WILL LOOP TO FIND THE LOWEST VALID PACKET
	if (!we.isValid) {
		printf("Oh no, the lowest packet was invalid. What to do???\n");
	}

	safeSendto(client->socket, we.PDU, we.packetLength, 0, (struct sockaddr *) &(client->address), client->len);
}

// Processes a single RR/SREJ
bool processRRandSREJ(ConnectionData *client, Window *window) {
	uint8_t packet[MAX_BUFFER] = {0};
	int recvLen = 0;

	recvLen = safeRecvfrom(client->socket, packet, MAX_BUFFER, 0, (struct sockaddr *) &((*client).address), &((*client).len));

	// check for CRC error or flags other than (RR, 5) or (SREJ, 6)
	if (!validChecksum(packet, recvLen)) {
		// discard the packet and continue looking for more
		return false;
	
	} else if (packet[FLAG_LOC] == RR_FLAG) {
		// RR recv'd, get RR seqNum and modify window
		processRR(window, packet);
	
	} else if (packet[FLAG_LOC] == SREJ_FLAG) {
		// SREJ recv'd, get the entry/pdu and send it back
		processSREJ(client, window, packet);
	}

	return true;
}

// Main processing function to poll for SREJs and RRs from a client when the window is open
bool processRRandSREJOpen(ConnectionData *client, Window *window) {
	int currentSocket = 0;

	// while RR/SREJS coming in, process
	while ((currentSocket = pollCall(0)) != -1) {
		
		if (!processRRandSREJ(client, window)) {
			return false;
		}
		
	}

	return true;
}

// Handles the sending of data packets (flag 3) to the connected client
void handleDataSending(ConnectionData *client, Window *window, uint32_t *seqNum, uint8_t *dataPayload, uint32_t bytesRead) {
	uint8_t *PDU;

	int dataLen = bytesRead + HEADER_LEN;

	PDU = createPDU(*seqNum, DATA_FLAG, dataPayload, bytesRead);

	// Now add this PDU to the window
	addPDUToWindow(window, PDU, dataLen);

	// Now send this PDU and increment the sequeunce number
	safeSendto(client->socket, PDU, dataLen, 0, (struct sockaddr *) &(client->address), client->len);

	(*seqNum)++;
}

// Used for the initial setup and filename exchange of the client.
STATE processFilenameEX(ConnectionData *client, uint8_t *packet, int32_t recvLen, ClientInfo *metadata) {
	FILE *dataFile = NULL;
	uint8_t fnameLen = 0;
	STATE nextState = DONE;

	// verify valid checksum, if invalid, quit client and wait for a good filename packet request
	if (in_cksum((unsigned short *) packet, recvLen) != 0) {
		nextState = DONE;
	}

	// get the metadata [header | winsize | buffersize | fnamelen | fname]
	memcpy(&(metadata->windowSize), &packet[7], sizeof(uint32_t));
	metadata->windowSize = ntohl(metadata->windowSize);
	memcpy(&(metadata->bufferSize), &packet[11], sizeof(uint16_t));
	metadata->bufferSize = ntohs(metadata->bufferSize);
	memcpy(&fnameLen, &packet[13], sizeof(uint8_t));
	memcpy(&(metadata->fileName), &packet[14], fnameLen);

	// Try to open the file specified for reading (determines next packet sending)
	if ((dataFile = fopen(metadata->fileName, "r")) == NULL) {
		// couldn't open the file, bad filename, send the response
		sendFilenameResponse(client, FNAME_BAD);
		nextState = DONE;

	} else {
		// Good file! send the response
		metadata->dataFile = dataFile;
		sendFilenameResponse(client, FNAME_GOOD);
		nextState = SEND_DATA;
	}

	return nextState;
}

// Forms and sends the filename response packet back to the connected client (flag 8)
void sendFilenameResponse (ConnectionData *client, uint8_t response) {
	// form the packet and send (seq num doesn't matter, flag=8, len=header+1)
	uint8_t *packet = createPDU(0, FNAME_RESPONSE_FLAG, &response, 1);
	
	safeSendto(client->socket, packet, HEADER_LEN+1, 0, (struct sockaddr *) &(client->address), client->len);
}

// Process the SREJ packet by getting the correct packet from the window and sending it to the requesting client
void processSREJ(ConnectionData *client, Window *window, uint8_t *packet) {
	windowEntry we;
	uint32_t srej = getOffsetSeqNumber(packet);
	// get the sequence number for SREJ and get that window entry
	we = handleSREJ(*window, srej);
	// send off the pdu 	
	safeSendto(client->socket, we.PDU, we.packetLength, 0, (struct sockaddr *) &(client->address), client->len);
}

// Process the RR by getting the RR number and sliding the window
void processRR(Window *window, uint8_t *packet) {
	uint32_t rr = getOffsetSeqNumber(packet);
	// update window
	handleRR(window, rr);
}

// Populates the client info on the first recv
int safeRecvBuf(int serverSocket, uint8_t *pkt, int len, ConnectionData *client, uint8_t *flag, uint32_t *seqNum) {
	(*client).len = sizeof(struct sockaddr_in6);
	return safeRecvfrom(serverSocket, pkt, len, *flag, (struct sockaddr *) &((*client).address), &((*client).len));
}

// Closes the socket created for the client, removes it from the pollset, and frees memory allocated for the window table
void cleanExit(Window *window, ConnectionData *client) {
	close(client->socket);
	removeFromPollSet(client->socket);
	// free(window->table);
	//exit(1);
}

// CL argument checking
int checkArgs(int argc, char *argv[])
{
	// Checks args and returns port number
	int portNumber = 0;

	if (argc > 3 || argc < 2)
	{
		fprintf(stderr, "Usage %s error-rate [optional port number]\n", argv[0]);
		exit(-1);
	}
	
	if (argc == 3)
	{
		portNumber = atoi(argv[2]);
	}
	
	return portNumber;
}

// Gets and converts the errRate string passed on the CL
double getRate(char * argv[]) {
	double rate = 0.0;
	rate = atof(argv[1]);

	// rate == 0 || ---- error rate can be zero
	if (rate > 1.0 || rate < 0) {
		printf("error-rate must be between 0-1\n");
		exit(1);
	}

	return rate;
}