//File name : TftpClt.cpp
//
#include <Winsock2.h>
#include <stdio.h>
#include <string.h>
#include <conio.h>
#include <stdlib.h>
#include "DStruct.h"

BOOL ErrorCmd(int argc,char *argv[]);	//測試命令行,打印幫助信息
void displayHelp();						//displayHelp 打印幫助信息函數
void SetErrMsg(u_short m_ErrorNum);     //SetErrMsg 設置出錯信息函數
int  InitSock(SOCKET m_socket, SOCKADDR_IN addrTo, char **argv);

void main(int argc, char *argv[])
{
	int Repeat = 0;
	int j,TotalBuf=0,nBytes=512;
	SOCKET m_socket = 0;
	SOCKADDR_IN addrTo,addrFrom;
	int len =sizeof(SOCKADDR_IN);

	//初始化階段
	memset(&addrTo, 0 ,sizeof(addrTo));
	if(InitSock(m_socket, addrTo, argv));
	{
		return;
	}
	if(ErrorCmd(argc,argv) == TRUE)
		return;
	//printf("%s\n",argv[2]);
	RWRQ *RwRq = new RWRQ;
//	char right[3];
//==================打開文件,設置RwRq==============//
	FILE *pFile;
	if(strcmp(argv[3],"GET") == 0)
	{	
		RwRq->pcode = htons(RRQ);
		if(NULL == (pFile = fopen(argv[5],"wb")))
		{
			SetErrMsg(2);
			return;
		}
		strcpy(RwRq->fileandmode, argv[4]);
	}
	else if(strcmp(argv[3],"PUT") == 0)
	{
		RwRq->pcode = htons(WRQ);
		if(NULL == (pFile = fopen(argv[4],"rb")))
		{
			SetErrMsg(u_short(1));
			printf("\t%s\n",argv[4]);
			return;
		}
		strcpy(RwRq->fileandmode, argv[5]);
	}

	printf("Conneting %s...\n",inet_ntoa(addrTo.sin_addr));

	strcpy(&(RwRq->fileandmode[strlen(RwRq->fileandmode)+1]),"octet");
	//發送請求包
	sendto(m_socket,(char *)RwRq,strlen(RwRq->fileandmode)+strlen("octet")+6,
		0,(SOCKADDR *)&addrTo,len);
	
	DataBuf *DBuf = new DataBuf;
	AckBuf  *ABuf = new AckBuf;

	DBuf->pcode = htons(DAT);
	ABuf->pcode = htons(ACK);
	u_short Dblock=0;
	u_short Ablock=0;

	int TimeOut = 0;
	fd_set ReadSet;
	timeval t;
	t.tv_sec = 5;
	t.tv_usec = 0;
	ErrorBuf *EBuf = new ErrorBuf;
	
	char buffer[1024];
	u_long TimeCount = ::GetTickCount();
//--------------------------------處理階段-----------------------------------------------
	while(1)
	{
		FD_ZERO(&ReadSet);
		FD_SET(m_socket,&ReadSet);

		if(select(m_socket,&ReadSet,NULL,NULL,&t))
		{
			Repeat=0;
			/*--------------RRQ----------------*/
			if(RwRq->pcode == htons(RRQ))
			{		
	
				nBytes = recvfrom(m_socket,(char*)buffer,sizeof(buffer),
					0,(SOCKADDR*)&addrFrom,&len);
				buffer[nBytes] = '\0';	
						
				if((u_short)++Dblock == (u_short)ntohs(((DataBuf*)buffer)->block) 
					&& ntohs(((DataBuf*)buffer)->pcode) == DAT)
				{
					if(TotalBuf == 0)
						printf("Dowloading %s ...\n",argv[4]);
					TimeOut = 0;
					TotalBuf += nBytes - 4;
					ABuf->block = ((DataBuf*)buffer)->block;
					sendto(m_socket,(char*)ABuf,sizeof(ABuf),0,(SOCKADDR*)&addrFrom,len);
				
					
					if(nBytes < 512)
					{
						TimeCount = u_long((::GetTickCount() - TimeCount)/1000);
						if(TimeCount == 0) TimeCount = 1;
						FD_ZERO(&ReadSet);
						FD_SET(m_socket,&ReadSet);

						if(select(m_socket,&ReadSet,NULL,NULL,&t))
						{ //5秒內還收到DATA說明最后一個ACK主機沒收到，重發一次
							nBytes = recvfrom(m_socket,(char*)buffer,sizeof(buffer),
								0,(SOCKADDR*)&addrFrom,&len);
							sendto(m_socket,(char*)ABuf,sizeof(ABuf),
								0,(SOCKADDR*)&addrFrom,len);
						}
						else
						{	//超過5秒認為ACK已經被主機收到
							printf("Transfer Successful: %d Bytes in %ld seconds,%ld Bytes/s\n",
								TotalBuf,TimeCount,u_long(TotalBuf/TimeCount));
						}

						for(j=0;j<nBytes-4;j++)
							fputc(int(((DataBuf*)buffer)->buf[j]),pFile);
						fclose(pFile);
						///printf("waiting for DBuf %d\n",Dblock);
						return;
					}
					continue;
				}
				else
				{
					Dblock--;
					
					if(((DataBuf*)buffer)->pcode == htons(ERR))
					{
					//	SetErrMsg(ntohs(((ErrorBuf*)buffer)->ErrorCode));
						printf("Server Error : %s\n", ((ErrorBuf*)buffer)->ErrMsg);
						return;
					}
					
					if(Dblock == 0)
					{
						if(TimeOut == 5)
						{
							TimeOut = 0;
							printf("Connect Repeat 5 times,Server Lost !\n");
							return;
						}
						sendto(m_socket,(char *)RwRq,sizeof(RWRQ),0,(SOCKADDR *)&addrTo,len);
					}
					else
					{
						ABuf->block = Dblock;
						sendto(m_socket,(char*)ABuf,sizeof(ABuf),0,(SOCKADDR*)&addrFrom,len);
					}
					if(TimeOut == 5) 
					{
						TimeOut = 0;
						printf("ACK Repeat 5 times,Server Lost !\n");
						return;
					}
					++TimeOut;
					continue;
				}//Else
			}
			/*--------------WRQ------------------上傳*/
			if(RwRq->pcode == ntohs(WRQ))
			{
				memset(buffer,0,sizeof(buffer));
				recvfrom(m_socket,(char*)buffer,sizeof(buffer),
					0,(SOCKADDR*)&addrFrom,&len);
					
				if(((AckBuf*)buffer)->pcode == htons(ACK) && 
					((AckBuf*)buffer)->block == htons(Ablock))
				{
					if(TotalBuf == 0 && nBytes != 0)
						printf("Uploading %s ...\n",argv[4]);

					TimeOut = 0;
					
					if(nBytes<512) 
					{
						TimeCount = u_long((::GetTickCount() - TimeCount)/1000);
						if(TimeCount == 0) TimeCount = 1;
						printf("Transfer Successful: %ld Bytes in %ld seconds,%ld Bytes/s\n",
							TotalBuf,TimeCount,u_long(TotalBuf/TimeCount));
						fclose(pFile);
						return;
					}
					((DataBuf*)DBuf)->block = htons(++Ablock);
					memset(DBuf->buf,0,sizeof(DBuf->buf));
					nBytes = 0;
					for(j=0;j<512 ;j++)
					{	
						((DataBuf*)DBuf)->buf[j] = fgetc(pFile);
						if(feof(pFile))
							break;	
						nBytes++;
					//	printf("%d\n",nBytes);
					}
					//printf("af");
					sendto(m_socket,(char*)DBuf,nBytes+4,
						0,(SOCKADDR*)&addrFrom,len);
				//	printf("DBuf->block : %d \tnBytes : %d\n",ntohs(DBuf->block),nBytes);
					
					TotalBuf += nBytes;
					continue;
				}
				else
				{
					if(((ErrorBuf*)buffer)->pcode == ntohs(ERR))
					{
						SetErrMsg(ntohs(((ErrorBuf*)buffer)->ErrorCode));
						return;
					}
					TimeOut++;
					if(TimeOut == 5) 
					{
						printf("DATA Repeat 5 times,Server Lost !\n");
						TimeOut = 0;
						return;
					}
					if(Ablock == 0)
						sendto(m_socket,(char *)RwRq,sizeof(RWRQ),0,(SOCKADDR *)&addrTo,len);
					else
					{
						sendto(m_socket,(char *)DBuf,sizeof(DBuf),0,(SOCKADDR *)&addrFrom,len);
					}
					continue;
				}
			}//if(RwRq->pcode = WRQ)
		}//if(select(m_socket,&m_sockset,NULL,NULL,&t))
		else//
		{	
			Repeat++;
			
			if(Repeat == 5)
			{//select 超時五次則認為與服務器斷開了
				printf("Time Out !  Repeat %d Times,Host Lost !\n",Repeat);
				Repeat = 0;
				return;
			}
			printf("Time Out !  Repeat %d Times\n",Repeat);
		}
	}//while(1)
	return;
}
//測試命令行,打印幫助信息
BOOL ErrorCmd(int argc,char **argv)
{
	for(int i=3;i<argc;i++)
	{
		if(i == 4) break;
		char *p = argv[i];
		for( ; p < argv[i] + strlen(argv[i]); p++ )
		{
			if( islower( *p ) )
				*p = _toupper( *p );
		}
	}
	if((argc< 6) || 0 != strcmp(argv[1],"-i") || (0 != strcmp(argv[3],"GET")
		&& 0 != strcmp(argv[3],"PUT")) )
	{
		displayHelp();
		return TRUE;
	}
	return FALSE;
}
//displayHelp 打印幫助信息函數
void displayHelp()
{
	printf("Transfers files to and from a remote computer running the TFTP service.\n");

	printf("TFTPClt [-i] host [GET | PUT] source [destination]\n");

  	printf("-i              Specifies binary image transfer mode (also called\n");
    printf("                octet). In binary image mode the file is moved\n");
    printf("                literally, byte by byte. Use this mode when\n");
    printf("                transferring binary files.\n");
  	printf("host            Specifies the local or remote host.\n");
	printf("GET             Transfers the file destination on the remote host to\n");
    printf("                the file source on the local host.\n");
	printf("PUT             Transfers the file source on the local host to\n");
    printf("                the file destination on the remote host.\n");
    printf("source          Specifies the file to transfer.\n");
	printf("destination     Specifies where to transfer the file.\n");
}
//SetErrMsg 設置出錯信息函數
void SetErrMsg(u_short m_ErrorNum)
{

	switch(m_ErrorNum)
	{
	case 0:
		printf("Not defined, see error message (if any).\n");		
		break;
	case 1:
		printf("Local Error: File not found.");
		break;
	case 2:
		printf("Access violation.\n");
		break;
	case 3:
		printf("Disk full or allocation exceeded.\n");
		break;
	case 4:
		printf("Illegal TFTP operation mode.\n");
		break;
	case 5:
		printf("Unknown transfer ID.\n");
		break;
	case 6:
		printf("Server Error: File already exists.\n");
		break;
	case 7:
		printf("No such user.\n");
		break;
	}
}

int InitSock(SOCKET m_socket, SOCKADDR_IN addrTo, char ** argv)
{
	DWORD wVersionRequested;
	WSADATA wsaData;
	int err;
 
	wVersionRequested = MAKEWORD( 1, 1 );
 
	err = WSAStartup( (u_short)wVersionRequested, &wsaData );
	if ( err != 0 ) {
		return err;
	} 

	if ( LOBYTE( wsaData.wVersion ) != 1 ||
			HIBYTE( wsaData.wVersion ) != 1 ) {
		WSACleanup( );
		return 1; 
	}

	m_socket = socket(AF_INET,SOCK_DGRAM,0);
    if(m_socket ==SOCKET_ERROR)
	{
		printf("SOCKET INITIAL ERROR\n");
		return m_socket;
	}
   
	addrTo.sin_addr.S_un.S_addr   = inet_addr(argv[2]);
	addrTo.sin_family = AF_INET;
	addrTo.sin_port   = htons(69);
	return 0;
}