#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "ymodem.h"

/* Common routines */
#define IS_AF(c)       ((c >= 'A') && (c <= 'F'))
#define IS_af(c)       ((c >= 'a') && (c <= 'f'))
#define IS_09(c)       ((c >= '0') && (c <= '9'))
#define ISVALIDHEX(c)  IS_AF(c) || IS_af(c) || IS_09(c)
#define ISVALIDDEC(c)  IS_09(c)
#define CONVERTDEC(c)  (c - '0')
#define CONVERTHEX_alpha(c)  (IS_AF(c) ? (c - 'A'+10) : (c - 'a'+10))
#define CONVERTHEX(c)   (IS_09(c) ? (c - '0') : CONVERTHEX_alpha(c))

/**
 * @brief  Receive byte from sender
 * @param  c: Character
 * @param  timeout: Timeout
 * @retval 0: Byte received
 *        -1: Timeout
 */
static int get_byte(struct ymodem_xfer *uart, uint8_t *c, uint32_t timeout)
{
    if (uart->getc)
        return uart->getc(uart, c, timeout);
    return -1;
}

/*
 * @brief  Send a byte
 * @param  c: Character
 * @retval 0: Byte sent
 */
static int put_byte (struct ymodem_xfer *uart, uint8_t c)
{
    if (uart->putc)
        return uart->putc(uart, c);
    return -1;
}

/**
 * @brief  Update CRC16 for input byte
 * @param  CRC input value
 * @param  input byte
 * @retval None
 */
static uint16_t crc16(uint16_t crcIn, uint8_t byte)
{
    uint32_t crc = crcIn;
    uint32_t in = byte | 0x100;

    do {
        crc <<= 1;
        in <<= 1;
        if(in & 0x100)
            ++crc;
        if(crc & 0x10000)
            crc ^= 0x1021;
    } while (!(in & 0x10000));
    return crc & 0xffffu;
}


/**
 * @brief  Cal CRC16 for YModem Packet
 * @param  data
 * @param  length
 * @retval None
 */
static uint16_t do_crc16(const uint8_t* data, uint32_t size)
{
    uint32_t crc = 0;
    const uint8_t* dataEnd = data+size;

    while(data < dataEnd)
        crc = crc16(crc, *data++);

    crc = crc16(crc, 0);
    crc = crc16(crc, 0);

    return crc&0xffffu;
}

struct rxdescr {
    struct ymodem_xfer *xfer;

    int     last_error;

    int     errors;
    int     eot;
    int     abort;

    uint8_t seq;

    char    fname[FILE_NAME_LENGTH + 1];
    char    fsize[FILE_SIZE_LENGTH + 1];

    uint8_t data[PACKET_1K_SIZE + PACKET_OVERHEAD];
};

/*
 * recv a ymodem frame over xfer.
 *
 * return : < 0 on error.
 *        : = 0 on eot or abort.
 *        : > 0 packet size.
 */
static int recv_frame(struct rxdescr *rx, uint32_t timeout)
{
    uint8_t c;
    uint16_t ccs, rcs;
    int i = 0, packet_size;

    if (get_byte(rx->xfer, &c, timeout) != 0) {
        return -1;
    }

    switch (c) {
        case SOH:
            packet_size = PACKET_SIZE;
            break;
        case STX:
            packet_size = PACKET_1K_SIZE;
            break;
        case EOT:
            ++rx->eot;
            return 0;
        case CA:
            if ((get_byte(rx->xfer, &c, timeout) == 0) && (c == CA)) {
                ++rx->abort;
                return 0;
            } else {
                return -1;
            }
        case ABORT1:
        case ABORT2:
            ++rx->abort;
            return 0;
        default:
            return -1;
    }

    rx->data[i++] = c;
    for (i = 1; i < (packet_size + PACKET_OVERHEAD); i++) {
        if (get_byte(rx->xfer, rx->data + i, timeout) != 0) {
            return -2;
        }
    }

    if (rx->data[PACKET_SEQNO_INDEX] != ((rx->data[PACKET_SEQNO_COMP_INDEX] ^ 0xff) & 0xff)) {
        return -3;
    }

    /* crc in packet */
    rcs = rx->data[packet_size + PACKET_OVERHEAD - 2];
    rcs <<= 8;
    rcs |= rx->data[packet_size + PACKET_OVERHEAD - 1];

    /* computed crc from data */
    ccs = do_crc16(&rx->data[PACKET_HEADER], packet_size);
    if (ccs != rcs) {
#ifdef YMODEM_DEBUG
        printf("crc mismatch: %04x, expected: %04x\n", ccs, rcs);
#endif
        return -4;
    }

    return (packet_size);
}

