#include <stdio.h>
#include <string.h>

#include "lwip/inet.h"
#include "lwip/err.h"
#include "lwip/udp.h"
#include "lwip/mem.h"

#include "tftpd.h"
#include "tftputils.h"

tftp_connection_args gArgs;

/* tftp_errorcode error strings */
char *tftp_errorcode_string[] =
    {
        "not defined",
        "file not found in SPI Flash",
        "command not supported access violation",
        "check the SPI interface or disk full",
        "illegal operation",
        "unknown transfer id",
        "file already exists in SPI Flash",
        "no such user",
    };

err_t tftp_send_message (struct udp_pcb *upcb,
		                 struct ip_addr *to_ip,
		                 int to_port,
		                 char *buf,
		                 int buflen)
{

    err_t err;
    struct pbuf *pkt_buf;

    pkt_buf = pbuf_alloc(PBUF_TRANSPORT, buflen, PBUF_POOL);
    if (!pkt_buf)
    {
        iprintf("error allocating pkt buffer\n\r");
        return ERR_MEM;
    }
    // Copy the original data buffer over to the packet buffer's payload
    memcpy(pkt_buf->payload, buf, buflen);
    /* send message */
    err = udp_sendto(upcb, pkt_buf, to_ip, to_port);
    pbuf_free(pkt_buf);
    return err;
}

/* construct an error message into buf using err as the error code */
int tftp_construct_error_message(char *buf,
		                         tftp_errorcode err)
{
    int errorlen;

    tftp_set_opcode(buf, TFTP_ERROR);
    tftp_set_errorcode(buf, err);

    tftp_set_errormsg(buf, tftp_errorcode_string[err]);
    errorlen = strlen(tftp_errorcode_string[err]);

    /* return message size */
    return 4 + errorlen + 1;
}

/* construct and send an error message back to client */
int tftp_send_error_message(struct udp_pcb *upcb,
		                    struct ip_addr *to,
		                    int to_port,
		                    tftp_errorcode err)
{
    char buf[512];
    int n;

    n = tftp_construct_error_message(buf, err);
    return tftp_send_message(upcb, to, to_port, buf, n);
}

/* construct and send a data packet */
int tftp_send_data_packet(struct udp_pcb *upcb,
		                  struct ip_addr *to,
		                  int to_port,
		                  int block,
		                  char *buf,
		                  int buflen)
{
    char packet[TFTP_DATA_PKT_LEN_MAX];
  
    tftp_set_opcode(packet, TFTP_DATA);
    tftp_set_block(packet, block);
    tftp_set_data_message(packet, buf, buflen);

    return tftp_send_message(upcb, to, to_port, packet, buflen + 4);
}

int tftp_send_ack_packet(struct udp_pcb *upcb,
		                 struct ip_addr *to,
		                 int to_port,
		                 int block)
{
    char packet[TFTP_ACK_PKT_LEN]; // create the maximum possible size packet that a TFTP ACK packet can be

    tftp_set_opcode(packet, TFTP_ACK); // define the first two bytes of the packet
    tftp_set_block(packet, block);
    return tftp_send_message(upcb, to, to_port, packet, TFTP_ACK_PKT_LEN);
}

void tftp_cleanup(struct udp_pcb *upcb,
		          tftp_connection_args *args)
{
    /* cleanup the args */
    f_close(&args->fd);

    /* close the connection */
    udp_remove(upcb);
}

void tftp_send_next_block(struct udp_pcb *upcb,
		                  tftp_connection_args *args,
                          struct ip_addr *to_ip,
                          u16_t to_port)
{
    int res, s2;
    res = f_read(&(args->fd), args->data, sizeof(args->data), (unsigned int *)&s2);
    if(res)
    {
	    iprintf("Can not read from file\n");
		/* we are done */
		return tftp_cleanup(upcb, args);
    }
    args->data_len -= s2;
    if (args->data_len <= 0)
    {
	    iprintf("closing connection, ret = %d\n\r", args->data_len);
		/* we are done */
		return tftp_cleanup(upcb, args);
    }
    /* send the data */
    tftp_send_data_packet(upcb, to_ip, to_port, args->block, args->data, args->data_len);
}



