#include "rs485_boot.h"

#include <string.h>

#define SOF0 0xA5U
#define SOF1 0x5AU
#define HEADER_SIZE 8U
#define CRC_SIZE 4U
#define MSG_BEGIN 0x0BU
#define MSG_DATA 0x0CU
#define MSG_END 0x0DU
#define MSG_ABORT 0x0EU
#define MSG_STATUS 0x0FU
#define MSG_NACK 0x80U

static uint16_t get_u16(const uint8_t *p) { return (uint16_t)(p[0] | ((uint16_t)p[1] << 8U)); }
static uint32_t get_u32(const uint8_t *p) { return (uint32_t)p[0] | ((uint32_t)p[1] << 8U) | ((uint32_t)p[2] << 16U) | ((uint32_t)p[3] << 24U); }
static void put_u32(uint8_t *p, uint32_t v) { p[0]=(uint8_t)v; p[1]=(uint8_t)(v>>8U); p[2]=(uint8_t)(v>>16U); p[3]=(uint8_t)(v>>24U); }

uint32_t rs485_boot_crc32(const uint8_t *data, size_t length)
{
    uint32_t crc = 0xFFFFFFFFUL;
    size_t i;
    for (i = 0U; i < length; ++i) {
        uint8_t bit;
        crc ^= data[i];
        for (bit = 0U; bit < 8U; ++bit) crc = (crc >> 1U) ^ ((crc & 1U) ? 0xEDB88320UL : 0U);
    }
    return crc ^ 0xFFFFFFFFUL;
}

void rs485_boot_abort(rs485_boot_t *boot)
{
    if (boot == NULL) return;
    boot->state = RS485_BOOT_IDLE;
    boot->image_size = 0U;
    boot->image_crc32 = 0U;
    boot->next_offset = 0U;
}

bool rs485_boot_init(rs485_boot_t *boot, const rs485_boot_config_t *config,
                     const rs485_boot_port_t *port, void *port_user)
{
    if ((boot == NULL) || (config == NULL) || (port == NULL) ||
        (config->app_size == 0U) || (port->transmit == NULL) ||
        (port->erase == NULL) || (port->write == NULL) || (port->read == NULL)) return false;
    (void)memset(boot, 0, sizeof(*boot));
    boot->config = *config;
    boot->port = *port;
    boot->port_user = port_user;
    if ((boot->config.max_block_size == 0U) || (boot->config.max_block_size > RS485_BOOT_MAX_BLOCK))
        boot->config.max_block_size = RS485_BOOT_MAX_BLOCK;
    return true;
}

static bool send_frame(rs485_boot_t *b, uint8_t type, uint16_t seq, const uint8_t *payload, uint16_t size)
{
    uint32_t crc;
    uint16_t i;
    const uint16_t total = (uint16_t)(HEADER_SIZE + size + CRC_SIZE);
    if (size > RS485_BOOT_MAX_PAYLOAD) return false;
    b->tx[0]=SOF0; b->tx[1]=SOF1; b->tx[2]=RS485_BOOT_VERSION; b->tx[3]=type;
    b->tx[4]=(uint8_t)(seq>>8U); b->tx[5]=(uint8_t)seq; b->tx[6]=(uint8_t)(size>>8U); b->tx[7]=(uint8_t)size;
    for (i=0U; i<size; ++i) b->tx[HEADER_SIZE+i]=payload[i];
    crc=rs485_boot_crc32(&b->tx[2], (size_t)6U+size);
    put_u32(&b->tx[HEADER_SIZE+size], crc);
    return b->port.transmit(b->port_user, b->tx, total);
}

static void nack(rs485_boot_t *b, uint8_t request, uint16_t seq, uint16_t result)
{
    uint8_t p[3]={(uint8_t)result,(uint8_t)(result>>8U),request};
    (void)send_frame(b, MSG_NACK, seq, p, 3U);
}

static bool matches(rs485_boot_t *b, uint32_t offset, const uint8_t *data, uint16_t size)
{
    uint16_t done=0U;
    while (done<size) {
        uint16_t n=(uint16_t)(size-done); uint16_t i;
        if (n>sizeof(b->scratch)) n=sizeof(b->scratch);
        if (!b->port.read(b->port_user, offset+done, b->scratch, n)) return false;
        for(i=0U;i<n;++i) if(b->scratch[i]!=data[done+i]) return false;
        done=(uint16_t)(done+n);
    }
    return true;
}

