#include "tftpserver.h"

struct CONTEXT *contextlist;

int main(int argc, char* argv[])
{	
	int sock;	
	struct sockaddr_in ServAddr;
	struct sockaddr_in ClientAddr;
	struct in_addr *clientaddr;
	unsigned int ClientAddrLen;
	unsigned short ServPort;
	unsigned short *opcode;
	int ret=0;
	unsigned char msg[600];
	int msgsize;

	contextlist = NULL;
	ServPort = atoi(argv[1]);	

	if((sock=socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP))<0)
	{
		printf("\n Error : socket() failed\n");
		exit(-1);
	}
	memset(&ServAddr, 0, sizeof(ServAddr));
	ServAddr.sin_family = AF_INET;	
	ServAddr.sin_addr.s_addr = htonl(INADDR_ANY);
	ServAddr.sin_port = htons(ServPort);

	if(bind(sock, (struct sockaddr*)&ServAddr, sizeof(ServAddr))<0)
	{
		printf("\n Error : bind() failed\n");
		exit(-1);
	}

	for(;;)
	{
		ClientAddrLen = sizeof(ClientAddr);
		memset(msg, 0, sizeof(msg));

		if( (msgsize = recvfrom(sock, msg, 600, 0, (struct sockaddr*)&ClientAddr, 
			&ClientAddrLen) ) < 0 )
		{
			printf("\n Error : recvfrom() failed !!\n");
			exit(-1);
		}

		opcode = (unsigned short*)msg;
		switch(ntohs(*opcode))
		{
			case RRQ_OPCODE :
				printf("\n recv RRQ[%d] \n", msgsize);
				ret = process_RRQ(sock, msg, msgsize, &ClientAddr, ClientAddrLen);
				break;
			case WRQ_OPCODE :
				ret = process_WRQ(sock, msg, msgsize, &ClientAddr, ClientAddrLen);
				break;
			case DATA_OPCODE :
				ret = process_DATA(sock, msg, msgsize, &ClientAddr, ClientAddrLen);
				break;
			case ACK_OPCODE :
				printf("\n recv ACK[%d] \n", msgsize);
				ret = process_ACK(sock, msg, msgsize, &ClientAddr, ClientAddrLen);
				break;
			case ERROR_OPCODE :
				ret = process_ERROR(sock, msg, msgsize, &ClientAddr, ClientAddrLen);
				break;
			default :
				printf("\n ERROR :  Receive Unknown OpCode Packet !! -> Ignore ");
				ret=1;
		}
		if( ret <0)
		{
			remove_all_context();
			close(sock);
			exit(-1);
		}
	}
	return 0;
}

int process_RRQ(int socket, char *msg, int recvsize, 
		struct sockaddr_in *clientaddr, int clientaddrsize)
{
	FILE *fp;
	struct RRQ *rrqpacket;
	struct DATA	*datapacket;
	char sendmsg[600];
	size_t readsize;
	int sendsize;
	struct CONTEXT *pContext;
	
	if( recvsize != sizeof(struct RRQ))
	{
		printf("\n Error : Bad RRQ size: %d ", recvsize);
		send_error(socket, clientaddr, clientaddrsize, 0, "Bad Packet Size");
		return 1;
	}

	rrqpacket = (struct RRQ *)msg;
	if(strcmp(rrqpacket->Mode, BIN) != 0)
	{
		printf("\n Error : Bad RRQ mode: %s field", rrqpacket->Mode);
		send_error(socket, clientaddr, clientaddrsize, 0, "Illegal Mode");
		return 1;
	}
	printf("\n process_RRQ : filename %s", rrqpacket->FileName);

	if((fp = fopen(rrqpacket->FileName, "r+b"))==NULL)
	{
		printf("\n Error : file[%s] not found", rrqpacket->FileName);
		send_error(socket, clientaddr, clientaddrsize, 1, "File not found");
		return 1;
	}
	
