#include "rs_boot_transport.h"

#define SOF0 0xA5U
#define SOF1 0x5AU
#define VERSION 1U
#define HEADER 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 rsbt_u16_t u16(const rsbt_byte_t *p)
{ return (p[0] & 0xFFU) | ((p[1] & 0xFFU) << 8); }

static rsbt_u32_t u32(const rsbt_byte_t *p)
{
    return (rsbt_u32_t)(p[0] & 0xFFU) |
           ((rsbt_u32_t)(p[1] & 0xFFU) << 8) |
           ((rsbt_u32_t)(p[2] & 0xFFU) << 16) |
           ((rsbt_u32_t)(p[3] & 0xFFU) << 24);
}

static void put32(rsbt_byte_t *p, rsbt_u32_t v)
{
    p[0] = (rsbt_byte_t)(v & 0xFFU); p[1] = (rsbt_byte_t)((v >> 8) & 0xFFU);
    p[2] = (rsbt_byte_t)((v >> 16) & 0xFFU); p[3] = (rsbt_byte_t)((v >> 24) & 0xFFU);
}

rsbt_u32_t rsbt_crc32(const rsbt_byte_t *p, rsbt_u16_t n)
{
    rsbt_u32_t crc = 0xFFFFFFFFUL;
    rsbt_u16_t i, bit;
    for (i = 0U; i < n; ++i) {
        crc ^= (rsbt_u32_t)(p[i] & 0xFFU);
        for (bit = 0U; bit < 8U; ++bit)
            crc = (crc >> 1) ^ ((crc & 1U) ? 0xEDB88320UL : 0U);
    }
    return crc ^ 0xFFFFFFFFUL;
}

void rsbt_abort(rsbt_t *b)
{
    if (!b) return;
    b->state = RSBT_IDLE; b->image_size = 0UL; b->flags = 0U;
    b->image_crc = 0UL; b->next_offset = 0UL;
}

int rsbt_init(rsbt_t *b, const rsbt_config_t *c, const rsbt_port_t *p, void *u)
{
    rsbt_u16_t i;
    if (!b || !c || !p || !c->app_size || !p->transmit || !p->erase || !p->write || !p->read) return 0;
    for (i = 0U; i < RSBT_MAX_FRAME; ++i) { b->parser[i] = 0U; b->tx[i] = 0U; }
    b->config = *c; b->port = *p; b->user = u; b->parser_length = 0U; b->crc_errors = 0UL;
    if (!b->config.max_block || b->config.max_block > RSBT_MAX_BLOCK) b->config.max_block = RSBT_MAX_BLOCK;
    rsbt_abort(b); return 1;
}

static int send_frame(rsbt_t *b, rsbt_byte_t type, rsbt_u16_t seq,
                      const rsbt_byte_t *payload, rsbt_u16_t size)
{
    rsbt_u16_t i, total; rsbt_u32_t crc;
    if (size > RSBT_MAX_PAYLOAD) return 0;
    total = HEADER + size + CRC_SIZE;
    b->tx[0]=SOF0; b->tx[1]=SOF1; b->tx[2]=VERSION; b->tx[3]=type;
    b->tx[4]=(seq>>8)&0xFFU; b->tx[5]=seq&0xFFU;
    b->tx[6]=(size>>8)&0xFFU; b->tx[7]=size&0xFFU;
    for(i=0U;i<size;++i) b->tx[HEADER+i]=payload[i]&0xFFU;
    crc=rsbt_crc32(&b->tx[2],6U+size); put32(&b->tx[HEADER+size],crc);
    return b->port.transmit(b->user,b->tx,total);
}

static void nack(rsbt_t *b, rsbt_byte_t req, rsbt_u16_t seq, rsbt_u16_t result)
{
    rsbt_byte_t p[3]; p[0]=result&0xFFU; p[1]=(result>>8)&0xFFU; p[2]=req;
    (void)send_frame(b,MSG_NACK,seq,p,3U);
}

static int matches(rsbt_t *b, rsbt_u32_t off, const rsbt_byte_t *p, rsbt_u16_t n)
{
    rsbt_u16_t done=0U,i,take;
    while(done<n){take=n-done;if(take>32U)take=32U;
        if(!b->port.read(b->user,off+done,b->scratch,take))return 0;
        for(i=0U;i<take;++i)if((b->scratch[i]&0xFFU)!=(p[done+i]&0xFFU))return 0;
        done+=take;} return 1;
}

