import java.io.*;
import java.net.*;
import java.nio.ByteBuffer;

public class TFTPServer {
    public static final int TFTPPORT = 4970;
    public static final int BUFSIZE = 516;
    public static final String READDIR = "/home/username/read/"; //custom address at your PC
    public static final String WRITEDIR = "/home/username/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;
    public static String Mode;
    public static final int OP_ERR0 = 0;
    public static String ERR0 = "Not defined, see error message(if any)";
    public static final int OP_ERR1 = 1;
    public static String ERR1 = "File not found.";
    public static final int OP_ERR2 = 2;
    public static String ERR2 = "Access Violation";
    public static final int OP_ERR3 = 3;
    public static String ERR3 = "Disk full of allocation exceeded";
    public static final int OP_ERR4 = 4;
    public static String ERR4 = "Illegal TFTP operation";
    public static final int OP_ERR5 = 5;
    public static String ERR5 = "Unknown transfer ID";
    public static final int OP_ERR6 = 6;
    public static String ERR6 = "File already exists";

    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();
        }
    }

    private void start() throws SocketException {
        byte[] buf = new byte[BUFSIZE];

        // Create socket
        DatagramSocket socket = new DatagramSocket(null);

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

        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);

                        System.out.printf("%s request for %s from %s using port %d\n",
                                (reqtype == OP_RRQ) ? "Read" : "Write",
                                clientAddress.getHostName(), clientAddress.getPort());

                        // Read request
                        if (reqtype == OP_RRQ) {
                            requestedFile.insert(0, READDIR);
                            HandleRQ(sendSocket, requestedFile.toString(), OP_RRQ);
                        }
                        // Write request
                        else {
                            requestedFile.insert(0, WRITEDIR);
                            HandleRQ(sendSocket, requestedFile.toString(), OP_WRQ);
                        }
                        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 receivePacket = new DatagramPacket(buf, buf.length);

        // Receive packet
        try {
            socket.receive(receivePacket);
        } catch (IOException e) {
            System.out.println("Error: " + e);
        }

        // Get client address and port from the packet
        InetSocketAddress socketAddress = new InetSocketAddress(receivePacket.getAddress(), receivePacket.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
        // byte[] buf;
        ByteBuffer wrap = ByteBuffer.wrap(buf);
        short opcode = wrap.getShort();
        String file = requestedFile.toString();

        file = file.substring(2, requestedFile.length());
        String[] parts = file.split("0");
        String fileName = (parts[0]);

        String Mode = parts[1];
        System.out.println("Opcode is: " + opcode);
        System.out.println("Requested file: " + fileName);
        System.out.println("Transfer mode is: " + Mode);
        return 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) {
        System.out.println("Requested file: " + requestedFile);
        File file = new File(requestedFile);

        if (opcode == OP_RRQ) {
            // See "TFTP Formats" in TFTP specification for the DATA and ACK packet contents

            FileInputStream in = null;

            try {
                in = new FileInputStream(requestedFile);

                if (send_DATA_receive_ACK(sendSocket, in)) {
                    System.out.println("Sending...");
                    //System.out.println("Successfully sent. BlockNum is: " + blockNum);
                } else {
                    System.out.println("ERROR: Connection lost");
                    send_ERR(sendSocket, OP_ERR0, "Connection to server lost");
                }
            } catch (FileNotFoundException f) {
                System.out.println("ERROR, FILE NOT FOUND" + f);
                send_ERR(sendSocket, OP_ERR1, ERR1);
            }
        }
        else if (opcode == OP_WRQ) {
            if (file.exists()) {
                System.out.println("Error, file already exists");
                send_ERR(sendSocket, OP_ERR6, ERR6);
            } else {

                short blockNum = 0;
                ByteBuffer buffer = ByteBuffer.allocate(BUFSIZE);
                buffer.putInt(OP_ACK);
                buffer.putShort(blockNum);
                DatagramPacket ackP = new DatagramPacket(buffer.array(), 4);

                boolean result = receive_DATA_send_ACK(sendSocket, ackP, blockNum);

                while (result) {
                    try {
                        FileOutputStream out = new FileOutputStream(file);
                    } catch (FileNotFoundException e) {
                        send_ERR(sendSocket, OP_ERR3, "Could not create file.");

                    }
                }
            }
        }
        else {
                System.err.println("Invalid request. Sending an error packet.");
                // See "TFTP Formats" in TFTP specification for the ERROR packet contents
                send_ERR(sendSocket, opcode, "Invalid request. Sending an error packet.");
            }
    }

    /**
     * To be implemented
     */
    private boolean send_DATA_receive_ACK(DatagramSocket sendSocket, FileInputStream in ) {
        short blockNum = 1;
        byte[] buf = new byte[512];
        int length=0;
        try
        {length = in.read(buf);}
        catch(IOException e)
        {System.out.println("ERROR IOE "+e);}

        ByteBuffer wrap = ByteBuffer.wrap(buf);
        short shortVal = OP_DAT;

        // put the opcode and block number in the bytebuffer
        wrap.putShort(shortVal);
        wrap.putShort(blockNum);
        wrap.put(buf, 0, length);
        byte[] wrapArr =wrap.array();

        int waitTimeForAck = 10; //staring with 10 milliseconds

        //The Send DATA part
        DatagramPacket sender = new DatagramPacket(wrapArr, 4+length);


        byte[] ACKreceived = new byte[4]; //The ACK received is 2 byte (opt) + 2 byte (block) = 4
        DatagramPacket receivedACKPacket = new DatagramPacket(ACKreceived, ACKreceived.length);
        try {
            sendSocket.send(sender);
            sendSocket.setSoTimeout(waitTimeForAck);
        }
        catch(IOException e)
        {
            System.out.println("ERROR SENDING PACKET "+e);
        }

        try{

            sendSocket.receive(receivedACKPacket);} catch(IOException e){
            System.err.println("No acknowledgment packet received!");
            e.printStackTrace();
            return false;
        }

        return true;
    }

    private boolean receive_DATA_send_ACK(DatagramSocket sendSocket, DatagramPacket ackP, short block) {

        byte[]rec=new byte[BUFSIZE];
       DatagramPacket recD= new DatagramPacket(rec, rec.length);
        try {
            sendSocket.receive(recD);
            ByteBuffer buffer = ByteBuffer.wrap(recD.getData());
            int opcode = buffer.getInt();
            short blockNum = buffer.getShort();
            System.out.println("Block recieved: "+ blockNum +" Block sent?: "+ block);
            sendSocket.send(ackP);
           // out.write(recD.getData(), 4, recD.length - 4);

           } catch (FileNotFoundException f) {
               System.out.println("ERROR, file could not be created" + f);
               send_ERR(sendSocket, OP_ERR0, "ERROR: FIle could not be created");
           }
           catch(IOException e) {
               System.out.println("Error in sending Ack: "+e);

           }
        return true;
        }



    private void send_ERR(DatagramSocket sendSocket, int opcode, String msg)
    {
        ByteBuffer err = ByteBuffer.allocate(BUFSIZE);
        err.put((byte) 0);
        err.putInt(OP_ERR);
        err.putInt(opcode);
        err.put((byte) 0);
        err.put(msg.getBytes());
        err.put((byte) 0);
        byte[] error=err.array();
        DatagramPacket errorP = new DatagramPacket(error, error.length);
        try {
            sendSocket.send(errorP);
        } catch (IOException e) {
            System.err.println("ERROR WITH ERROR: Error packet did not send");
        }

    }

}