import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;

public class TFTPServer 
{
	private static final String HOME_DIR = System.getProperty("user.home");


	public static final int TFTPPORT = 4970;
	public static final int BUFSIZE = 516;
	public static final int DATASIZE = 512;
	public static final String READDIR =  HOME_DIR + "/tftp/read/"; //custom address at your PC
	public static final String WRITEDIR = HOME_DIR + "/tftp/write/"; //custom address at your PC
	// OP codes
	public static final int OP_RRQ = 1;
	public static final int OP_WRQ = 2;
	public static final int OP_DAT = 3;
	public static final int OP_ACK = 4;
	public static final int OP_ERR = 5;
	// Error codes
	public static final short ERR_NDE = 0;
	public static final short ERR_FNF = 1;
	public static final short ERR_ACV = 2;
	public static final short ERR_ITO = 4;
	public static final short ERR_FAE = 6;

	public static void main(String[] args) {
		if (args.length < 0)
		{
			System.err.printf("usage: java %s\n", TFTPServer.class.getCanonicalName());
			System.exit(1);
		}
		//Starting the server
		try 
		{
			TFTPServer server= new TFTPServer();
			server.start();
		}
		catch (SocketException e) 
			{e.printStackTrace();}
		catch (IOException e)
			{e.printStackTrace();}
	}
	
	private void start() throws IOException 
	{
		byte[] buf= new byte[BUFSIZE];

		// Set timeouts for the session
		TimeOutThread.setTimeouts();
		
		// Create socket
		DatagramSocket socket= new DatagramSocket(null);
		

		// Create local bind point 
		SocketAddress localBindPoint= new InetSocketAddress(TFTPPORT);
		socket.bind(localBindPoint);


		// Create directories for read and write if non-existent
		Files.createDirectories(Paths.get(READDIR));
		Files.createDirectories(Paths.get(WRITEDIR));


		// Print out of status
		System.out.printf("Listening at port %d for new requests\n", TFTPPORT);


		// Loop to handle client requests 
		while (true) 
		{
			final InetSocketAddress clientAddress = receiveFrom(socket, buf);
			

			// If clientAddress is null, an error occurred in receiveFrom()
			if (clientAddress == null) 
				continue;

			final StringBuffer requestedFile= new StringBuffer();
			final int reqtype = ParseRQ(buf, requestedFile);

			new Thread() 
			{
				public void run() 
				{
					try 
					{
						DatagramSocket sendSocket= new DatagramSocket(0);

						// Connect to client
						sendSocket.connect(clientAddress);						

						
						// Add printout to confirm client connection and type.
						System.out.printf("%s request for %s from %s using port %d\n",
								(reqtype == OP_RRQ)?"Read":"Write",
								clientAddress.getHostName(),
								clientAddress.getAddress(),
								clientAddress.getPort());


						// Read request
						if (reqtype == OP_RRQ) 
						{
							requestedFile.insert(0, READDIR);
							HandleRQ(sendSocket, requestedFile.toString(), OP_RRQ);
						}


						// Write request
						else if(reqtype == OP_WRQ)
						{                       
							requestedFile.insert(0, WRITEDIR);
							HandleRQ(sendSocket,requestedFile.toString(),OP_WRQ);  
						}
						else
						{
							send_ERR("Malformed request", ERR_ITO, sendSocket);
						}


						// Close socket to conclude the connection
						sendSocket.close();

					}
					catch (SocketException e) 
						{e.printStackTrace();}
				}
			}.start();
		}
	}
	




	/***********************************************************************************
	 * Reads the first block of data, i.e., the request for an action (read or write). *
	 * @param socket (socket to read from)																						 *
	 * @param buf (where to store the read data)																		   *
	 * @return socketAddress (the socket address of the client)												 *
	 ***********************************************************************************/
	private InetSocketAddress receiveFrom(DatagramSocket socket, byte[] buf) 
	{
		// Create datagram packet
		DatagramPacket dp = new DatagramPacket(buf, BUFSIZE);


		// Receive packet
		try
		{socket.receive(dp);}
		catch (IOException e)
		{e.printStackTrace();}

		// Get client address and port from the packet
		InetSocketAddress socketAddress = new InetSocketAddress(dp.getAddress(), dp.getPort());

		return socketAddress;
	}



	
	/********************************************************************************
	 * Parses the request in buf to retrieve the type of request and requestedFile. *
	 * 																																							*
	 * @param buf (received request)																								*
	 * @param requestedFile (name of file to read/write)														*
	 * @return opcode (request type: RRQ or WRQ)																		*
	 ********************************************************************************/
	private int ParseRQ(byte[] buf, StringBuffer requestedFile) 
	{
		// See "TFTP Formats" in TFTP specification for the RRQ/WRQ request contents
		ByteBuffer bb = ByteBuffer.wrap(buf);
		short opcode = bb.getShort();

		// Keep adding the next byte (as char) to str buffer until null encountered.
		byte nxtByt;
		while((nxtByt = bb.get()) != (byte)0) {
			requestedFile.append((char)nxtByt);	
		}

		// TODO - Here is a good spot for handling Mode
		
		return (int)opcode;
	}




