#include "xlrutftp.h"
//#include "qgnetworkinterface.h"

#include <QUdpSocket>
#include <QNetworkDatagram>
#include <QTimer>
#include <QFile>
#include <QFileInfo>
#include <QQueue>

#include <QDebug>

#define dbg qDebug

#include <string.h>

#define TFTP_MAXPKTLEN  1024

#define DEFAULT_TGT_PORT    50069

//--
#define TFTP_RRQ    01          /* read request */
#define TFTP_WRQ    02          /* write request */
#define TFTP_DATA   03          /* data packet */
#define TFTP_ACK    04          /* acknowledgement */
#define TFTP_ERROR  05          /* error code */
#define TFTP_OACK   06          /* option acknowledgement */

#define  TFTP_EUNDEF      0       /* not defined */
#define  TFTP_ENOTFOUND   1       /* file not found */
#define  TFTP_EACCESS     2       /* access violation */
#define  TFTP_ENOSPACE    3       /* disk full or allocation exceeded */
#define  TFTP_EBADOP      4       /* illegal TFTP operation */
#define  TFTP_EBADID      5       /* unknown transfer ID */
#define  TFTP_EEXISTS     6       /* file already exists */
#define  TFTP_ENOUSER     7       /* no such user */
#define  TFTP_EBADOPTION  8        /* bad option */
#define  TFTP_ECANCELLED 99      /* cancelled by administrator */

typedef struct {
  unsigned char  th_opcode[2];      /* packet type */
  union {
    unsigned char tu_block[2];   /* block # */
    unsigned char tu_code[2];    /* error code */
    unsigned char tu_stuff[1];    /* request packet stuff */    
    char tu_name[1];
  } th_u;
} tftphdr_t;


typedef struct
{
  tftphdr_t header;
  char    th_data[2048];     /* data or error string */
} tftppkt_t;

static void tftp_htons(unsigned char* const p, unsigned short v)
{
  p[0]=(unsigned char)(v>>8);
  p[1]=(unsigned char)(v);
}

static unsigned short tftp_ntos(const unsigned char* n)
{
  return n[1]|(n[0]<<8);
}

//--

class XlruTftp::Implementation
{
public:
    XlruTftp* parent;

    int tgt_port;
    QHostAddress tgtAddress;

    QUdpSocket rxSocket;
    QUdpSocket txSocket;

    QTimer txTimer;
    QTimer stsTimer;

    unsigned int data_len;

    unsigned long timeout_sts_ms;
    unsigned long timeout_first_ms;
    unsigned long timeout_last_ms;
    unsigned long timeout_data_ms;

    QString lastErrorString;
    int lastErrorCode;
    QHostAddress lastTargetAddress;
    int lastTargetPort;

    QString lastStatusString;

    QByteArray fileBuffer;
    const char* data_pointer;
    int data_size;
    unsigned short block_counter;

    tftppkt_t pkt;

    int state;
    int rx_mode;
    char* rx_data_pointer;
    int rx_size;

    int sts_timeout_counter;

    QHostAddress localIP;
    QHostAddress broadcastIP;

    qint64 sendPacket(const tftppkt_t& pkt, unsigned int size, QUdpSocket* s=0)
    {
        int total_size=size+sizeof pkt.header;
        if (s==0)
            s=&txSocket;
        dbg()<<"->"<<tftp_ntos(pkt.header.th_opcode)<<tftp_ntos(pkt.header.th_u.tu_stuff)<<size<<total_size<<tgtAddress<<tgt_port;
        return s->writeDatagram(reinterpret_cast<const char*>(&pkt.header.th_opcode[0]), total_size, tgtAddress, tgt_port);
    }

    Implementation(XlruTftp* controller):
        parent(controller)
    {
        memset(&pkt, 0,sizeof pkt);

        data_len=TFTP_MAXPKTLEN;

        timeout_first_ms=10*1000;
        timeout_last_ms=30*1000;
        timeout_data_ms=10*1000;
        timeout_sts_ms=1*1000;
        txTimer.setSingleShot(true);
        stsTimer.setInterval(timeout_sts_ms);

        rx_mode=0;
        state=0;
    }

