/*
 * Copyright (c) 2023 - 2024 Stefan Strobel <stefan.strobel@shimatta.net>
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

/**
 * @file    shellmatta_ymodem.c
 * @brief   ymodem functions of shellmatta
 * @author  Simon Fischer <dev@s-fischer.net>
 */

#include "shellmatta_ymodem.h"
#include "shellmatta_crc.h"
#include "shellmatta_utils.h"
#include <stddef.h>
#include <string.h>

/** @brief symbols needed for ymodem protocol */
typedef enum {
    YMODEM_NULL    = 0x00, /**< NULL-terminator                 */
    YMODEM_SOH     = 0x01, /**< Start of header                 */
    YMODEM_STX     = 0x02, /**< Start of header for 1k packet   */
    YMODEM_EOT     = 0x04, /**< End of transmission             */
    YMODEM_ACK     = 0x06, /**< Acknowledge                     */
    YMODEM_NAK     = 0x15, /**< Negative acknowledge            */
    YMODEM_CA      = 0x18, /**< Cancel                          */
    YMODEM_SPACE   = 0x20, /**< Space                           */
    YMODEM_CRC     = 0x43  /**< 'C' to indicate CRC type        */
} YMODEM_SYMBOLS;

#define YMODEM_PACKET_SIZE          128u    /**< default packet size of ymodem transmission     */
#define YMODEM_PACKET_SIZE_1K       1024u   /**< extended packet size of ymodem transmission    */
#define YMODEM_CRC_SIZE             2u      /**< CRC size of the ymodem packet                  */
#define YMODEM_POLL_NUMBER          2u      /**< Number of polls to be received                 */

/**
 * @brief               forwards the given character to write-function without formatting
 * @param[in]   handle  shellmatta handle of the instance
 * @param[in]   c       character to be sent
*/
static void shellmatta_ymodem_control(shellmatta_handle_t handle, const uint8_t c)
{
    shellmatta_instance_t *inst = (shellmatta_instance_t*)handle;
    SHELLMATTA_WRITE((char*)&c, 1);
}

/**
 * @brief                       reset function for the ymodem module
 * @param[in, out]  handle      shellmatta instance handle
 * @param[in]       doCancel    flag to execute the cancel-callback
*/
static void shellmatta_ymodem_reset(shellmatta_handle_t handle, bool doCancel)
{
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;

    /** -# call cancel callback function */
    if (doCancel)
    {
        /* send cancel symbol */
        shellmatta_ymodem_control(handle, YMODEM_CA);
        if (NULL != inst->ymodem.cancelCallback)
        {
            inst->ymodem.cancelCallback(handle);
        }
    }
    /** -# reset instance variables */
    inst->ymodem.state = SHELLMATTA_YMODEM_INACTIVE;
    inst->ymodem.byteCounter = 0u;
    inst->ymodem.packetCounter = 0u;
    inst->ymodem.totalBytesReceived = 0u;
    inst->ymodem.fileSize = 0u;
    inst->ymodem.pauseRequested = false;
    inst->ymodem.pollCyclesLeft = 0u;

#ifdef SHELLMATTA_TRANSPORT
    /** .-# reenable transport layer optional mode */
    inst->transportLayer.suspendOptional = false;
#endif

    (void)memset((void *)&inst->ymodem.packet, 0, sizeof(shellmatta_ymodem_packet_t));
}

shellmatta_retCode_t processPacket(shellmatta_handle_t handle)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;
    char *fileName;
    uint32_t packetSize;

    /** -# read filesize and name from first packet - ignore rest */
    if(0u == inst->ymodem.packetCounter)
    {
        fileName = (char*)inst->ymodem.packet.packetData;

        ret = utils_shellAsciiToUInt32((char*)&inst->ymodem.packet.packetData[strlen(fileName) + 1u],
                                       &inst->ymodem.fileSize,
                                       10u);

        /** -# pass filename and size to the callback */
        inst->ymodem.recvHeaderCallback(handle, inst->ymodem.fileSize, fileName);
    }
    else
    {
        /** -# calculate packet size - when it is the last packet this is limited by the file size */
        if((inst->ymodem.totalBytesReceived + inst->ymodem.packet.size) > inst->ymodem.fileSize)
        {
            packetSize = inst->ymodem.fileSize % inst->ymodem.packet.size;
        }
        else
        {
            packetSize = inst->ymodem.packet.size;
        }

        /** -# pass data to the application using the callback */
        inst->ymodem.recvPacketCallback(handle, inst->ymodem.packet.packetData, packetSize, inst->ymodem.packetCounter);
    }

    return ret;
}

