/*
 * xmodem.c
 *
 *  Created on: Jun 23, 2023
 *      Author: true
 *
 * originally Copyright (c) 2018 Ferenc Nemeth - https://github.com/ferenc-nemeth - 21 Dec 2018
 * modified for use with HK32F and my coding style.
 *
 * under MIT license.
 */

#include "xmodem.h"
#include "flash.h"
#include "usart.h"



/* Global variables. */
static uint8_t  xmodem_packet_number;         		/**< Packet number counter. */
static uint32_t xmodem_flash_w_addr; 				/**< Address where we have to write. */
static uint8_t  x_first_packet_rcvd;   				/**< First packet or not. */


/* Local functions. */
static uint16_t xmodem_calc_crc(uint8_t *data, uint16_t length);
static BL_xmodem_status xmodem_handle_packet(uint8_t size);
static BL_xmodem_status xmodem_error_handler(uint8_t *error_number, uint8_t max_error_number);



/**
 * @brief   This function is the base of the Xmodem protocol.
 *          When we receive a header from UART, it decides what action it shall take.
 * @param   void
 * @return  void
 */
BL_xmodem_status xmodem_receive(void)
{
	volatile BL_xmodem_status status = X_OK;
	uint8_t error_number = 0u;

	x_first_packet_rcvd  = 0;
	xmodem_packet_number = 1u;
	xmodem_flash_w_addr  = USER_APP_START_ADDR;

	uint8_t x_c_sent 	 = 0;

	/* Loop until there isn't any error (or until we jump to the user application). */
	while (status == X_OK) {
		uint8_t header = 0x00u;

		/* Get the header from UART. */
		BL_comm_status status_c = comm_rx(&header, 1);

		/* Spam the host (until we receive something) with ACSII "C", to notify it, we want to use CRC-16. */
		if ((status_c != COMM_OK) && (x_first_packet_rcvd == 0)) {
			// return with timeout if we've tried too many times
			// this is so main can repeat our banner
			x_c_sent++;
			if (x_c_sent >= 240) {
				return X_TIMEOUT;
			}

			comm_tx_byte(X_C);
		}

		/* Uart timeout or any other errors. */
		else if ((status_c != COMM_OK) && (x_first_packet_rcvd)) {
			status = xmodem_error_handler(&error_number, X_MAX_ERRORS);
		}

		/* Do nothing. */
		else {

		}

		/* The header can be: SOH, STX, EOT and CAN. */
		BL_xmodem_status packet_status = X_ERROR;

		switch(header) {
			/* 128 or 1024 bytes of data. */
			case X_SOH:
			case X_STX: {
				/* If the handling was successful, then send an ACK. */
				packet_status = xmodem_handle_packet(header);
				if (packet_status == X_OK) {
					comm_tx_byte(X_ACK);
				}

				/* If the error was flash related, then immediately set the error counter to max (graceful abort). */
				else if (packet_status == X_ERROR_FLASH) {
					error_number = X_MAX_ERRORS;
					status = xmodem_error_handler(&error_number, X_MAX_ERRORS);
				}

				/* Error while processing the packet, either send a NAK or do graceful abort. */
				else {
					status = xmodem_error_handler(&error_number, X_MAX_ERRORS);
				}

				break;
			}

			/* End of Transmission. */
			case X_EOT: {
				/* ACK, feedback to user (as a text), then jump to user application. */
				comm_tx_byte(X_ACK);
				return X_OK;
			}

			/* Abort from host. */
			case X_CAN: {
				return X_ERROR;
			}

			/* Wrong header. */
			default: {
				if (status_c == COMM_OK) {
					status = xmodem_error_handler(&error_number, X_MAX_ERRORS);
				}
				break;
			}
		}
	}

	return status;
}

/**
 * @brief   Calculates the CRC-16 for the input package.
 * @param   *data:  Array of the data which we want to calculate.
 * @param   length: Size of the data, either 128 or 1024 bytes.
 * @return  status: The calculated CRC.
 */
static uint16_t xmodem_calc_crc(uint8_t *data, uint16_t length)
{
    uint16_t crc = 0u;

    while (length) {
        length--;
        crc = crc ^ ((uint16_t)*data++ << 8u);
        for (uint8_t i = 0u; i < 8u; i++) {
            if (crc & 0x8000u) {
                crc = (crc << 1u) ^ 0x1021u;
            } else {
                crc = crc << 1u;
            }
        }
    }

    return crc;
}

/**
 * @brief   This function handles the data packet we get from the xmodem protocol.
 * @param   header: SOH or STX.
 * @return  status: Report about the packet.
 */
