#ifndef _TFTP2_H
#define _TFTP2_H

#define WINDOWS_OS

#include <iostream>
#include <fstream>
#include <string>
#include "udp.h"
//#include "udp_stream.h"
#include "bsstream.h"
#ifdef WINDOWS_OS
#include <windows.h>
#include <winsock.h>
#else
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#endif

namespace tftp_tools {

using namespace std;
using namespace net_tools;

class tftp {
    string server;
    int port;

    int read_buffer(ifstream & fin, char * buffer, size_t size) {
       int ret = 0;
       char ch = fin.get();
       if(fin) {
          for(size_t i = 0; i < size; i++) {
             buffer[i] = ch;
             ret++;
             if((i+1) < size)
                ch = fin.get();
             if(!fin) 
                break;
          }
       }
       return ret;
    }

public:
   tftp() { }

   tftp(const string & server, const int port = 69) {
      this->server = server;
      this->port = port;
   }

   tftp(const tftp & cp) {
      server = cp.server;
      port = cp.port;
   }

   ~tftp() {
   }

   tftp & operator=(const tftp & cp) {
      server = cp.server;
      port = cp.port;
      return *this;
   }

   bool get(const string & file_name) {
      char buf[516];
      bool ret = false;
      udp_socket udp(server, port, udp_socket::UDP_CLIENT);
      memset(buf, 0, sizeof(buf));  
      *((unsigned short*)buf) = htons((u_short)1);
      //sprintf(buf+2, "%s%c%s%c", file_name.c_str(), '\0', "netascii", '\0');
      sprintf(buf+2, "%s%c%s%c", file_name.c_str(), '\0', "octet", '\0');
      //int n = udp.send(buf, 2+file_name.size()+1+8+1); 
      int n = udp.send(buf, 2+file_name.size()+1+5+1); 
      ofstream fout(file_name.c_str(), ios::out | ios::ate | ios::trunc | ios::binary);
      int total=0;
      short temp;
      char ch1,ch2;
      if(fout) {
         do {
            memset(buf, 0, sizeof(buf));
            n = udp.recive(buf, sizeof(buf));
            fout.write(&buf[4], n-4);
            fout.flush();
            ch1 = buf[2];
            ch2 = buf[3];
            memset(buf, 0, sizeof(buf));
            *((unsigned short*)buf) = htons((u_short)4);
            sprintf(&buf[2], "%c%c", /*'\0', (char)4,*/ ch1, ch2);
            udp.send(buf, 4);
         }
         while(n == 516);
         fout.close();
      }
      return ret;
   }

   bool put(const string & file_name) {
      char buf[516];
      bool ret = false;
      udp_socket udp(server, port, udp_socket::UDP_CLIENT);
      //udp_stream out(server, port, udp_socket::UDP_CLIENT);
      memset(buf, 0, sizeof(buf));
      *((unsigned short*)buf) = htons((u_short)2);
      //sprintf(buf+2, "%s%c%s%c", file_name.c_str(), '\0', "netascii", '\0');
      sprintf(buf+2, "%s%c%s%c", file_name.c_str(), '\0', "octet", '\0');
      //udp.send(buf, 2+file_name.size()+1+8+1);
      int n = udp.send(buf, 2+file_name.size()+1+5+1);
      memset(buf, 0, sizeof(buf));
      udp.recive(buf, 4);
      short temp;
      temp = *((short*)buf);
      ifstream fin(file_name.c_str(), ios::in | ios::binary);
      short count = 1;
      memset(buf, 0, sizeof(buf));
      *((short*)buf) = htons((u_short)3);
      *((short*)&buf[2]) = htons((u_short)count);
      int size = 0;
      n = read_buffer(fin, &buf[4], sizeof(buf)-4);
      while(n > 0) {
         size += n;
         udp.send(buf, n+4);
         memset(buf, 0, sizeof(buf));
         udp.recive(buf, 4);
         temp = *((short*)buf);
         count++;
         memset(buf, 0, sizeof(buf));
         *((short*)buf) = htons((u_short)3);
         *((short*)&buf[2]) = htons((u_short)count);
         n = read_buffer(fin, &buf[4], sizeof(buf)-4);
      }
      fin.close();
      return ret;
   }

/*
   bool put(const string & file_name) {
      datagram buf(516);
      buf.clear();
      bool ret = false;
      udp_stream udp(server, port, udp_socket::UDP_CLIENT);
      *((unsigned short*)buf.ptr()) = htons((u_short)2);
      sprintf(buf.ptr()+2, "%s%c%s%c", file_name.c_str(), '\0', "octet", '\0');
      udp << buf;
      buf.clear();
      udp >> buf;
      short temp;
      temp = *((short*)buf.ptr());
      ifstream fin(file_name.c_str(), ios::in | ios::binary);
      int n;
      short count = 1;
      buf.clear();
      *((short*)buf.ptr()) = htons((u_short)3);
      *((short*)buf.ptr()+2) = htons((u_short)count);
      int size = 0;
      n = read_buffer(fin, buf.ptr()+4, buf.size()-4);
      while(n > 0) {
         size += n;
         udp << buf;
         buf.clear();
         udp >> buf;
         temp = *((short*)buf.ptr());
         count++;
         buf.clear();
         *((short*)buf.ptr()) = htons((u_short)3);
         *((short*)buf.ptr()+2) = htons((u_short)count);
         n = read_buffer(fin, buf.ptr()+4, buf.size()-4);
      }
      fin.close();
      return ret;
   }
*/
};

}


#endif /* _TFTP2_H */