	// Context creation ...
	if( (pContext=(struct CONTEXT*)malloc(sizeof(struct CONTEXT))) == NULL)
	{
		printf("\n ERROR: malloc error");
		send_error(socket, clientaddr, clientaddrsize, 0, "malloc error");
		fclose(fp);
		return -1;
	}
	memset(pContext, 0, sizeof(struct CONTEXT));
	
	// Context initialize ...
	pContext->State = SEND_DATA_STATE;
	pContext->RecvNumber = 1;
	pContext->SendNumber = 1;
	pContext->CurrentByteSize = 0;
	
	// Send first DATA packet ..
	datapacket = (struct DATA *)sendmsg;
	datapacket->OpCode = htons(DATA_OPCODE);


	readsize = fread(datapacket->Data, 1, 512, fp);

	if( readsize <= 0)
	{
		printf("\n ERROR: fread() error");
		send_error(socket, clientaddr, clientaddrsize, 0, "fread error");
		fclose(fp);
		free(pContext);
		return -1;
	}
	datapacket->BlockNumber = htons(pContext->SendNumber);  // == 1
	sendsize = sendto(socket, sendmsg, sizeof(struct DATA)+readsize-1, 0, 
		(struct sockaddr*)clientaddr, clientaddrsize);
	printf("\n Send DATA[%d]  BlockNumber[%u]\n", sendsize, ntohs(datapacket->BlockNumber));

	// state processing ...
	if (readsize < 512)
		pContext->State = WAITING_CLOSE_STATE;
	else
		pContext->State = SEND_DATA_STATE;
	pContext->CurrentByteSize=readsize;
	strcpy(pContext->FileName, rrqpacket->FileName);
	pContext->fp = fp;
	memcpy( &(pContext->ClientAddr), clientaddr, sizeof(struct sockaddr_in) );

	// Context insertion ..
	pContext->next=contextlist;
	contextlist=pContext;

	return 1;
}

int process_ACK(int socket, char *msg, int recvsize, 
		struct sockaddr_in *clientaddr, int clientaddrsize)
{
	FILE *fp;
	struct ACK *ackpacket;
	struct DATA	*datapacket;
	char sendmsg[600];
	size_t readsize;
	int sendsize;
	struct CONTEXT *pContext;

	ackpacket = (struct ACK *)msg;
	pContext = find_context(clientaddr);

	if( pContext == NULL)
	{
		printf("\n Error : recv ACK but No Context.. ignore  ");
		send_error(socket, clientaddr, clientaddrsize, 0, "ACK but No Context");
		return 1;
	}
	
	if( recvsize != sizeof(struct ACK))
	{
		printf("\n Error : Bad ACK size: %d ", recvsize);
		send_error(socket, clientaddr, clientaddrsize, 0, "ACK but Bad Packet Size");
		remove_context(pContext);
		return 1;
	}

	if( pContext->State != SEND_DATA_STATE && pContext->State != WAITING_CLOSE_STATE)
	{
		printf("\n Error : ACK but bad State %d", pContext->State);
		send_error(socket, clientaddr, clientaddrsize, 0, "ACK but Bad State");
		remove_context(pContext);
		return 1;
	}
	if( pContext->SendNumber != ntohs(ackpacket->BlockNumber) )
	{
		printf("\n Error : recv ACK but bad Ack num %u", ntohs(ackpacket->BlockNumber));
		send_error(socket, clientaddr, clientaddrsize, 0, "ACK but Bad Block Number");
		remove_context(pContext);
		return 1;
	}

	// State == WAITING_CLOSE_STATE
	if ( pContext->State == WAITING_CLOSE_STATE)
	{
		printf("\n All Data Sending ...\n");
		fclose(pContext->fp);
		remove_context(pContext);
		return 1;
	}


	// Make and Send  DATA packet ..
	datapacket = (struct DATA *)sendmsg;
	datapacket->OpCode = htons(DATA_OPCODE);
	readsize = fread(datapacket->Data, 1, 512, pContext->fp);