int ymodem_recv(struct ymodem_xfer *uart)
{
    struct rxdescr *rx;
    uint8_t *fp;
    int pkt_size, i;
    unsigned int fsize, recv_bytes = 0;

    if (!uart->on_info || !uart->on_data) {
#ifdef YMODEM_DEBUG
        printf("please set `on_info` and `on_data`\n");
#endif
        return -1;
    }
    /* allocate enough memory. */
    rx = malloc(sizeof *rx);
    if (rx == NULL) {
#ifdef YMODEM_DEBUG
        printf("%s: out of memory.\n", __func__);
#endif
        return -1;
    }
    memset(rx, 0, sizeof *rx);
    rx->xfer = uart;

retry:
    /* wait to start */
    while ((pkt_size = recv_frame(rx, NAK_TIMEOUT)) <= 0) {
        put_byte(uart, WANTCRC);
    }

    if (rx->data[PACKET_SEQNO_INDEX] != 0) { /* seq mismatch */
        put_byte(uart, NAK);
        goto retry;
    }

    fp = &rx->data[PACKET_HEADER];
    /* packet 0, filename packet. */
    if (*fp == 0) { /* have no filename, it's not a vlid ymodem first packet */
        put_byte(uart, CA);
        put_byte(uart, CA);
#ifdef YMODEM_DEBUG
        printf("no filename retry\n");
#endif
        goto retry;
    }

    /* send req 0 response */
    put_byte(uart, ACK);
    put_byte(uart, WANTCRC);

    /* parse filename */
    i = 0;
    while ((*fp != 0) && (i < FILE_NAME_LENGTH)) {
        rx->fname[i++] = *fp++;
    }
    rx->fname[i++] = '\0';

    /* parse file size */
    i = 0;
    ++fp;
    while ((*fp != ' ') && (i < FILE_SIZE_LENGTH)) {
        rx->fsize[i++] = *fp++;
    }
    rx->fsize[i++] = '\0';
    fsize = strtoul(rx->fsize, NULL, 0);

    /* call file info callback */
    rx->last_error = uart->on_info(uart, rx->fname, fsize);

#ifdef YMODEM_DEBUG
    printf("%s : (%d) bytes (last_error %d)\n", rx->fname, fsize, rx->last_error);
#endif

    /* initialize ymodem state */
    rx->eot    = 0;
    rx->abort  = 0;
    rx->errors = 0;
    rx->seq    = 1;

    /* start to recv file contents */
    fp = &rx->data[PACKET_HEADER];
    for (;;) {
        pkt_size = recv_frame(rx, NAK_TIMEOUT);

        if (rx->abort > 0) {
#ifdef YMODEM_DEBUG
            printf("abort\n");
#endif
            break;
        }

        if (rx->last_error) {
            put_byte(uart, CA);
            put_byte(uart, CA);
#ifdef YMODEM_DEBUG
            printf("last_error break\n");
#endif
            break;
        }

        if (pkt_size < 0) {
            ++rx->errors;
            if (rx->errors >= MAX_ERRORS) {
                put_byte(uart, CA);
                put_byte(uart, CA);
#ifdef YMODEM_DEBUG
                printf("errors break\n");
#endif
                break;
            } else {
                put_byte(uart, NAK);
#ifdef YMODEM_DEBUG
                printf("send NAK\n");
#endif
                continue;
            }
        }

        rx->errors = 0;

        if (rx->eot == 2 && rx->data[PACKET_SEQNO_INDEX] == 0) {
#ifdef YMODEM_DEBUG
            printf("end of ymodem protocol.\n");
#endif
            put_byte(uart, ACK);
            break;
        }

        if (pkt_size > 0) { /* on data */
            if (rx->data[PACKET_SEQNO_INDEX] == 0) { /* all zeros */
                for (i = 0; i < pkt_size; ++i) {
                    if (fp[i] != 0) {
                        break;
                    }
                }
                if ((i == pkt_size) && rx->eot > 0) { /* end of stream */
                    put_byte(uart, ACK);
                    break;
                }
            }

            if (rx->seq != rx->data[PACKET_SEQNO_INDEX]) {
                put_byte(uart, NAK);
#ifdef YMODEM_DEBUG
                printf("data phase: seq mismatch %d, expected %d\n", rx->data[PACKET_SEQNO_INDEX], rx->seq);
#endif
            } else {
                put_byte(uart, ACK);

                rx->last_error = uart->on_data(uart, fp, pkt_size);

                ++rx->seq;
                recv_bytes += pkt_size;
            }
            continue;
        }

        /* eot */
        if (pkt_size == 0) {
            if (rx->eot == 1) {
                put_byte(uart, NAK);
                continue;
            }

            if (rx->eot == 2) {
                put_byte(uart, ACK);
                put_byte(uart, WANTCRC);
            }
        }
    }

    free(rx);

    return (recv_bytes);
}

