// rcopy - client side			    
// By Connor McKee

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/uio.h>
#include <sys/time.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <strings.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>

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

#define MAXBUF 80
#define ACK_EOF 1
#define NO_ACK_EOF 0
#define GOT_EOF 1
#define NO_EOF 0


// rcopy states
typedef enum STATE {
	INIT, DONE, FILENAME, FILE_OK, RECV_DATA, EOF_ACK
} STATE;

// Information grabbed from CL arguments for use throughout state machine
typedef struct {
	char *fromFilename;
	char *toFilename;
	FILE *toFile;
	uint32_t windowSize;
	uint16_t bufferSize;
	double errRate;
} rcopyInfo;

void talkToServer(int socketNum, struct sockaddr_in6 * server);
STATE establishConnection (ConnectionData *server, rcopyInfo metadata, uint32_t *clientSeqNum);
void processCopy(int socket, rcopyInfo metadata, struct sockaddr_in6 server);
void sendFilenameRequest(int socket, struct sockaddr_in6 *server, uint32_t sequenceNumber, uint8_t *fileNamePayload, uint8_t fnameLen);
STATE processFilenameResponse(ConnectionData *server, rcopyInfo metadata);
STATE processFilenameRetry(ConnectionData *server, int *retries);
STATE processRecvData(ConnectionData *server, rcopyInfo *metadata, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum, Window *window);
STATE processOpenFile (ConnectionData *server, rcopyInfo *metadata);
STATE processEOFACK(ConnectionData *server, rcopyInfo *metadata, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum, Window *window);
void writeDataToFile(ConnectionData *server, rcopyInfo *metadata, uint8_t *PDU, int recvLen);
int flushBuffer(ConnectionData *server, rcopyInfo *metadata, Window *window, uint32_t *expectedSeqNum);
int processSREJ(ConnectionData *server, rcopyInfo *metadata, Window *window, int acknowledgeEOF, int gotEOF, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum);
void sendDataRR(ConnectionData *server, Window *window, uint32_t *clientSeqNum, uint32_t seqNum);
void sendSREJ(ConnectionData *server, uint32_t *clientSeqNum, uint32_t expectedSeqNum);
void sendEOFACK(ConnectionData *server, uint32_t *clientSeqNum);
void cleanExit(Window *window, rcopyInfo *metadata);
int readFromStdin(char * buffer);
int checkArgs(int argc, char * argv[], rcopyInfo *metadata);
double getRate(char * argv[]);


int main (int argc, char *argv[])
 {
	int socketNum = 0;				
	struct sockaddr_in6 server;		// Supports 4 and 6 but requires IPv6 struct
	int portNumber = 0;
	rcopyInfo metadata;
	
	setupPollSet();

	portNumber = checkArgs(argc, argv, &metadata);
	
	sendErr_init(metadata.errRate, DROP_ON, FLIP_ON, DEBUG_ON, RSEED_ON);
	
	socketNum = setupUdpClientToServer(&server, argv[6], portNumber);
	
	addToPollSet(socketNum);

	processCopy(socketNum, metadata, server);
	
	close(socketNum);

	return 0;
}

void processCopy(int socket, rcopyInfo metadata, struct sockaddr_in6 serverAddr) {
	ConnectionData server;
	server.address = serverAddr;
	server.socket = socket;
	uint32_t clientSeqNum = 0;
	uint32_t expectedSeqNum = 0;
	uint32_t eofSeqNum = 0;
	//Window window;
	Window window = createWindow(metadata.windowSize);
	STATE state = INIT;

	while (state != DONE) {
		switch (state) {
			case INIT:
				state = establishConnection(&server, metadata, &clientSeqNum);
				break;

			case FILENAME:
				state = processFilenameResponse(&server, metadata);
				break;

			case FILE_OK:
				state = processOpenFile(&server, &metadata);
				break;

			case RECV_DATA:
				state = processRecvData(&server, &metadata, &clientSeqNum, &expectedSeqNum, &eofSeqNum, &window);
				break;

			case EOF_ACK:
				state = processEOFACK(&server, &metadata, &clientSeqNum, &expectedSeqNum, &eofSeqNum, &window);
				break;

			case DONE:
				// need to cleanly exit (free the table)
				printf("Finished the file transfer. Exiting.\n");
				cleanExit(&window, &metadata);
				exit(0);
				break;

			default:
				printf("in default!! Shouldn't be here\n");
				break;

		}
	}
}