	/************************************************************
	 * Handles RRQ and WRQ requests 														*
	 * 																													*
	 * @param sendSocket (socket used to send/receive packets)	*
	 * @param requestedFile (name of file to read/write)				*
	 * @param opcode (RRQ or WRQ)																*
	 ************************************************************/
	private void HandleRQ(DatagramSocket sendSocket, String requestedFile, int opcode) 
	{
		// See "TFTP Formats" in TFTP specification for the DATA and ACK packet contents
		if(opcode == OP_RRQ)
		{
			try 
			{ 
				// Create the abstract file
				File fileToSend = new File(requestedFile);

				// Check if valid file
				if(Files.notExists(fileToSend.toPath())){
					throw new NullPointerException();
				}
				
				if (!fileToSend.canRead()){
					throw new IllegalAccessError();
				}
				

				// Proceed if a real file. Loop for all blocks until end of file.
				AtomicInteger blockNo = new AtomicInteger(1);
				while(send_DATA_receive_ACK(fileToSend, sendSocket, blockNo)) {
					blockNo.incrementAndGet();
				};
			}
			catch (IllegalAccessError ia) {
				// Access violation error needs to be sent
				System.out.println("Illegal access attempted!");
				send_ERR("Access violation", ERR_ACV, sendSocket);
			}
			catch (NullPointerException e) 
			{
				// File not found error needs to be sent
				send_ERR("File not found", ERR_FNF, sendSocket);
			}
			

			
		}
		else if (opcode == OP_WRQ) 
		{
			// Check if file already exsists - no overwriting allowed
			if (Files.exists(Paths.get(requestedFile)))
			{
				send_ERR("File already exists!", ERR_FAE, sendSocket);
			}
			else
			{
				try (FileOutputStream fileStream = new FileOutputStream(new File(requestedFile))){
					
					
					// Create vars
					AtomicInteger retrans = new AtomicInteger(0);
					AtomicInteger blockNo = new AtomicInteger(1);
					Ack firstAck = new Ack((short)OP_ACK, (short)0);


					// Send the first ACK and keep receiving until the end of file.
					sendSocket.send(new DatagramPacket(firstAck.getAckArray(), firstAck.getAckArray().length));
					while(receive_DATA_send_ACK(fileStream, sendSocket, blockNo, retrans)){
						blockNo.incrementAndGet();
					};


				} catch (IOException e) {
					send_ERR("Something went wrong. Could not receive the file.", ERR_NDE, sendSocket);
				}
				
			}
		}
		else 
		{
			System.err.println("Invalid request. Sending an error packet.");
			// See "TFTP Formats" in TFTP specification for the ERROR packet contents
			send_ERR("Cannot make out the request.", ERR_NDE, sendSocket);
			return;
		}		
	}

	


	/**************************************************
	 * Send Datagrams and handle ACKs from the client	*
	 * 																								*
	 * @param fileToSend - a java.nio.file File				*
	 * @param socket - data socket										*
	 * @return - boolean to nofify end of file				*
	 **************************************************/
	private boolean send_DATA_receive_ACK(File fileToSend, DatagramSocket socket, AtomicInteger blockNo)
	{
		ByteBuffer datagramBuff;

		// Write
		try{
			// Set timeout to prevent infinite blocking while receiving.
			socket.setSoTimeout(TimeOutThread.READRCV);


			// Declare vars for the session
			byte[] datagram;
			int writtenBytes = 0;
			long offset = (blockNo.get() - 1) * DATASIZE;
			int nxtByte;


			// Check file size and prepare block with maximum 512 bytes 
			// of data per block, according to RFC1350 hence 511.
			if (fileToSend.length() > (DATASIZE - 1))
			{
				datagram = new byte[BUFSIZE];
				datagramBuff = ByteBuffer.wrap(datagram);
			}
			else
			{
				datagram = new byte[(4 + (int)fileToSend.length())];
				datagramBuff = ByteBuffer.wrap(datagram);
			}


			// Write the headers for a packet and add 4 bits to written bytes (2xshort)
			datagramBuff.putShort((short)OP_DAT);
			datagramBuff.putShort((short)blockNo.get());
			writtenBytes += Ack.ACKSIZE;
			

			// Keep reading the file until end of stream is reached or buffer is full
			// Start is offset depending on block number.
			FileInputStream fileStream = new FileInputStream(fileToSend);
			fileStream.skip(offset);
			while((nxtByte = fileStream.read()) != -1 && writtenBytes < BUFSIZE) {
				datagramBuff.put((byte)nxtByte);
				writtenBytes ++;
			}
			fileStream.close();


			// Time-out thread to prevent getting stuck in infinite retransmission
			Thread timeOut = new Thread(new TimeOutThread());


			
			// ACK response from the client
			Ack ack = new Ack();
			

			// Start the timer and attempt to send data and receive an ACK
			timeOut.start();
			do{
				socket.send(new DatagramPacket(datagram, writtenBytes));
				try
				{socket.receive(new DatagramPacket(ack.getAckArray(), 4));}
				catch (SocketTimeoutException ste)
				{System.out.println("ACK timeout. Trying again...");}

				// Check if acknowlegement has correct opcode and block number
				ack.setAckOp();
				ack.setAckBlock();
				if(ack.getAckOp() == (short)OP_ACK && ack.getAckBlock() == blockNo.get()){break;}
			}


			// Retransmit until the package is acknowledged or the server times out.
			// while(ack.getAckOp() != (short)OP_ACK);
			while(ack.getAckOp() != (short)OP_ACK && timeOut.isAlive());


			// Malformed ack handler.
			if (ack.getAckOp() != (short)OP_ACK || ack.getAckBlock() != blockNo.get()){
				if (timeOut.isAlive()) 
				{throw new IllegalArgumentException();}
				else 
				{throw new TimeoutException();}
				
			}


			// Check if more packets need to be sent
			if(nxtByte == -1) {
				System.out.printf("%s sent successfully\n", fileToSend.getName());
				blockNo.set(0);
				return false;
			} else {
				return true;
			}

		} catch (TimeoutException te) { 
			// Handle timeout.
			String error = "Client unresponsive. Time out reached";
			System.out.println(error);
			send_ERR(error, ERR_NDE, socket);
			return false;

		} catch (Exception e) {
			// Handle errors from socket sending,reading a file stream, or malformed ACK.
			// Allow the stream to continue. Terminal output disabled due to annoyin output.

			// String error = "Malformed acknowlegement. Cannot send the block.";
			// System.out.println(error);
			blockNo.decrementAndGet();
			return true;
		}
	}
	