/* write callback */
void wrq_recv_callback(void *_args,
		               struct udp_pcb *upcb,
		               struct pbuf *pkt_buf,
		               struct ip_addr *addr,
		               u16_t port)
{
    unsigned int res, s2;
    int flag_connect = 0;
    tftp_connection_args *args = (tftp_connection_args *)_args;

    // Does this packet have any valid data to write?
    if ((pkt_buf->len > TFTP_DATA_PKT_HDR_LEN) &&
        (tftp_extract_block(pkt_buf->payload) == (args->block + 1)))
    {
//Todo f_seek	  
	    res = f_write(&(args->fd), (char *)pkt_buf->payload+TFTP_DATA_PKT_HDR_LEN, pkt_buf->tot_len-TFTP_DATA_PKT_HDR_LEN, &s2);
	    if (res)
	    {
            iprintf("write to file error\n\r");
            tftp_send_error_message(upcb, addr, port, TFTP_ERR_FILE_NOT_FOUND);
                        pbuf_free(pkt_buf);
            return tftp_cleanup(upcb, args);
        }
        //iprintf("%s",(char *)pkt_buf->payload+TFTP_DATA_PKT_HDR_LEN);
        args->block++;
        (args->tot_bytes)+=(pkt_buf->tot_len-TFTP_DATA_PKT_HDR_LEN);
        flag_connect = 1;

    }
#if 0
    else if(tftp_extract_block(pkt_buf->payload) == (args->block + 1))
    {
        args->block++;
    }
#endif

    tftp_send_ack_packet(upcb, addr, port, args->block);


    if( (pkt_buf->tot_len < TFTP_DATA_PKT_LEN_MAX) &&
    		(flag_connect))
    {
        pbuf_free(pkt_buf);
	
	    res = f_close(&(args->fd));
	    if(res)
	    {
		    iprintf("Can't close the file\n");
	    }
        return tftp_cleanup(upcb, args);
    }
    else
    {
        pbuf_free(pkt_buf);
        return;
    }
}

int tftp_process_write(struct udp_pcb *upcb,
		               struct ip_addr *to,
		               int to_port,
		               char *fname)
{
    FIL fd;
    unsigned int res;
    tftp_connection_args *args;
    args = &gArgs;

    res = f_open(&fd,fname, FA_OPEN_EXISTING);
    f_close(&fd);
    if(!res)
    {
        f_unlink(fname);
	  //      iprintf("file %s already exists: May overwrite\n", fname);
    }
    res = f_open(&fd, fname, FA_CREATE_ALWAYS | FA_WRITE);
    if (res)
    {
        iprintf("unable to open file %s for writing\n", fname);
        tftp_send_error_message(upcb, to, to_port, TFTP_ERR_DISKFULL);
        udp_remove(upcb);
        return 1;
    }
    args->op = TFTP_WRQ;
    args->fd = fd;
    args->to_ip.addr = to->addr;
    args->to_port = to_port;
    args->block = 0;
    args->tot_bytes = 0;
    args->data_len = 600;

    /* set callback for receives on this UDP PCB (Protocol Control Block) */
    udp_recv(upcb, wrq_recv_callback, args);

    /* initiate the write transaction by sending the first ack */
    tftp_send_ack_packet(upcb, to, to_port, args->block);

    return 0;
}

void process_tftp_request(struct pbuf *pkt_buf,
		                  struct ip_addr *addr,
		                  u16_t port)
{
    tftp_opcode op = tftp_decode_op(pkt_buf->payload);
    char fname[512];
    struct udp_pcb *upcb;
    err_t err;

    /* create new UDP PCB structure */
    upcb = udp_new();
    if (!upcb)
    {
        iprintf("Error creating PCB. Out of Memory\n\r");
        return;
    }

    /* bind to port 0 to receive next available free port */
    // NOTE!  This is how TFTP works.  There is a UDP PCB for the standard port
    // 69 which al transactions begin communication on, however, _all_ subsequent
    // transactions for a given "stream" occur on another port!
    err = udp_bind(upcb, IP_ADDR_ANY, 0);
    if (err != ERR_OK)
    {
        iprintf("Unable to bind to port %d: err = %d\n\r", port, err);
        return;
    }

    switch (op)
    {
        case TFTP_WRQ:
            tftp_extract_filename(fname, pkt_buf->payload);
            tftp_process_write(upcb, addr, port, fname);
        break;
        default:
            /* send a generic access violation message */
            tftp_send_error_message(upcb, addr, port, TFTP_ERR_ACCESS_VIOLATION);
            iprintf("TFTP unknown request op: %d\n", op);
            udp_remove(upcb); // no need to use tftp_cleanup because no "tftp_connection_args" struct has been malloc'd
        break;
    }
    return;
}

void recv_callback(void *arg,
		           struct udp_pcb *upcb,
		           struct pbuf *pkt_buf,
		           struct ip_addr *addr,
		           u16_t port)
{
    /* process new connection request */
    process_tftp_request(pkt_buf, addr, port);

    pbuf_free(pkt_buf);
}

int tftpd_init()
{
    struct udp_pcb *upcb;
    err_t err;
    unsigned port = 69;

    /* create new UDP PCB structure */
    upcb = udp_new();
    if (!upcb)
    {
        iprintf("Error creating PCB. Out of Memory\n\r");
        return -1;
    }

    /* Bind this PCB to port 69 */
    err = udp_bind(upcb, IP_ADDR_ANY, port);
    if (err != ERR_OK)
    {
        iprintf("Unable to bind to port %d: err = %d\n\r", port, err);
        return -2;
    }

    udp_recv(upcb, recv_callback, NULL);

    return 0;
}