    int extractProgress(const QString& str)
    {
        bool ok=false;
        QString tmp=str.section(' ', 0, 0);
        int n=tmp.toInt(&ok, 0);
        if (ok)
            return n;
        else
            return -1;
    }

    int rxDataReady()
    {
        while(rxSocket.hasPendingDatagrams())
        {
            QNetworkDatagram d=rxSocket.receiveDatagram();
            unsigned int size=d.data().size();
            if (size<sizeof(tftphdr_t))
                continue;
            const tftppkt_t& rxpkt=*reinterpret_cast<const tftppkt_t*>(d.data().constData());
            auto opcode=tftp_ntos(rxpkt.header.th_opcode);


            //if ((d.senderAddress()!=tgtAddress) || (d.senderPort()!=tgt_port))
            //    continue;

            lastTargetAddress=d.senderAddress();
            lastTargetPort=d.senderPort();

            dbg()<<"RRQ-ACK"<<lastTargetAddress<<lastTargetPort<<opcode<<state;

            if (rx_mode)
            {
                tftppkt_t reply;
                if ((opcode==TFTP_ACK) || (opcode==TFTP_OACK))
                {
                    tftp_htons(reply.header.th_opcode, TFTP_ACK);
                    reply.header.th_u.tu_block[0]=0;
                    reply.header.th_u.tu_block[1]=0;

                    sendPacket(reply, 4, &rxSocket);
                }
                else if (opcode==TFTP_DATA)
                {
                    int rx_data_size=size-sizeof(rxpkt.header);
                    if (rx_data_size>0)
                    {
                        memcpy(rx_data_pointer, rxpkt.th_data, size-sizeof(rxpkt.header));
                        rx_data_pointer+=rx_data_size;
                    }
                    if (size<data_len)
                    {
                        rx_mode=0;
                        qDebug()<<"RRQ: Terminated";
                        Q_EMIT parent->receiveTerminated(true);
                    }
                    tftp_htons(reply.header.th_opcode, TFTP_ACK);
                    reply.header.th_u.tu_block[0]=rxpkt.header.th_u.tu_block[0];
                    reply.header.th_u.tu_block[1]=rxpkt.header.th_u.tu_block[1];

                    sendPacket(reply, 4, &rxSocket);
                }
                else
                {
                    stsTimer.stop();
                    return 1;
                }

                //rxSocket.writeDatagram((const char*)&reply, sizeof reply, lastTargetAddress, lastTargetPort);
                //sendPacket(reply, 4, &rxSocket);
            }
            if (opcode!=TFTP_ERROR)
            {
                lastErrorCode=1000;
                lastErrorString=QString("invalid opcode %1").arg(opcode);
                stsTimer.stop();
                return 1;
            }
            unsigned short ecode=tftp_ntos(rxpkt.header.th_u.tu_code);
            if (ecode==TFTP_EBADID)
            {
                if (sts_timeout_counter)
                    --sts_timeout_counter;

                lastStatusString=QString::fromLatin1(&rxpkt.th_data[0]);
                if (txTimer.isActive())
                    txTimer.start();
                return 0;
            }
            else if (ecode==TFTP_ENOTFOUND)
            {
                lastStatusString=QString::fromLatin1(&rxpkt.th_data[0]);
                if (lastStatusString=="ready")
                {
                    stsTimer.stop();
                }
                return 0;
            }
            else if (ecode==TFTP_EEXISTS)
            {
                lastErrorCode=ecode;
                lastErrorString=QString::fromLatin1(&rxpkt.th_data[0]);
                return 2;
            }
            else
            {
                stsTimer.stop();
                lastErrorCode=ecode;
                lastErrorString=QString::fromLatin1(&rxpkt.th_data[0]);
                return 1;
            }
        }
        return 0;
    }