static shellmatta_retCode_t ymodem_stateMachine(shellmatta_handle_t handle, uint8_t byte)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;
    uint16_t computedCrc;

    switch(inst->ymodem.state)
    {
        case SHELLMATTA_YMODEM_INACTIVE:
            /** -# skip state machine if inactive */
            break;
        case SHELLMATTA_YMODEM_WAIT_FOR_START:
            /** -# Wait for start character */
            inst->ymodem.byteCounter = 0u;
            switch(byte)
            {
                case YMODEM_SOH:
                    inst->ymodem.packet.size = YMODEM_PACKET_SIZE;
                    inst->ymodem.state = SHELLMATTA_YMODEM_RECEIVE_HEADER;
                    break;
                case YMODEM_STX:
                    inst->ymodem.packet.size = YMODEM_PACKET_SIZE_1K;
                    inst->ymodem.state = SHELLMATTA_YMODEM_RECEIVE_HEADER;
                    break;

                case YMODEM_EOT:
                    /** -# ACK the successful file reception */
                    shellmatta_ymodem_control(handle, YMODEM_ACK);
                    shellmatta_ymodem_control(handle, YMODEM_CRC);

                    /** -# handle additional EOTs in WAIT FOR END state */
                    inst->ymodem.state = SHELLMATTA_YMODEM_WAIT_FOR_END;
                    inst->ymodem.pollCyclesLeft = YMODEM_POLL_NUMBER;
                    break;
                default:
                    /** -# ignore unexpected characters on start */
                    break;
            }
            break;
        case SHELLMATTA_YMODEM_RECEIVE_HEADER:
            if(0u == inst->ymodem.byteCounter)
            {
                inst->ymodem.packet.packetNumber = byte;
                inst->ymodem.byteCounter ++;
            }
            else
            {
                if(((0xffu - byte) != inst->ymodem.packet.packetNumber) ||
                  (inst->ymodem.packetCounter % 256u != inst->ymodem.packet.packetNumber))
                {
                    /** -# return error on packet number mismatch or on unexpected packet numbers */
                    inst->ymodem.state = SHELLMATTA_YMODEM_WAIT_FOR_START;
                    shellmatta_ymodem_control(handle, YMODEM_NAK);
                    ret = SHELLMATTA_ERROR;
                }
                else
                {
                    /** -# start receiving the payload */
                    inst->ymodem.state = SHELLMATTA_YMODEM_RECEIVE_DATA;
                    inst->ymodem.byteCounter = 0u;
                }
            }
            break;
        case SHELLMATTA_YMODEM_RECEIVE_DATA:
            inst->ymodem.packet.packetData[inst->ymodem.byteCounter] = byte;
            inst->ymodem.byteCounter ++;
            
            /** -# load payload until the packet is full */
            if(inst->ymodem.byteCounter >= inst->ymodem.packet.size)
            {
                inst->ymodem.state = SHELLMATTA_YMODEM_RECEIVE_CRC;
                inst->ymodem.byteCounter = 0u;
                inst->ymodem.packet.crc = 0u;
            }

            break;
        case SHELLMATTA_YMODEM_RECEIVE_CRC:
            inst->ymodem.byteCounter ++;
            inst->ymodem.packet.crc |= (uint16_t)byte << (8u * (YMODEM_CRC_SIZE - inst->ymodem.byteCounter));

            if(inst->ymodem.byteCounter >= YMODEM_CRC_SIZE)
            {
                inst->ymodem.state = SHELLMATTA_YMODEM_WAIT_FOR_START;

                /** -# check CRC */
                computedCrc = crc16Calc((const char*)inst->ymodem.packet.packetData,
                                        inst->ymodem.packet.size);
                if(computedCrc != inst->ymodem.packet.crc)
                {
                    shellmatta_ymodem_control(handle, YMODEM_NAK);
                    ret = SHELLMATTA_ERROR;
                }
                else
                {
                    ret = processPacket(handle);
                    if(SHELLMATTA_OK == ret)
                    {
                        if(0u != inst->ymodem.packetCounter)
                        {
                            /** -# Calculate the total bytes received */
                            inst->ymodem.totalBytesReceived += inst->ymodem.packet.size;
                        }
                        if(true != inst->ymodem.pauseRequested)
                        {
                            /** -# ACK the successful packet reception */
                            shellmatta_ymodem_control(handle, YMODEM_ACK);
                            if(0u == inst->ymodem.packetCounter)
                            {
                                /** -# send addional CRC flag after packet 0 */
                                shellmatta_ymodem_control(handle, YMODEM_CRC);
                            }
                        }
                        inst->ymodem.packetCounter ++;
                    }
                }
            }
            break;
        case SHELLMATTA_YMODEM_WAIT_FOR_END:

            inst->ymodem.pollCyclesLeft = YMODEM_POLL_NUMBER;
            if(YMODEM_EOT ==  byte)
            {
                /** -# ACK the successful file reception */
                shellmatta_ymodem_control(handle, YMODEM_ACK);
                shellmatta_ymodem_control(handle, YMODEM_CRC);
            }
            break;
        default:
            /** -# unexpected state - should never happen */
            break;
    }

    return ret;
}