/**
 * @brief  Prepare the first block
 * @param  timeout
 *     0: end of transmission
 * @retval None
 */
static void init_packet(uint8_t *data, const uint8_t* fileName, uint32_t *length)
{
    uint16_t i, j;
    uint8_t file_ptr[10];

    /* Make first three packet */
    data[0] = SOH;
    data[1] = 0x00;
    data[2] = 0xff;

    /* Filename packet has valid data */
    for (i = 0; (fileName[i] != '\0') && (i < FILE_NAME_LENGTH);i++) {
        data[i + PACKET_HEADER] = fileName[i];
    }

    data[i + PACKET_HEADER] = 0x00;

    snprintf((char *)file_ptr, sizeof file_ptr, "%lu", *length);
    for (j =0, i = i + PACKET_HEADER + 1; file_ptr[j] != '\0' ; ) {
        data[i++] = file_ptr[j++];
    }

    for (j = i; j < PACKET_SIZE + PACKET_HEADER; j++) {
        data[j] = 0;
    }
}

/**
 * @brief  Prepare the data packet
 * @param  timeout
 *     0: end of transmission
 * @retval None
 */
static void prepare_packet(uint8_t *sourceBuf, uint8_t *data, uint8_t pktNo, uint32_t sizeBlk)
{
    uint16_t i, size, packetSize;
    uint8_t* filePtr;

    /* Make first three packet */
    packetSize = sizeBlk >= PACKET_1K_SIZE ? PACKET_1K_SIZE : PACKET_SIZE;
    size = sizeBlk < packetSize ? sizeBlk :packetSize;
    if (packetSize == PACKET_1K_SIZE)
    {
        data[0] = STX;
    }
    else
    {
        data[0] = SOH;
    }
    data[1] = pktNo;
    data[2] = (~pktNo);
    filePtr = sourceBuf;

    /* Filename packet has valid data */
    for (i = PACKET_HEADER; i < size + PACKET_HEADER;i++) {
        data[i] = *filePtr++;
    }
    if ( size  <= packetSize) {
        for (i = size + PACKET_HEADER; i < packetSize + PACKET_HEADER; i++) {
            data[i] = 0x1A; /* EOF (0x1A) or 0x00 */
        }
    }
}

/**
 * @brief  Cal Check sum for YModem Packet
 * @param  data
 * @param  length
 * @retval None
 */
static uint8_t CalChecksum(const uint8_t* data, uint32_t size)
{
    uint32_t sum = 0;
    const uint8_t* dataEnd = data+size;

    while(data < dataEnd )
        sum += *data++;

    return (sum & 0xffu);
}

/**
 * @brief  Transmit a data packet using the ymodem protocol
 * @param  data
 * @param  length
 * @retval None
 */
static void send_packet(struct ymodem_xfer *uart, uint8_t *data, uint16_t length)
{
    uint16_t i;
    i = 0;
    while (i < length) {
        put_byte(uart, data[i]);
        i++;
    }
}

/**
 * @brief  Transmit a file using the ymodem protocol
 * @param  buf: Address of the first byte
 * @retval The size of the file
 */