    bool dataReady()
    {
        if (state<1)
        {
            lastErrorCode=1000;
            lastErrorString=QString("Unexpected datagram");
            txSocket.reset();
            return false;
        }
        while(txSocket.hasPendingDatagrams())
        {
            QNetworkDatagram d=txSocket.receiveDatagram();
            if ((d.senderAddress()!=tgtAddress))// || (d.senderPort()!=tgt_port))
                continue;
            unsigned int size=d.data().size();
            if (size<sizeof(tftphdr_t))
                continue;

            const tftppkt_t& rxpkt=*reinterpret_cast<const tftppkt_t*>(d.data().constData());
            auto opcode=tftp_ntos(rxpkt.header.th_opcode);
            auto bknum=tftp_ntos(rxpkt.header.th_u.tu_block);
            dbg()<<"WRQ-ACK"<<tgtAddress<<tgt_port<<opcode<<bknum<<state;

            if ((opcode==TFTP_ACK) || (opcode==TFTP_OACK))
            {
                tgt_port=d.senderPort();

                if (state==2)
                {
                    dbg()<<"Finished!";
                    lastErrorString="Transfer completed";
                    lastErrorCode=0;
                    stsTimer.stop();
                    txTimer.stop();
                    state=0;
                    return false;
                }
                ++block_counter;
                int frag_size=data_size>data_len ? data_len: data_size;
                tftp_htons(pkt.header.th_opcode, TFTP_DATA);
                tftp_htons(pkt.header.th_u.tu_block, block_counter);

                memcpy(&pkt.th_data[0], data_pointer, frag_size);

                data_pointer+=frag_size;
                data_size-=frag_size;

                qint64 r=sendPacket(pkt, frag_size); //txSocket.writeDatagram(reinterpret_cast<char*>(&pkt.header.th_opcode[0]), frag_size+sizeof pkt.header, tgtAddress, tgt_port);
                if (frag_size<data_len)
                {
                    txTimer.start(timeout_last_ms);
                    stsTimer.start(timeout_sts_ms);
                    state=2;
                }
                else
                    txTimer.start(timeout_data_ms);
            }
            else if (opcode==TFTP_ERROR)
            {
                state=0;
                lastErrorCode=tftp_ntos(rxpkt.header.th_u.tu_code);
                lastErrorString=QString::fromLatin1(&rxpkt.th_data[0]);
                txTimer.stop();
                return false;
            }
            else
            {
                state=0;
                lastErrorCode=1000;
                lastErrorString=QString("Invadid opcode (%1)").arg(opcode);
                txTimer.stop();
                return false;
            }
        }
        return true;
    }

    bool send(const QString& localFile, const QString& remoteFile)
    {
        QFile file(localFile);
        if (!file.open(QIODevice::ReadOnly))
        {
            lastErrorString=file.errorString();
            return false;
        }
        QFileInfo finfo(localFile);


        fileBuffer=file.readAll();
        file.close();

        data_size=fileBuffer.size();
        if (data_size==0)
        {
            return false;
        }

        QString rFile(remoteFile.isEmpty() ? finfo.fileName() : QString(remoteFile).arg(finfo.fileName()));

        return send(fileBuffer.constData(), data_size, rFile);
    }

    int addNameAndOptions(char* p, const char* name)
    {
        const char* opt[]={0, "octet", "blksize", "1024", 0};
        opt[0]=name;
        unsigned int max_len=TFTP_MAXPKTLEN;
        unsigned int added_len=0;

        for(int i=0; opt[i]; ++i)
        {
            unsigned int len=strlen(opt[i]);
            strncpy(p, opt[i], max_len);
            added_len+=len;
            p+=len;
            ++p;
        }
        return added_len+2;
    }