static void process(rs485_boot_t *b, uint8_t type, uint16_t seq, const uint8_t *p, uint16_t size)
{
    if (type==MSG_BEGIN) {
        uint32_t base;
        if(size!=16U){nack(b,type,seq,RS485_BOOT_INVALID_LENGTH);return;}
        b->image_size=get_u32(p); b->image_crc32=get_u32(&p[4]); base=get_u32(&p[12]);
        if((b->image_size==0U)||(b->image_size>b->config.app_size)||((base!=0U)&&(base!=b->config.app_address))){nack(b,type,seq,RS485_BOOT_INVALID_ARGUMENT);return;}
        b->next_offset=0U; b->state=RS485_BOOT_RECEIVING;
        if(!b->port.erase(b->port_user,b->image_size)){b->state=RS485_BOOT_FAILED;nack(b,type,seq,RS485_BOOT_INTERNAL);return;}
        put_u32(b->scratch,0U); (void)send_frame(b,type,seq,b->scratch,4U); return;
    }
    if(type==MSG_DATA){
        uint32_t offset; uint16_t n,expected=0U,i;
        if((b->state!=RS485_BOOT_RECEIVING)||(size<RS485_BOOT_DATA_HEADER_SIZE)){nack(b,type,seq,RS485_BOOT_BUSY);return;}
        offset=get_u32(p); n=get_u16(&p[4]); for(i=0U;i<n && (uint16_t)(RS485_BOOT_DATA_HEADER_SIZE+i)<size;++i) expected=(uint16_t)(expected+p[RS485_BOOT_DATA_HEADER_SIZE+i]);
        if((n==0U)||(n>b->config.max_block_size)||(size!=(uint16_t)(RS485_BOOT_DATA_HEADER_SIZE+n))||((offset&1U)!=0U)||(offset>b->image_size)||(n>b->image_size-offset)){nack(b,type,seq,RS485_BOOT_INVALID_LENGTH);return;}
        if(expected!=get_u16(&p[6])){nack(b,type,seq,RS485_BOOT_INVALID_ARGUMENT);return;}
        if(offset<b->next_offset){if((offset+n>b->next_offset)||!matches(b,offset,&p[8],n)){nack(b,type,seq,RS485_BOOT_INVALID_ARGUMENT);return;}}
        else {if(offset!=b->next_offset){nack(b,type,seq,RS485_BOOT_INVALID_ARGUMENT);return;} if(!b->port.write(b->port_user,offset,&p[8],n)){b->state=RS485_BOOT_FAILED;nack(b,type,seq,RS485_BOOT_INTERNAL);return;} b->next_offset+=n;}
        put_u32(b->scratch,b->next_offset); (void)send_frame(b,type,seq,b->scratch,4U); return;
    }
    if(type==MSG_END){
        if((b->state!=RS485_BOOT_RECEIVING)||(size!=8U)||(get_u32(p)!=b->image_size)||(get_u32(&p[4])!=b->image_crc32)||(b->next_offset!=b->image_size)||((b->port.image_valid!=NULL)&&!b->port.image_valid(b->port_user,b->image_size,b->image_crc32))){b->state=RS485_BOOT_FAILED;nack(b,type,seq,RS485_BOOT_INVALID_ARGUMENT);return;}
        b->state=RS485_BOOT_READY; (void)send_frame(b,type,seq,NULL,0U); if(b->port.reboot!=NULL)b->port.reboot(b->port_user); return;
    }
    if(type==MSG_ABORT){rs485_boot_abort(b);(void)send_frame(b,type,seq,NULL,0U);return;}
    if(type==MSG_STATUS){b->scratch[0]=(uint8_t)b->state;put_u32(&b->scratch[1],b->next_offset);(void)send_frame(b,type,seq,b->scratch,5U);return;}
    nack(b,type,seq,RS485_BOOT_UNSUPPORTED);
}

static void discard(rs485_boot_t *b, uint16_t n)
{
    if(n>=b->parser_length){b->parser_length=0U;return;}
    b->parser_length=(uint16_t)(b->parser_length-n);
    (void)memmove(b->parser,&b->parser[n],b->parser_length);
}

void rs485_boot_feed(rs485_boot_t *b, const uint8_t *data, size_t length)
{
    size_t input;
    if((b==NULL)||((data==NULL)&&(length!=0U)))return;
    for(input=0U;input<length;++input){
        if(b->parser_length==RS485_BOOT_MAX_FRAME)b->parser_length=0U;
        b->parser[b->parser_length++]=data[input];
        for(;;){
            uint16_t payload,total,i; uint32_t got;
            while((b->parser_length>=2U)&&((b->parser[0]!=SOF0)||(b->parser[1]!=SOF1)))discard(b,1U);
            if(b->parser_length<HEADER_SIZE)break;
            if(b->parser[2]!=RS485_BOOT_VERSION){discard(b,1U);continue;}
            payload=(uint16_t)(((uint16_t)b->parser[6]<<8U)|b->parser[7]);
            if(payload>RS485_BOOT_MAX_PAYLOAD){discard(b,1U);continue;}
            total=(uint16_t)(HEADER_SIZE+payload+CRC_SIZE); if(b->parser_length<total)break;
            got=get_u32(&b->parser[HEADER_SIZE+payload]);
            if(got!=rs485_boot_crc32(&b->parser[2],(size_t)6U+payload)){++b->crc_errors;discard(b,1U);continue;}
            i=(uint16_t)(((uint16_t)b->parser[4]<<8U)|b->parser[5]);
            process(b,b->parser[3],i,&b->parser[HEADER_SIZE],payload); discard(b,total);
        }
    }
}