static void process(rsbt_t *b, rsbt_byte_t type, rsbt_u16_t seq,
                    const rsbt_byte_t *p, rsbt_u16_t size)
{
    rsbt_u32_t off,base; rsbt_u16_t n,sum,i;
    if(type==MSG_BEGIN){if(size!=16U){nack(b,type,seq,2U);return;}
        b->image_size=u32(p);b->image_crc=u32(&p[4]);b->flags=p[8]&RSBT_FLAG_MIRROR_RAM;base=u32(&p[12]);
        if((b->flags&RSBT_FLAG_MIRROR_RAM)&&!b->port.write_ram){nack(b,type,seq,0x11U);return;}
        if(!b->image_size||b->image_size>b->config.app_size||(base&&base!=b->config.app_address)){nack(b,type,seq,1U);return;}
        b->next_offset=0UL;b->state=RSBT_RECEIVING;
        if(!b->port.erase(b->user,b->image_size)){b->state=RSBT_FAILED;nack(b,type,seq,7U);return;}
        put32(b->scratch,0UL);(void)send_frame(b,type,seq,b->scratch,4U);return;}
    if(type==MSG_DATA){if(b->state!=RSBT_RECEIVING||size<8U){nack(b,type,seq,5U);return;}
        off=u32(p);n=u16(&p[4]);sum=0U;for(i=0U;i<n&&8U+i<size;++i)sum=(sum+(p[8U+i]&0xFFU))&0xFFFFU;
        if(!n||n>b->config.max_block||size!=8U+n||off>b->image_size||n>b->image_size-off){nack(b,type,seq,2U);return;}
        if(sum!=u16(&p[6])){nack(b,type,seq,1U);return;}
        if(off<b->next_offset){if(off+n>b->next_offset||!matches(b,off,&p[8],n)){nack(b,type,seq,1U);return;}}
        else {if(off!=b->next_offset||!b->port.write(b->user,off,&p[8],n)||
            ((b->flags&RSBT_FLAG_MIRROR_RAM)&&!b->port.write_ram(b->user,off,&p[8],n))){b->state=RSBT_FAILED;nack(b,type,seq,7U);return;}b->next_offset+=n;}
        put32(b->scratch,b->next_offset);(void)send_frame(b,type,seq,b->scratch,4U);return;}
    if(type==MSG_END){if(b->state!=RSBT_RECEIVING||size!=8U||u32(p)!=b->image_size||u32(&p[4])!=b->image_crc||b->next_offset!=b->image_size||(b->port.image_valid&&!b->port.image_valid(b->user,b->image_size,b->image_crc))){b->state=RSBT_FAILED;nack(b,type,seq,1U);return;}
        b->state=RSBT_READY;(void)send_frame(b,type,seq,0,0U);if(b->port.reboot)b->port.reboot(b->user);return;}
    if(type==MSG_ABORT){rsbt_abort(b);(void)send_frame(b,type,seq,0,0U);return;}
    if(type==MSG_STATUS){b->scratch[0]=b->state;put32(&b->scratch[1],b->next_offset);b->scratch[5]=b->flags;(void)send_frame(b,type,seq,b->scratch,6U);return;}
    nack(b,type,seq,0x11U);
}

static void discard(rsbt_t *b, rsbt_u16_t n)
{ rsbt_u16_t i;if(n>=b->parser_length){b->parser_length=0U;return;}for(i=n;i<b->parser_length;++i)b->parser[i-n]=b->parser[i];b->parser_length-=n; }

void rsbt_feed(rsbt_t *b, rsbt_byte_t byte)
{
    rsbt_u16_t payload,total,seq; rsbt_u32_t got;
    if(!b)return;if(b->parser_length==RSBT_MAX_FRAME)b->parser_length=0U;
    b->parser[b->parser_length++]=byte&0xFFU;
    for(;;){while(b->parser_length>=2U&&(b->parser[0]!=SOF0||b->parser[1]!=SOF1))discard(b,1U);
        if(b->parser_length<HEADER)return;if(b->parser[2]!=VERSION){discard(b,1U);continue;}
        payload=((b->parser[6]&0xFFU)<<8)|(b->parser[7]&0xFFU);if(payload>RSBT_MAX_PAYLOAD){discard(b,1U);continue;}
        total=HEADER+payload+CRC_SIZE;if(b->parser_length<total)return;got=u32(&b->parser[HEADER+payload]);
        if(got!=rsbt_crc32(&b->parser[2],6U+payload)){++b->crc_errors;discard(b,1U);continue;}
        seq=((b->parser[4]&0xFFU)<<8)|(b->parser[5]&0xFFU);process(b,b->parser[3],seq,&b->parser[HEADER],payload);discard(b,total);}
}