/*
 * Copyright (c) 1990, 1999 Erick Engelke
 */
#include <rtos.h>
#include <net.h>
#include <stdio.h>
#include <ctype.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <syslog.h>
#include <tftp.h>
#include <sys\stat.h>

#define TFTP_LOCAL_PORT 10069
#define TFTP_PACKET_LEN 512 + 8

/*
 * get_tftp_response
 *      - if TIMEOUT return -1
 *      - sends multiple retries as needed in case data lost
 */
static int get_tftp_response( udp_Socket *out, udp_Socket *in, char *buffer, int *buflen )
{
#define TFTP_TIMEOUT 5 /* seconds */
#define TFTP_SLEEP 20  /* milliseconds */
#define TFTP_RETRIES 5

    int timeout;
    DWORD  timecount;
    int retrycount = TFTP_RETRIES;
    int templen;

    if ( out == NULL )
        out = in;           /* handle post-setup sockets */

    timeout = TFTP_TIMEOUT * TFTP_SLEEP;
    timecount = set_timeout( 1 );
    sock_fastwrite( out, buffer, *buflen );
    while  ((templen = sock_dataready( in )) == 0 ) {
        tcp_tick( NULL );
//        rt_sleep( TFTP_SLEEP );
//        timeout -= TFTP_SLEEP;
//        if ( timeout < 0 ) {
        rt_sleep( 0 );      // give up a bit of a slice
        if ( chk_timeout( timecount )) {
//            timeout = TFTP_TIMEOUT * TFTP_SLEEP;
            timeout = set_timeout( 1 );
            if ( retrycount -- < 1 )
                return( -1 );
            sock_fastwrite( out, buffer, *buflen );
        }
    }
    sock_fastread( in, buffer, templen );
    *buflen = templen;
    return( 0 );
}

// getblock - returns received length
static int getblock( char *filename, int blocknum, void *buffer )
{
    FILE *f;
    int result = 0;

    dos_enter();
    f = fopen( filename, "rb");
    if ( f != NULL ) {
        if ( 0 == fseek( f, (blocknum - 1) * 512L, SEEK_SET ))
            result = fread( buffer, 1, 512, f );
        fclose( f );
    }
    dos_exit();
    return( result );
}

// putblock - returns 0 on success
static int putblock( char *filename, int blocknum, void *buffer, int length )
{
    FILE *f;
    int result = 0;

    dos_enter();
    f = fopen( filename, "a+b");
    if ( f != NULL ) {
        if ( 0 == fseek( f, (blocknum - 1) * 512L, SEEK_SET ))
            fwrite( buffer, 1, length, f );
        else result = 3;    // disk full or alloc exceeded
        fclose( f );
    } else result = 3;
    dos_exit();

    return( result ? 3 : 0 );
}


int tftp( DWORD host, int op, char * file )
{
    udp_Socket *s, *s2;
    char *buf;
    int datalen;
    RRQ_Packet*t;
    DATA_Packet *d;
    ERROR_Packet *e;
    char *p;
    int result;
    int ongoing = 1;        /* means we are still doing stuff */
    int block;
    DWORD filesize;
    long bytesleft;
    struct stat statbuf;
    int i;

    s = kcalloc( sizeof( udp_Socket ), 1 );
    if ( s == NULL ) return( -1 );
    s2 = kcalloc( sizeof( udp_Socket ), 1 );
    if ( s2 == NULL ) {
        kfree(s);
        return( -1 );
    }
    buf = kcalloc( TFTP_PACKET_LEN , 1 );
    if ( buf == NULL ) {
        kfree( s );
        kfree( s2 );
        return( -1 );
    }
    d = t = e = buf;

    if ( op == TFTP_GET )
        unlink( file );
    else {
        dos_enter();
        i = stat( file , &statbuf );
        dos_exit();
        if ( i != 0 ) return -1 ;   /* file not found */
        filesize = statbuf.st_size;
        block = 0;
    }

    if (!udp_open( s, TFTP_LOCAL_PORT, host, TFTP_PORT, NULL )) {
        kfree( s );
        kfree( s2 );
        kfree( buf );
        return( -1 );
    }
    udp_open( s2, TFTP_LOCAL_PORT, 0, 0, NULL );

    // build the file request packet
    memset( buf, 0, TFTP_PACKET_LEN );
    t->op2 = ( op == TFTP_PUT ) ? TFTP_WRQ : TFTP_RRQ;
    strcpy( t->filename, file );
    p = strchr( t->filename, 0 );
    strcpy( ++p, "octet");
    p = strchr( p, 0 );
    datalen = p - t + 1;


// sock_write( s, buf, datalen );

    result = 0;

    while ( 1 ) {

        if ( get_tftp_response( s, s2, buf, &datalen) == -1 ) {
            result = -1;
            break;
        }

        /* we don't need original socket any more */
        if ( s != NULL ) {
            sock_close( s );
            kfree( s );
            s = NULL;
        }
        if ( !ongoing ) break;      /* after we have acked last packet, we quit */

        if ( d->op2 == TFTP_ERROR ) {
            result = -1;
            break;
        } else if (d->op2 == TFTP_DATA ) {
            /* TFTP_GET */
            datalen = datalen - sizeof( ACK_Packet );
            d->op2 = TFTP_ACK;
            block = intel16( d->block) ;

            putblock( file, block, d->data, datalen );

            /* 512 bytes means there is more to read */
            if ( datalen != 512 ) {
                sock_write( s2, buf, sizeof( ACK_Packet ));
                break;
            }
            datalen = sizeof( ACK_Packet );
        } else if ( d->op2 == TFTP_ACK ) {
            /* TFTP_PUT */
            if ( intel16( d->block) == block )
                block++;                /* advance, else repeat */

            bytesleft = filesize - (( block - 1) * 512L);
            if ( bytesleft <= 0 )   /* hit end of file */
                break;

            d->op2 = TFTP_DATA;
            d->block = intel16( block );

            /* try to read block from disk */
            i = getblock( file, block, d->data );
            if ( i == 0 ) {
                result = -1;
                break;
            }
            datalen = sizeof( ACK_Packet ) + i;

        } else {
            /* error */
            result = -1;
            ongoing = 0;
            e->op2 = TFTP_ERROR;
            e->err1 = e->err2 = 0;
            *e->errmsg = 0;
            datalen = sizeof( ERROR_Packet );
        }

    } 


    if ( s != NULL ) {
        sock_close( s );
        kfree( s );
    }
    sock_close( s2 );
    kfree( buf );
    kfree( s2 );
    return( result );
}