shellmatta_retCode_t shellmatta_ymodem_processByte(shellmatta_handle_t handle, char byte)
{
    shellmatta_instance_t *inst = (shellmatta_instance_t*)handle;
    shellmatta_retCode_t ret;

    /** -# check if session is cancelled -accept only if ymodem is between packets */
    if(('\x03' == byte) && (inst->ymodem.state == SHELLMATTA_YMODEM_WAIT_FOR_START))
    {
        /** -# explicitly reset ymodem with cancel, if it was active */
        shellmatta_ymodem_reset(handle, true);

        utils_terminateInput(inst);

        ret = SHELLMATTA_ERROR;
    }
    else
    {
        ret = ymodem_stateMachine(handle, (uint8_t)byte);
    }

    return ret;
}

shellmatta_retCode_t shellmatta_ymodem_poll(shellmatta_handle_t handle)
{
    shellmatta_retCode_t ret = SHELLMATTA_ERROR;
    shellmatta_instance_t *inst = (shellmatta_instance_t*)handle;

    switch (inst->ymodem.state)
    {
    case SHELLMATTA_YMODEM_WAIT_FOR_START:
        /** -# send ymodem symbol to start transmission */
        shellmatta_ymodem_control(handle, YMODEM_CRC);
        ret = SHELLMATTA_OK;
        break;
    case SHELLMATTA_YMODEM_WAIT_FOR_END:
        if(inst->ymodem.pollCyclesLeft > 1u)
        {
            inst->ymodem.pollCyclesLeft --;
        }
        else
        {
            ret = SHELLMATTA_OK;
            /** -# check if the received data matches the file size */
            if((inst->ymodem.totalBytesReceived < inst->ymodem.fileSize) ||
                ((inst->ymodem.totalBytesReceived - inst->ymodem.fileSize) >= YMODEM_PACKET_SIZE_1K))
            {
                ret = SHELLMATTA_ERROR;
            }

#ifdef SHELLMATTA_TRANSPORT
            /** .-# reenable transport layer optional mode */
            inst->transportLayer.suspendOptional = false;
#endif

            shellmatta_ymodem_reset(handle, false);
            inst->ymodem.transmissionCompleteCallback(handle, ret);
            (void)utils_terminateInput(inst);
        }
        break;
    default:
        /* nothing to do */
        break;
    }

    return ret;
}