    bool send(const void* data_address, unsigned int len, const QString& remoteFile)
    {
        if (state!=0)
        {
            qDebug()<<"TFTP busy!";
            return false;
        }
        data_pointer=(const char*)data_address; //fileBuffer.constData();
        data_size=len;

        block_counter=0;
        sts_timeout_counter=0;

        memset(&pkt, 0, sizeof pkt);
        tftp_htons(pkt.header.th_opcode, TFTP_WRQ);
        const char* pname=remoteFile.toLatin1().constData();
#if 0
        unsigned int name_len=strlen(pname);
        strncpy((char*)pkt.header.th_u.tu_stuff, pname, TFTP_MAXPKTLEN);
        unsigned char* p=&pkt.header.th_u.tu_stuff[name_len+1];
        strcpy((char*)p, "octet");
        qint64 msg_len=name_len+strlen("octet")+1;  //+(sizeof pkt.header)
#endif
        unsigned int add_len=addNameAndOptions((char*)pkt.header.th_u.tu_stuff, pname);
        qint64 msg_len=1+add_len;

        state=0;

        qint64 r=sendPacket(pkt, msg_len); //txSocket.writeDatagram(reinterpret_cast<char*>(&pkt.header.th_opcode[0]), msg_len, tgtAddress, tgt_port);
        //dbg()<<"WRQ"<<tgtAddress<<tgt_port;
        if (r<1)
        {
            lastErrorString=txSocket.errorString();
            lastErrorCode=txSocket.error();
            return false;
        }
        state=1;
        txTimer.start(timeout_first_ms);
        return true;
    }

    bool txTimeOut()
    {
        state=0;
        return false;
    }

    bool stsTimeOut()
    {
        if (state==2)
        {
            dbg()<<"Request status";
            ++sts_timeout_counter;
            if (sts_timeout_counter>10)
            {
                lastErrorCode=100;
                lastErrorString="status timeout";
                return false;
            }
            memset(&pkt, 0, sizeof pkt);
            tftp_htons(pkt.header.th_opcode, TFTP_RRQ);
            static const char sts_fname[]="$status$";
            strcpy((char*)pkt.header.th_u.tu_stuff, sts_fname);
            qint64 msg_len=sizeof(sts_fname); //(sizeof pkt.header)
            sendPacket(pkt, msg_len, &rxSocket); //rxSocket.writeDatagram(reinterpret_cast<char*>(&pkt.header.th_opcode[0]), msg_len, tgtAddress, tgt_port);
            stsTimer.start();
        }
        return true;
    }

    void pingServer(int port)
    {
        memset(&pkt, 0, sizeof pkt);
        tftp_htons(pkt.header.th_opcode, TFTP_RRQ);

        static const char sts_fname[]="$list$0octet";
        strcpy((char*)pkt.header.th_u.tu_stuff, sts_fname);
        ((char*)pkt.header.th_u.tu_stuff)[6]=0;

        //static const char sts_fname[]="$list$octet";
        //strcpy((char*)pkt.header.th_u.tu_stuff, sts_fname);
        qint64 msg_len=sizeof(sts_fname); //(sizeof pkt.header)
        bool ok=rxSocket.writeDatagram((char*)&pkt, msg_len+sizeof pkt.header, broadcastIP, port); //QgNetworkInterface::networkBroadcast(), port);
        if (!ok)
            qDebug()<<"ping fial:"<<rxSocket.errorString();
    }

    bool receive(void* data_address, unsigned int len, const QString& remoteFile)
    {
        rx_data_pointer=(char*)data_address; //fileBuffer.constData();
        //rx_data_size=len;

        block_counter=0;
        sts_timeout_counter=0;

        memset(&pkt, 0, sizeof pkt);
        tftp_htons(pkt.header.th_opcode, TFTP_RRQ);
        const char* pname=remoteFile.toLatin1().constData();

        unsigned int add_len=addNameAndOptions((char*)pkt.header.th_u.tu_stuff, pname);
        qint64 msg_len=1+add_len;

        state=0;
        rx_mode=1;

        qint64 r=sendPacket(pkt, msg_len, &rxSocket); //txSocket.writeDatagram(reinterpret_cast<char*>(&pkt.header.th_opcode[0]), msg_len, tgtAddress, tgt_port);
        //dbg()<<"WRQ"<<tgtAddress<<tgt_port;
        if (r<1)
        {
            lastErrorString=txSocket.errorString();
            lastErrorCode=txSocket.error();
            return false;
        }
        state=0;
        //rxTimer.start(timeout_first_ms);
        return true;

    }

};