static BL_xmodem_status xmodem_handle_packet(uint8_t header)
{
	BL_xmodem_status status = X_OK;
	uint16_t size = 0;
	uint32_t i;

	/* 2 bytes for packet number, 1024 for data, 2 for CRC*/
	uint8_t rcvd_packet_number[X_PACKET_NUMBER_SIZE];
	uint8_t rcvd_packet_data[X_PACKET_1024_SIZE];
	uint8_t rcvd_packet_crc[X_PACKET_CRC_SIZE];

	/* Get the size of the data. */
	if (header == X_SOH) {
		size = X_PACKET_128_SIZE;
	} else if (header == X_STX) {
		size = X_PACKET_1024_SIZE;
	} else {
		/* Wrong header type. This shoudn't be possible... */
		status |= X_ERROR;
	}

	BL_comm_status status_c = COMM_OK;

	/* Get the packet number, data and CRC from UART. */
	status_c |= comm_rx(&rcvd_packet_number[0u], X_PACKET_NUMBER_SIZE);
	status_c |= comm_rx(&rcvd_packet_data[0u], size);
	status_c |= comm_rx(&rcvd_packet_crc[0u], X_PACKET_CRC_SIZE);

	/* Merge the two bytes of CRC. */
	uint16_t crc_rcvd = ((uint16_t)rcvd_packet_crc[X_PACKET_CRC_HIGH_INDEX] << 8u)
			| ((uint16_t)rcvd_packet_crc[X_PACKET_CRC_LOW_INDEX]);

	/* We calculate it too. */
	uint16_t crc_calc = xmodem_calc_crc(&rcvd_packet_data[0u], size);

	/* Communication error. */
	if (status_c != COMM_OK) {
		status |= X_ERROR_UART;
	}

	/* If it is the first packet, then mark it. If you want to erase all flash, do it here. */
	if ((status == X_OK)) { // && (x_first_packet_rcvd == 0)) {
		x_first_packet_rcvd = 1;
	}

	/* Error handling and flashing. */
	if (status == X_OK) {
		if (xmodem_packet_number != rcvd_packet_number[0u]) {
			/* Packet number counter mismatch. */
			status |= X_ERROR_NUMBER;
		}

		if ((rcvd_packet_number[X_PACKET_NUMBER_INDEX] + rcvd_packet_number[X_PACKET_NUMBER_COMPLEMENT_INDEX]) != 255) {
			/* The sum of the packet number and packet number complement aren't 255. */
			/* The sum always has to be 255. */
			status |= X_ERROR_NUMBER;
		}

		if (crc_calc != crc_rcvd) {
			/* The calculated and received CRC are different. */
			status |= X_ERROR_CRC;
		}
	}

	/* Do the actual flashing if there weren't any errors. */
	if (status == X_OK) {
		/* first, erase pages depending on size */
		for (i = 0; i < size; i += FLASH_ERASE_PAGE_SIZE) {
			if (flash_erase_page(xmodem_flash_w_addr + i) != FLASH_OK) {
				/* Flashing error. */
				status |= X_ERROR_FLASH;
			}
		}

		/* then write data to these cleared pages */
		if (status == X_OK) {
			if (flash_write(xmodem_flash_w_addr, rcvd_packet_data, size) != FLASH_OK) {
				/* Flashing error. */
				status |= X_ERROR_FLASH;
			}
		}
	}

	/* Raise the packet number and the address counters (if there weren't any errors). */
	if (status == X_OK) {
		xmodem_packet_number++;
		xmodem_flash_w_addr += size;
	}

	return status;
}

/**
 * @brief   Handles the xmodem error.
 *          Raises the error counter, then if the number of the errors reached critical, do a graceful abort, otherwise send a NAK.
 * @param   *error_number:    Number of current errors (passed as a pointer).
 * @param   max_error_number: Maximal allowed number of errors.
 * @return  status: X_ERROR in case of too many errors, X_OK otherwise.
 */
static BL_xmodem_status xmodem_error_handler(uint8_t *error_number, uint8_t max_error_number)
{
	BL_xmodem_status status = X_OK;

	/* Raise the error counter. */
	(*error_number)++;

	/* If the counter reached the max value, then abort. */
	if ((*error_number) >= max_error_number) {
		/* Graceful abort. */
		comm_tx_byte(X_CAN);
		comm_tx_byte(X_CAN);
		status = X_ERROR;
	}

	/* Otherwise send a NAK for a repeat. */
	else {
		comm_tx_byte(X_NAK);
		status = X_OK;
	}

	return status;
}