/**
 * @brief                                           Initialise the ymodem prior to actually receiving data
 * @param[in, out]  handle                          shellmatta instance handle
 * @param[in]       recvBuffer                      pointer to the buffer to save the received payload in
 * @param[in]       recvPacketCallback              pointer to the file size variable
 * @param[in]       recvPacketCallback              pointer to the packet size variable
 * @param[in]       transmissionCompleteCallback    callback functions for the ymodem module
 * @return          errorcode   #SHELLMATTA_OK
 *                              #SHELLMATTA_USE_FAULT (param err)
 * @note            Disables the tranport layer if inactive or sets it to mandatory if active
*/
shellmatta_retCode_t shellmatta_ymodem_init(shellmatta_handle_t             handle,
                                            uint8_t*                        recvBuffer,
                                            shellmatta_ymodem_cancel_t      cancelCallback,
                                            shellmatta_ymodem_recvHeader_t  recvHeaderCallback,
                                            shellmatta_ymodem_recvPacket_t  recvPacketCallback,
                                            shellmatta_ymodem_complete_t    transmissionCompleteCallback)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;

    /** -# check parameters for plausibility  */
    if(     (NULL               != inst)
        &&  (SHELLMATTA_MAGIC   == inst->magic)
        &&  (NULL               != cancelCallback)
        &&  (NULL               != recvHeaderCallback)
        &&  (NULL               != recvPacketCallback)
        &&  (NULL               != transmissionCompleteCallback))
    {
        /** -# use instances buffer if no buffer is given */
        if(NULL == recvBuffer)
        {
            if(inst->bufferSize <= YMODEM_PACKET_SIZE_1K)
            {
                /** -# return use fault if buffer is too small */
                ret = SHELLMATTA_USE_FAULT;
            }
            else
            {
                inst->ymodem.packet.packetData = (uint8_t*)inst->buffer;
            }
        }

        (void)memset((void *)&inst->ymodem, 0, sizeof(shellmatta_ymodem_t));

        /** -# store buffer */
        inst->ymodem.packet.packetData = recvBuffer;

        /** -# init callbacks */
        inst->ymodem.cancelCallback = cancelCallback;
        inst->ymodem.recvHeaderCallback = recvHeaderCallback;
        inst->ymodem.recvPacketCallback = recvPacketCallback;
        inst->ymodem.transmissionCompleteCallback = transmissionCompleteCallback;

        inst->ymodem.state = SHELLMATTA_YMODEM_WAIT_FOR_START;

#ifdef SHELLMATTA_TRANSPORT
        /** -# suspend the transport layer being optional while ymodem is running */
        inst->transportLayer.suspendOptional = true;
#endif

        /** -# send initial ymodem symbol to start transmission */
        shellmatta_ymodem_control(handle, YMODEM_CRC);
    }
    else
    {
        ret = SHELLMATTA_USE_FAULT;
    }

    return ret;
}

/**
 * @brief                   pauses the shellmatta reception until #shellmatta_ymodem_resume resume is called
 * @param[in, out]  handle  shellmatta instance handle
 * @return      errorcode   #SHELLMATTA_OK
 *                          #SHELLMATTA_USE_FAULT (param err)
 * @note                    This can be used when an application needs processing time before accepting the next packet
*/
shellmatta_retCode_t shellmatta_ymodem_pause(shellmatta_handle_t handle)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;

    /** -# check parameters for plausibility  */
    if(     (NULL               != inst)
        &&  (SHELLMATTA_MAGIC   == inst->magic))
    {
        inst->ymodem.pauseRequested = true;
    }
    else
    {
        ret = SHELLMATTA_USE_FAULT;
    }

    return ret;
}

/**
 * @brief                   Resume the ymodem module
 * @param[in, out]  handle  shellmatta instance handle
 * @return      errorcode   #SHELLMATTA_OK
 *                          #SHELLMATTA_USE_FAULT (param err)
 * @note                    This can be used when an application needs processing time before accepting the next packet
*/
shellmatta_retCode_t shellmatta_ymodem_resume(shellmatta_handle_t handle)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;

    /** -# check parameters for plausibility  */
    if(     (NULL               != inst)
        &&  (SHELLMATTA_MAGIC   == inst->magic))
    {
        /** -# ACK the successful packet reception */
        shellmatta_ymodem_control(handle, YMODEM_ACK);
        if(1u == inst->ymodem.packetCounter)
        {
            /** -# send addional CRC flag after packet 0 */
            shellmatta_ymodem_control(handle, YMODEM_CRC);
        }
        inst->ymodem.pauseRequested = false;
    }
    else
    {
        ret = SHELLMATTA_USE_FAULT;
    }

    return ret;
}

/**
 * @brief                   Resets the ymodem module
 * @param[in]   doCancel    Set this flag to execute the cancel-callback function within the ymodem-reset function
 * @return      errorcode   #SHELLMATTA_OK
 *                          #SHELLMATTA_USE_FAULT (param err)
 * @note                    call this function after file transmission is done or cancelled
*/
shellmatta_retCode_t shellmatta_ymodem_cancel(shellmatta_handle_t handle, bool doCancel)
{
    shellmatta_retCode_t ret = SHELLMATTA_OK;
    shellmatta_instance_t *inst = (shellmatta_instance_t *)handle;

    /** -# check parameters for plausibility  */
    if(     (NULL               != inst)
        &&  (SHELLMATTA_MAGIC   == inst->magic))
    {
        shellmatta_ymodem_reset(handle, doCancel);
        /* clear any possibly leftover inputs */
        utils_clearInput((shellmatta_instance_t*)handle);
    }
    else
    {
        ret = SHELLMATTA_USE_FAULT;
    }

    return ret;
}