XlruTftp::XlruTftp(QObject *parent):
    QObject(parent),
    p_(*new Implementation(this))
{

    connect(&p_.txSocket, &QUdpSocket::readyRead, this,
            [this]()
    {
        //while(p_.txSocket.hasPendingDatagrams())
        {
            bool ok=p_.dataReady();
            if (!ok)
            {
                if (p_.lastErrorCode==0)
                    Q_EMIT sendTerminated();
                else
                    Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
            }
            else if (p_.block_counter==1)
                Q_EMIT progressInfo(-1, "Uploading");
        }

    });

    connect(&p_.rxSocket, &QUdpSocket::readyRead, this,
            [this]()
    {
        //while(p_.rxSocket.hasPendingDatagrams())
        {
            int ok=p_.rxDataReady();
            if (ok==2)
            {
                Q_EMIT serverHello(p_.lastTargetAddress.toIPv4Address(), p_.lastTargetPort, p_.lastErrorString);
            }
            else if (ok)
                    Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
            else
            {
                int n=p_.extractProgress(p_.lastStatusString);
                Q_EMIT progressInfo(n, p_.lastStatusString);
            }
        }
    });

    connect(&p_.txTimer, &QTimer::timeout, this,
            [this]()
    {
       bool ok=p_.txTimeOut();
       if (!ok)
       {
            dbg()<<"Tx timeout";
            Q_EMIT transferError(100, "Timeout");
       }

    });
    connect(&p_.stsTimer, &QTimer::timeout, this,
            [this]()
    {
       bool ok=p_.stsTimeOut();
       if (false)
           Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
    });
}

XlruTftp::~XlruTftp()
{
    p_.txTimer.stop();
    p_.txTimer.disconnect();
    p_.stsTimer.stop();
    p_.stsTimer.disconnect();
    p_.rxSocket.disconnect();
    p_.txSocket.disconnect();

    p_.rxSocket.abort();
    p_.txSocket.abort();

    delete &p_;
}

void XlruTftp::bindRemote(const QHostAddress& adr, unsigned int dst_port)
{
    p_.tgtAddress=adr;
    p_.tgt_port=dst_port ? dst_port : DEFAULT_TGT_PORT;
    qDebug()<<"TFTP Bind:"<<p_.tgtAddress<<p_.tgt_port;
}

void XlruTftp::bind(const QHostAddress& adr)
{
    p_.localIP=adr;
    unsigned int bip=adr.toIPv4Address();
    bip|=0x0FF;

    p_.broadcastIP=QHostAddress(bip);

    p_.rxSocket.bind(adr);
    p_.txSocket.bind(adr);
}

QString XlruTftp::errorString()
{
    return p_.lastErrorString;
}

int XlruTftp::errorCode()
{
    return p_.lastErrorCode;
}

class TftpOperation
{
public:
    int type;

    void* data;
    unsigned int size;
    QString remoteFile;
    unsigned int result_size;
    unsigned int results;
    QString resultMessage;
};

class TftpQueue: public QQueue<TftpOperation>
{
public:
};


bool XlruTftp::sendFile(const QString& localFileName, const QString& remoteFileName)
{
    bool ok=p_.send(localFileName, remoteFileName);
    if (ok)
        Q_EMIT sendBegin();
    else
        Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
    return ok;
}

bool XlruTftp::sendMemory(void* data, unsigned int size, const QString& remoteFileName)
{
    bool ok=p_.send(data, size, remoteFileName);
    if (ok)
        Q_EMIT sendBegin();
    else
        Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
    return ok;
}

void XlruTftp::pingServer(int port)
{
    p_.pingServer(port);
}

bool XlruTftp::receiveMemory(void* data, unsigned int max_size, const QString& remoteFileName)
{
    bool ok=p_.receive(data, max_size, remoteFileName);
    if (ok)
    {
        Q_EMIT sendBegin();
    }
    else
        Q_EMIT transferError(p_.lastErrorCode, p_.lastErrorString);
    return true;
}