STATE processEOFACK(ConnectionData *server, rcopyInfo *metadata, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum, Window *window) {
	STATE nextState = DONE;

	// if the expected seq num is the same as the eof seq num, done! send the EOF ACK, otherwise, there's some missing data
	if (*expectedSeqNum == *eofSeqNum) {
		// send the EOF ACK 
		sendEOFACK(server, clientSeqNum);
		return DONE;
	}

	// since there's missing data, start processing the SREJs needed to correct the buffer. break when you find a seqnum == eofseqnum
	// Since the EOF was already seen, we don't care about getting a repeat EOF. Do not acknowledge subsequent EOFs sent
	processSREJ(server, metadata, window, NO_ACK_EOF, GOT_EOF, clientSeqNum, expectedSeqNum, eofSeqNum);

	// Done processing and fixing data, MAY NEED TO FLUSH THE BUFFER ONCE MORE IN CASE OF GOOD DATA ABOVE THE FINAL HOLE

	// Really done now, send the EOF ACK and cleanly exit
	sendEOFACK(server, clientSeqNum);

	return nextState;
}



STATE processRecvData(ConnectionData *server, rcopyInfo *metadata, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum, Window *window) {
	int recvLen = 0;
	uint32_t currentSeqNum = 0;
	int holeLoc = 0;
	uint8_t packet[MAX_BUFFER] = {0};
	STATE nextState = DONE;
	// loop to process data packets
	while (1) {
		if (pollCall(POLL_WAIT_10_SEC) != -1) {
			// there's some data to be received
			memset(packet, 0, MAX_BUFFER);
			recvLen = safeRecvfrom(server->socket, packet, MAX_BUFFER, 0, (struct sockaddr *) (&(server->address)), &(server->len));
			
			// first, ensure CRC validity. if not valid, discard and wait for more data
			if (!validChecksum(packet, recvLen)) {
				continue;
			}

			// Get the current packet's seqNum
			currentSeqNum = getSeqNumber(packet);

			// verify correct flag type
			// If EOF received, go to EOF ack state to cleanup writing data to file
			if(packet[FLAG_LOC] == EOF_FLAG) {
				// Set the EOF sequence number
				(*eofSeqNum) = currentSeqNum;
				nextState = EOF_ACK;
				break;
			} else if(packet[FLAG_LOC] != DATA_FLAG) {
				// discard and try receive more data
				continue;
			}

			// MAIN DATA PROCESSING
			if (currentSeqNum == *expectedSeqNum) {
				// Data is as it should be! 
				writeDataToFile(server, metadata, packet, recvLen);

				// before RR'ing, check to see if anything higher is in buffer. Write anything valid to file and update expected
				// loop to check holes, (any more valid packets above)
				(*expectedSeqNum)++;
				if ((holeLoc = flushBuffer(server, metadata, window, expectedSeqNum)) < 0) {
					// No holes found, expected sequence number was updated to reflect where we now are
				} else {
					// Hole found. This is the expected sequence number
					(*expectedSeqNum) = holeLoc;
				}

				sendDataRR(server, window, clientSeqNum, *expectedSeqNum);
				
			} else if(currentSeqNum < *expectedSeqNum) {
				// Got some old data that is already written to disk. RR the expectedSeqNum
				sendDataRR(server, window, clientSeqNum, *expectedSeqNum);
				
			} else {
				// currentSeq is bigger, we are missing some data! Need to send out SREJs for all missing data
				// Buffer this data, go to processing SREJS and flushing the buffer
				addPDUToWindow(window, packet, recvLen);

				// process SREJs, acknowledge the EOF if found and break to go into EOF ack state
				if (processSREJ(server, metadata, window, ACK_EOF, NO_EOF, clientSeqNum, expectedSeqNum, eofSeqNum)) {
					// Found an EOF when trying to process SREJs. eofSeqNum was set
					nextState = EOF_ACK;
					break;
				}
			}
		
		} else {
			// No data has been sent by the server for ten seconds. exit.
			printf("Data poll timed out. Server terminated\n");
			nextState = DONE;
			break;
		}
	}

	return nextState;
}