	// Temp addition of string to make this work
	private boolean receive_DATA_send_ACK(FileOutputStream fileToRec, DatagramSocket sock, AtomicInteger blockNo, AtomicInteger retrans)
	{
		byte[] fileBlock = new byte [BUFSIZE];
		Ack ack = new Ack((short)OP_ACK, (short)blockNo.get());
		DatagramPacket packet = new DatagramPacket(fileBlock, BUFSIZE);
		

		
		try {
			// Receive the data and wrap it in Datagram Packet
			sock.setSoTimeout(TimeOutThread.WRITERCV);
			sock.receive(packet);


			// Reset retrans for the next packet
			retrans.set(0);
			// retrans = Integer.valueOf(TimeOutThread.MAXTRANS);


			//Confirm the packet has correct OP code and continue transmission.
			if(ByteBuffer.wrap(fileBlock).getShort() != (short)OP_DAT){
				send_ERR("Incorrect Data OP code", ERR_ITO, sock);
				return true;
			}


			// Check if the block number is correct
			if(ByteBuffer.wrap(fileBlock).getShort(2) != (short)blockNo.get()){
				return true;
			}


			// Write bytes from file buffer to file stream.
			// Only writes as many bytes as were sent minus the header.
			fileToRec.write(fileBlock, Ack.ACKSIZE, (packet.getLength() - Ack.ACKSIZE));


			// Send the ACK
			sock.send(new DatagramPacket(ack.getAckArray(), ack.getAckArray().length));
			

			// Check if this was the last packet - indicated by size < 512bytes
			// Reset block counter to accept new files from the same source.
			if(packet.getLength() < BUFSIZE) {
				blockNo.set(1);
				return false;
			}
		} 
		catch (SocketTimeoutException ste){
			// Ensure block number correctness
			blockNo.decrementAndGet();
			System.out.printf("Retransmission number: %d\n",retrans.incrementAndGet());

			// Only allow MAXTRANS number of retransmissions
			// per packet to ensure no infinite loops are not created.
			// Total wait per packet is MAXTRANS * WRITERCV (30s default)
			// Send Error (Opcode 5) as a curtesy if timeout reached.
			if(retrans.get() < TimeOutThread.MAXTRANS)
			{ return true;}
			else 
			{
				send_ERR("Maximum number of retransmissions reached", ERR_NDE, sock);
				return false;
			}
		}
		catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
			return false;
		}
		return true;
	}
	

	// Prepare an error datagram and send to socket
	private void send_ERR(String errMsg, short ercode, DatagramSocket sock)
	{

		// Create the array with 5 bytes constant + variable error message
		ByteBuffer errorBuffer = ByteBuffer.wrap(new byte[ (5 + errMsg.length()) ]);
		
		// Populate the array with values according to RFC1350
		// 2 bytes  2 bytes   string   1 byte
		//|  05   | errCode | errMsg |   0   |
		errorBuffer.putShort((short)OP_ERR);
		errorBuffer.putShort(ercode);
		errorBuffer.put(errMsg.getBytes());
		errorBuffer.put((byte)0);


		// Send error to sock
		try {
			sock.send(new DatagramPacket(errorBuffer.array(), errorBuffer.array().length));
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
}