	if( readsize <= 0)
	{
		printf("\n ERROR: fread() error");
		send_error(socket, clientaddr, clientaddrsize, 0, "fread error");
		fclose(pContext->fp);
		remove_context(pContext);
		return -1;
	}
	++(pContext->SendNumber);  // increase the send block counter ..
	datapacket->BlockNumber = htons(pContext->SendNumber);   
	sendsize = sendto(socket, sendmsg, sizeof(struct DATA)+readsize-1, 0, 
		(struct sockaddr*)clientaddr, clientaddrsize);
	printf("\n Send DATA[%d]  BlockNumber[%u]\n", sendsize, ntohs(datapacket->BlockNumber));

	// state processing ...
	if (readsize < 512)
		pContext->State = WAITING_CLOSE_STATE;
	else
		pContext->State = SEND_DATA_STATE;
	pContext->CurrentByteSize += readsize;
	return 1;
}

int process_DATA(int socket, char *msg, int recvsize, 
		struct sockaddr_in *clientaddr, int clientaddrsize)
{
	printf("\n Not Implemented yet.. ");
	return 1;
}

int process_WRQ(int socket, char *msg, int recvsize, 
		struct sockaddr_in *clientaddr, int clientaddrsize)
{
	printf("\n Not Implemented yet.. ");
	return 1;
}

int process_ERROR(int socket, char *msg, int recvsize, 
		struct sockaddr_in *clientaddr, int clientaddrsize)
{
	struct ERROR *errorpacket;

	if( recvsize < sizeof(struct ERROR))
	{
		printf("\n Error Packet : Bad ERROR size: %d ", recvsize);
		return -1;
	}
	errorpacket = (struct ERROR *)msg;

	printf("\n ERROR packet recv  ErrorNumber: %d  ErrorMsg : %s", 
		ntohs(errorpacket->OpCode), errorpacket->ErrorData);

	return -1;

}

struct CONTEXT *find_context(struct sockaddr_in *clientaddr)
{
	struct CONTEXT *pContext;

	pContext=contextlist;
	while(pContext)
	{
		if( !memcmp(&(pContext->ClientAddr), clientaddr, sizeof(struct sockaddr_in)) )
			return pContext;
		pContext=pContext->next;
	}
	return NULL;
}

int remove_context(struct CONTEXT *pContext)
{
	struct CONTEXT *p, *q;

	p=q=contextlist;

	while(p)
	{
		if( p==pContext)
		{
			//first context ..
			if( p==q)
			{
				contextlist=p->next;
				free(p);
				return 1;
			}
			else
			{
				q->next=p->next;
				free(p);
				return 1;
			}
		}
		q=p;
		p=p->next;
	}
	return 0;
}

int remove_all_context(void)
{
	struct CONTEXT *p, *temp;

	if(!contextlist)
		return 0;

	p=contextlist;
	while(p)
	{
		temp=p;
		p=p->next;
		free(temp);
	}
	contextlist=NULL;
	return 0;
}

int send_error(int socket, struct sockaddr_in *clientaddr, int clientaddrsize, 
			unsigned short errorcode, char *errormsg)
{
	char sendmsg[128];
	struct ERROR *errorpacket;
	int	sendsize;

	errorpacket = (struct ERROR *)sendmsg;
	errorpacket->OpCode = htons(ERROR_OPCODE);
	errorpacket->ErrorNumber = htons(errorcode); // file not found ..
	strcpy(errorpacket->ErrorData, errormsg);

	sendsize = sendto(socket, errorpacket, sizeof(struct ERROR)+strlen(errormsg)-1, 0,
		(struct sockaddr*)clientaddr, clientaddrsize);
	
	printf("\n Send ERROR[%d]  OpCode[%u] ErrorNum[%u] ErrorMsg[%s]\n", sendsize, 
		ntohs(errorpacket->OpCode), ntohs(errorpacket->ErrorNumber), errorpacket->ErrorData);

	return 1;
}