// returns an indication if EOF was found in the middle of trying to process SREJs
int processSREJ(ConnectionData *server, rcopyInfo *metadata, Window *window, int acknowledgeEOF, int gotEOF, uint32_t *clientSeqNum, uint32_t *expectedSeqNum, uint32_t *eofSeqNum) {
	int recvLen = 0;
	int holeLoc = 0;
	uint32_t currentSeqNum = 0;
	uint8_t packet[MAX_BUFFER] = {0};

	// Only send SREJ once
	sendSREJ(server, clientSeqNum, *expectedSeqNum);

	while(1) {

		// Process recv
		if (pollCall(POLL_WAIT_10_SEC) != -1) {
			memset(packet, 0, MAX_BUFFER);
			recvLen = safeRecvfrom(server->socket, packet, MAX_BUFFER, 0, (struct sockaddr *) (&(server->address)), &(server->len));

			// Check for CRC validity
			if (!validChecksum(packet, recvLen)) {
				continue;
			}

			currentSeqNum = getSeqNumber(packet);

			// Need to check flags
			if(packet[FLAG_LOC] == EOF_FLAG) {
				// if user set to acknowledge, set and return now. Else, keep processing SREJs
				if (acknowledgeEOF) {
					(*eofSeqNum) = currentSeqNum;
					return 1;
				} else {
					// if we got the expected eof then we're done!
					if (*expectedSeqNum == *eofSeqNum) {
						return 1;
					} else {
						sendSREJ(server, clientSeqNum, *expectedSeqNum);
					}

					continue;
				}

			} else if(packet[FLAG_LOC] != DATA_FLAG) {
				// discard and try receive more data
				continue;
			
			} 

			if (currentSeqNum > *expectedSeqNum) {
				// buffer data 
				addPDUToWindow(window, packet, recvLen);

			} else if (currentSeqNum < *expectedSeqNum) {
				// Send an RR for expected
				sendDataRR(server, window, clientSeqNum, *expectedSeqNum);
			
			} else {
				writeDataToFile(server, metadata, packet, recvLen);
				// got expected, write and flush data above now. Check for holes. no holes -> send RR for new expected and exit. update expected and loop else
				(*expectedSeqNum)++;
				if ((holeLoc = flushBuffer(server, metadata, window, expectedSeqNum)) < 0) {
					// No holes found, expected sequence number was updated to reflect where we now are
					sendDataRR(server, window, clientSeqNum, *expectedSeqNum);
					break;
				} else {
					// Hole found. This is the expected sequence number
					*expectedSeqNum = holeLoc;
					// send SREJ
					sendSREJ(server, clientSeqNum, *expectedSeqNum);
				}
			}

		} else {
			printf("Server is unresponsive. Terminating.\n");
			cleanExit(window, metadata);
			exit(1);
		}
	}

	// EOF not found while processing
	return 0;
}


void sendEOFACK(ConnectionData *server, uint32_t *clientSeqNum) {
	uint8_t payload = 0;

	uint8_t *PDU = createPDU(*clientSeqNum, EOF_ACK_FLAG, &payload, 0);

	safeSendto(server->socket, PDU, HEADER_LEN, 0, (struct sockaddr *) (&(server->address)), server->len);

	(*clientSeqNum)++;
}

void sendSREJ(ConnectionData *server, uint32_t *clientSeqNum, uint32_t expectedSeqNum) {
	uint8_t payload[MAX_BUFFER] = {0};
	expectedSeqNum = htonl(expectedSeqNum);
	memcpy(payload, &expectedSeqNum, sizeof(uint32_t));

	uint8_t *PDU = createPDU(*clientSeqNum, SREJ_FLAG, payload, sizeof(uint32_t));

	safeSendto(server->socket, PDU, sizeof(uint32_t)+HEADER_LEN, 0, (struct sockaddr *) (&(server->address)), server->len);

	(*clientSeqNum)++;
}