int ymodem_send(struct ymodem_xfer *uart, uint8_t *buf, const uint8_t* sendFileName, uint32_t sizeFile)
{
    uint8_t packet_data[PACKET_1K_SIZE + PACKET_OVERHEAD];
    uint8_t filename[FILE_NAME_LENGTH];
    uint8_t *buf_ptr, tempCheckSum;
    uint16_t tempCRC;
    uint16_t blkNumber;
    uint8_t receivedC[2], CRC16_F = 0, i;
    uint32_t errors, ackReceived, size = 0, pktSize;

    errors = 0;
    ackReceived = 0;

    for (i = 0; i < (FILE_NAME_LENGTH - 1); i++) {
        filename[i] = sendFileName[i];
    }
    CRC16_F = 1;

    /* Prepare first block */
    init_packet(&packet_data[0], filename, &sizeFile);

    do {
        /* Send Packet */
        send_packet(uart, packet_data, PACKET_SIZE + PACKET_HEADER);

        /* Send CRC or Check Sum based on CRC16_F */
        if (CRC16_F) {
            tempCRC = do_crc16(&packet_data[3], PACKET_SIZE);
            put_byte(uart, tempCRC >> 8);
            put_byte(uart, tempCRC & 0xFF);
        } else {
            tempCheckSum = CalChecksum (&packet_data[3], PACKET_SIZE);
            put_byte(uart, tempCheckSum);
        }

        /* Wait for Ack and 'C' */
        if (get_byte(uart, &receivedC[0], 10000) == 0)  {
            if (receivedC[0] == ACK) {
                /* Packet transferred correctly */
                ackReceived = 1;
            }
        } else {
            errors++;
        }
    } while (!ackReceived && (errors < 0x0A));

    if (errors >=  0x0A) {
        return errors;
    }
    buf_ptr = buf;
    size = sizeFile;
    blkNumber = 0x01;
    /* Here 1024 bytes package is used to send the packets */


    /* Resend packet if NAK  for a count of 10 else end of communication */
    while (size) {
        /* Prepare next packet */
        prepare_packet(buf_ptr, &packet_data[0], blkNumber, size);
        ackReceived = 0;
        receivedC[0]= 0;
        errors = 0;
        do {
            /* Send next packet */
            if (size >= PACKET_1K_SIZE) {
                pktSize = PACKET_1K_SIZE;

            } else {
                pktSize = PACKET_SIZE;
            }
            send_packet(uart, packet_data, pktSize + PACKET_HEADER);
            /* Send CRC or Check Sum based on CRC16_F */
            /* Send CRC or Check Sum based on CRC16_F */
            if (CRC16_F) {
                tempCRC = do_crc16(&packet_data[3], pktSize);
                put_byte(uart, tempCRC >> 8);
                put_byte(uart, tempCRC & 0xFF);
            } else {
                tempCheckSum = CalChecksum (&packet_data[3], pktSize);
                put_byte(uart, tempCheckSum);
            }

            /* Wait for Ack */
            if ((get_byte(uart, &receivedC[0], 100000) == 0)  && (receivedC[0] == ACK)) {
                ackReceived = 1;
                if (size > pktSize) {
                    buf_ptr += pktSize;
                    size -= pktSize;
                    if (blkNumber > ((sizeFile / 1024) + 1)) { /* XXX */
                        return 0xFF; /*  error */
                    } else {
                        blkNumber++;
                    }
                } else {
                    buf_ptr += pktSize;
                    size = 0;
                }
            } else {
                errors++;
            }
        } while(!ackReceived && (errors < 0x0A));
        /* Resend packet if NAK  for a count of 10 else end of communication */

        if (errors >=  0x0A) {
            return errors;
        }

    }
    ackReceived = 0;
    receivedC[0] = 0x00;
    errors = 0;
    do {
        put_byte(uart, EOT);
        /* Send (EOT); */
        /* Wait for Ack */
        if ((get_byte(uart, &receivedC[0], 10000) == 0)  && receivedC[0] == ACK) {
            ackReceived = 1;
        } else {
            errors++;
        }
    } while (!ackReceived && (errors < 0x0A));

    if (errors >=  0x0A) {
        return errors;
    }

    /* Last packet preparation */
    ackReceived = 0;
    receivedC[0] = 0x00;
    errors = 0;

    packet_data[0] = SOH;
    packet_data[1] = 0;
    packet_data [2] = 0xFF;

    for (i = PACKET_HEADER; i < (PACKET_SIZE + PACKET_HEADER); i++) {
        packet_data [i] = 0x00;
    }

    do {
        /* Send Packet */
        send_packet(uart, packet_data, PACKET_SIZE + PACKET_HEADER);

        /* Send CRC or Check Sum based on CRC16_F */
        tempCRC = do_crc16(&packet_data[3], PACKET_SIZE);
        put_byte(uart, tempCRC >> 8);
        put_byte(uart, tempCRC & 0xFF);

        /* Wait for Ack and 'C' */
        if (get_byte(uart, &receivedC[0], 10000) == 0) {
            if (receivedC[0] == ACK) {
                /* Packet transferred correctly */
                ackReceived = 1;
            }
        } else {
            errors++;
        }
    } while (!ackReceived && (errors < 0x0A));

    /* Resend packet if NAK  for a count of 10  else end of communication */
    if (errors >=  0x0A) {
        return errors;
    }

    do {
        put_byte(uart, EOT);
        /* Send (EOT); */
        /* Wait for Ack */
        if ((get_byte(uart, &receivedC[0], 10000) == 0)  && receivedC[0] == ACK) {
            ackReceived = 1;
        } else {
            errors++;
        }
    } while (!ackReceived && (errors < 0x0A));

    if (errors >=  0x0A) {
        return errors;
    }
    return 0; /* file transmitted successfully */
}