// will return the location of a hole in the buffer
int flushBuffer(ConnectionData *server, rcopyInfo *metadata, Window *window, uint32_t *expectedSeqNum) {
	int i = *expectedSeqNum;
	int end = (*expectedSeqNum) + metadata->windowSize + 1;
	uint8_t localPacket[MAX_BUFFER];
	int localSize = 0;
	int index = 0;
	int hole = -1;

	// loop through window buffer to find good data. Go from expected around the circle back to starting pos.
	// if an entry is invalid, break, this is the new expected. If entry is valid, write it to disk and continue through buffer.
	for(; i < end; i++) {
		index = i%(metadata->windowSize);

		// Found an invalid entry in table (hole)
		if (!(window->table)[i%(metadata->windowSize)].isValid) {
			hole = i;
			break;
		}

		// copy the window entry pdu data to a local variable before writing (MEMORY BUG AVOIDANCE)
		memset(localPacket, 0, MAX_BUFFER);
		copyPDULocal(window, localPacket, &localSize, index);

		writeDataToFile(server, metadata, localPacket, localSize);

		markEntryInvalid(window, index);

		(*expectedSeqNum)++;
	}

	return hole;
}


// Sends an RR packet for the specified seqNum
void sendDataRR(ConnectionData *server, Window *window, uint32_t *clientSeqNum, uint32_t seqNum) {
	uint8_t payload[MAX_BUFFER] = {0};
	seqNum = htonl(seqNum);
	memcpy(payload, &seqNum, sizeof(uint32_t));

	uint8_t *PDU = createPDU(*clientSeqNum, RR_FLAG, payload, sizeof(uint32_t));

	safeSendto(server->socket, PDU, sizeof(uint32_t)+HEADER_LEN, 0, (struct sockaddr *) (&(server->address)), server->len);

	(*clientSeqNum)++;
}

// Write's a data packet payload to disk
void writeDataToFile(ConnectionData *server, rcopyInfo *metadata, uint8_t *PDU, int recvLen) {
	uint8_t data[MAX_BUFFER] = {0};
	// get the data to write to file
	memcpy(data, &PDU[HEADER_LEN], recvLen-HEADER_LEN);

	// Write to file
	if (fwrite(data, sizeof(uint8_t), recvLen-HEADER_LEN, metadata->toFile) == 0) {
		//printf("An error occurred while writing to the local file. Exiting\n");
		perror("Error occurred while trying to write to the local file: ");
		exit(1);
	}
}

// Attempts to open the local file specified for writing
STATE processOpenFile (ConnectionData *server, rcopyInfo *metadata) {
	FILE *toFile = NULL;
	STATE nextState = DONE;

	if((toFile = fopen(metadata->toFilename, "w")) == NULL) {
		// Error opening/creating file for writing
		printf("Error on open for output of file: <%s> -> ", metadata->toFilename);
		fflush(stdout);
		perror("");
		nextState = DONE;

	} else {
		// File has been opened for writing data to
		metadata->toFile = toFile;
		nextState = RECV_DATA;
	}

	return nextState;
}

// Processes the server's response to the filename exchange packet
STATE processFilenameResponse(ConnectionData *server, rcopyInfo metadata) {
	// process the reponse from the server. 
	// return to init if server does not reply, done if bad fname, file ok other
	static int retries = 0;
	uint8_t packet[MAX_BUFFER] = {0};
	int recvLen = 0;
	STATE nextState = INIT;

	if ((nextState = processFilenameRetry(server, &retries)) == FILE_OK) {
		// receive the filename response, server is trying to send a packet
		recvLen = safeRecvfrom(server->socket, packet, MAX_BUFFER, 0, (struct sockaddr *) (&(server->address)), &(server->len));

		// Check for CRC error
		if (in_cksum((unsigned short *) packet, recvLen) != 0) {
			nextState = INIT;

		// filename response lost, got a data packet
		} else if (packet[FLAG_LOC] == DATA_FLAG) {
			nextState = FILE_OK;

		// Check for bad filename response
		} else if (packet[HEADER_LEN] == FNAME_BAD) {
			printf("Error: file <%s> not found on server.\n", metadata.fromFilename);
			nextState = DONE;
		} 

	}

	return nextState;
}

// Tries to process the initial filename exchange packet up to ten times
STATE processFilenameRetry(ConnectionData *server, int *retries) {
	// calls poll with one second timeout (10 times) waiting for the file ok response from server
	// if timeout, we're headed back to init to try and resend the filename packet
	STATE nextState = INIT;
	int currentSocket = 0;
	
	// Check somewhere if the retry counter has exceeded the 10 tries? maybe check this in processfilenamereponse

	if ((currentSocket = pollCall(POLL_WAIT_1_SEC)) != -1) {
		// receiving a response, check that the server validated the requested filename
		nextState = FILE_OK;
	} else {
		// poll timed out after 1 second, increment the counter and return to INIT
		(*retries)++;
		nextState = INIT;
	}

	// if retries > 9, nextstate = done, server terminated or not valid
	if ((*retries) > 9) {
		nextState = DONE;
	}

	return nextState;
}


STATE establishConnection (ConnectionData *server, rcopyInfo metadata, uint32_t *clientSeqNum) {
	// attempts to establish connection to a server. Successfully will return the FILENAME next state
	server -> len = sizeof(struct sockaddr_in6);
	uint8_t *PDU;
	int payloadLen = 0;
	uint8_t filenamePayload[MAX_BUFFER];
	uint32_t windowSize = htonl(metadata.windowSize);
	uint16_t bufferSize = htons(metadata.bufferSize);
	uint8_t filenameLen = strlen(metadata.fromFilename) + 1;

	// form the payload before attaching the Header
	memcpy(filenamePayload, &windowSize, sizeof(uint32_t));
	memcpy(&filenamePayload[4], &bufferSize, sizeof(uint16_t));
	memcpy(&filenamePayload[6], &filenameLen, sizeof(uint8_t));
	memcpy(&filenamePayload[7], metadata.fromFilename, filenameLen);

	payloadLen = sizeof(uint32_t) + sizeof(uint16_t) + sizeof(uint8_t) + filenameLen;

	PDU = createPDU(*clientSeqNum, FNAME_REQUEST_FLAG, filenamePayload, payloadLen);

	safeSendto(server->socket, PDU, payloadLen+HEADER_LEN, 0, (struct sockaddr *) (&(server->address)), server->len);

	return FILENAME;
}

// Sends the initial filename request packet to the server
void sendFilenameRequest(int socket, struct sockaddr_in6 *server, uint32_t sequenceNumber, uint8_t *fileNamePayload, uint8_t fnameLen) {
	uint8_t *PDU;
	int serverAddrLen = sizeof(struct sockaddr_in6);
	int dataLen = fnameLen + 1;

	PDU = createPDU(sequenceNumber, FNAME_REQUEST_FLAG, fileNamePayload, dataLen);
	
	safeSendto(socket, PDU, dataLen + HEADER_LEN, 0, (struct sockaddr *) server, serverAddrLen);
}

// Error checks and converts the rate passed through CL
double getRate(char * argv[]) {
	double rate = 0.0;
	rate = atof(argv[5]);

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

	return rate;
}

// Checks the CL Args
int checkArgs(int argc, char * argv[], rcopyInfo *metadata)
{
	/* check command line arguments  */
	if (argc != 8)
	{
		printf("usage: %s from-filename to-filename window-size buffer-size error-rate remote-machine remote-port\n", argv[0]);
		exit(1);
	}
	
	// Do argument error checking here
	// Checks args, fills in metadata, and returns port number
	if (strlen(argv[1]) > (MAX_FNAME_LEN-1)) {
		printf("filename %s is too long. Exiting\n", argv[1]);
		exit(1);
	}

	if (strlen(argv[2]) > (MAX_FNAME_LEN-1)) {
		printf("filename %s is too long. Exiting\n", argv[2]);
		exit(1);
	}

	if (atoi(argv[4]) > 1400 || atoi(argv[4]) == 0) {
		printf("Buffer size invalid. Buffer size must be a number 1400 or less.\n");
		exit(1);
	}

	// Error check and get the error rate
	metadata -> errRate = getRate(argv);

	// Add the rest of the data to metadata
	metadata -> fromFilename = argv[1];
	metadata -> toFilename = argv[2];
	metadata -> windowSize = (uint32_t) atoi(argv[3]);
	metadata -> bufferSize = (uint16_t) atoi(argv[4]);
	
	int portNumber = 0;
	
	portNumber = atoi(argv[7]);
		
	return portNumber;
}

// Closes the file stream and frees memory allocated for the buffer
void cleanExit(Window *window, rcopyInfo *metadata) {
	// close the file stream
	fclose(metadata->toFile);
	free(window->table);
}