/**
 * @file app_ota.c
 * @brief BLE OTA receiver implementation.
 *
 * Flow: OTA_BEGIN (size/crc/version) -> OTA_DATA* (strictly sequential
 * offsets, acked per flushed 4KB sector) -> OTA_END (flush remainder, verify
 * whole-image CRC32 against flash, update bootsetting, reset).
 *
 * Notes:
 * - Runs in the BLE schedule task context; vTaskDelay() is available.
 * - Qflash erase/write disable interrupts for tens of ms per sector (SDK
 *   driver behavior, the flash algorithm executes from RAM). The BLE link
 *   survives this via the 5 s connection supervision timeout.
 * - Qflash_Write() requires a 4-byte aligned length; the tail of the image
 *   is padded with 0xFF up to the next word boundary before writing.
 */
#include "app_ota.h"
#include "app_ble_proto.h"
#include "dfu_layout.h"
#include "boot_crc32.h"
#include "n32wb03x.h"
#include "n32wb03x_qflash.h"

#include <string.h>
#include <stddef.h>

#include "FreeRTOS.h"
#include "task.h"

/* armlink execution-region base symbol: address of this running image */
extern uint32_t Image$$ER_IROM1$$Base;

#define OTA_SECTOR_BUF_SIZE     CAIIC_FLASH_SECTOR_SIZE   /* 4KB */
#define OTA_MIN_IMAGE_SIZE      4u
#define OTA_RESET_DELAY_MS      200u

/* OTA session state */
static uint8_t  s_active;                              /* session in progress */
static uint32_t s_target_base;                         /* opposite bank base */
static uint32_t s_total_size;
static uint32_t s_image_crc32;
static uint32_t s_version;
static uint32_t s_offset;                              /* bytes received so far */
static uint32_t s_sector_idx;                          /* sector being filled */
static uint8_t  s_sector_buf[OTA_SECTOR_BUF_SIZE];
static uint32_t s_sector_fill;
static uint8_t  s_qflash_ready;

/* ------------------------------------------------------------------ */
/* Helpers                                                            */
/* ------------------------------------------------------------------ */

static uint32_t ota_rd32(const uint8_t* p)
{
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8) |
           ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

/**
 * @brief Send an OTA_RSP frame: {cmd_echo, status, offset(LE)}.
 */
static void ota_respond(uint8_t seq, uint8_t cmd_echo, uint8_t status, uint32_t offset)
{
    uint8_t p[6];

    p[0] = cmd_echo;
    p[1] = status;
    p[2] = (uint8_t)(offset & 0xFFu);
    p[3] = (uint8_t)((offset >> 8) & 0xFFu);
    p[4] = (uint8_t)((offset >> 16) & 0xFFu);
    p[5] = (uint8_t)((offset >> 24) & 0xFFu);
    (void)app_ble_proto_send_frame(BLE_FRAME_OTA_RSP, seq, p, sizeof(p));
}

uint32_t app_ota_current_bank_base(void)
{
    return (uint32_t)&Image$$ER_IROM1$$Base;
}

void app_ota_abort(void)
{
    s_active = 0;
    s_offset = 0;
    s_sector_idx = 0;
    s_sector_fill = 0;
}

/**
 * @brief Erase and write the current sector buffer (full 4KB).
 * @return 0 on success, BLE_OTA_ST_FLASH_FAIL on error.
 */
static uint8_t ota_flush_full_sector(void)
{
    uint32_t addr = s_target_base + s_sector_idx * OTA_SECTOR_BUF_SIZE;

    if (Qflash_Erase_Sector(addr) != 0)
    {
        return BLE_OTA_ST_FLASH_FAIL;
    }
    if (Qflash_Write(addr, s_sector_buf, OTA_SECTOR_BUF_SIZE) != 0)
    {
        return BLE_OTA_ST_FLASH_FAIL;
    }
    return 0;
}

/* ------------------------------------------------------------------ */
/* Frame handlers                                                     */
/* ------------------------------------------------------------------ */

/**
 * @brief OTA_BEGIN: {total_size u32, image_crc32 u32, version u32}.
 */
static void ota_on_begin(uint8_t seq, const uint8_t* p, uint16_t len)
{
    uint32_t total, crc, version;
    uint32_t self, target;

    if (len != 12u)
    {
        ota_respond(seq, BLE_FRAME_OTA_BEGIN, BLE_OTA_ST_BAD_FRAME, 0);
        return;
    }

    total   = ota_rd32(p);
    crc     = ota_rd32(p + 4);
    version = ota_rd32(p + 8);

    /* Pick the opposite bank as the update target */
    self = app_ota_current_bank_base();
    if (self == CAIIC_APP1_BASE)
    {
        target = CAIIC_APP2_BASE;
    }
    else if (self == CAIIC_APP2_BASE)
    {
        target = CAIIC_APP1_BASE;
    }
    else
    {
        ota_respond(seq, BLE_FRAME_OTA_BEGIN, BLE_OTA_ST_BAD_STATE, 0);
        return;
    }

    if (total < OTA_MIN_IMAGE_SIZE || total > CAIIC_APP_BANK_SIZE)
    {
        ota_respond(seq, BLE_FRAME_OTA_BEGIN, BLE_OTA_ST_SIZE_TOO_BIG, 0);
        return;
    }

    if (!s_qflash_ready)
    {
        Qflash_Init();   /* copy the flash algorithm to RAM (once) */
        s_qflash_ready = 1;
    }

    s_active      = 1;
    s_target_base = target;
    s_total_size  = total;
    s_image_crc32 = crc;
    s_version     = version;
    s_offset      = 0;
    s_sector_idx  = 0;
    s_sector_fill = 0;

    ota_respond(seq, BLE_FRAME_OTA_BEGIN, BLE_OTA_ST_OK, 0);
}

/**
 * @brief OTA_DATA: {offset u32 + data}. Data must arrive strictly in order.
 *        An OTA_RSP(ok) is sent per flushed 4KB sector (flow control /
 *        progress); errors are acked immediately and the state is kept so
 *        the peer can resend.
 */
static void ota_on_data(uint8_t seq, const uint8_t* p, uint16_t len)
{
    uint32_t off;
    const uint8_t* data;
    uint32_t dlen;

    if (!s_active)
    {
        ota_respond(seq, BLE_FRAME_OTA_DATA, BLE_OTA_ST_BAD_STATE, 0);
        return;
    }
    if (len < 4u)
    {
        ota_respond(seq, BLE_FRAME_OTA_DATA, BLE_OTA_ST_BAD_FRAME, s_offset);
        return;
    }

    off  = ota_rd32(p);
    data = p + 4;
    dlen = (uint32_t)len - 4u;

    if (off != s_offset)
    {
        /* out-of-order: report the expected offset, keep the session */
        ota_respond(seq, BLE_FRAME_OTA_DATA, BLE_OTA_ST_BAD_STATE, s_offset);
        return;
    }
    if (s_offset + dlen > s_total_size)
    {
        ota_respond(seq, BLE_FRAME_OTA_DATA, BLE_OTA_ST_SIZE_TOO_BIG, s_offset);
        return;
    }

    while (dlen != 0)
    {
        uint32_t n = OTA_SECTOR_BUF_SIZE - s_sector_fill;
        if (n > dlen)
        {
            n = dlen;
        }
        memcpy(s_sector_buf + s_sector_fill, data, n);
        s_sector_fill += n;
        data          += n;
        dlen          -= n;
        s_offset      += n;

        if (s_sector_fill == OTA_SECTOR_BUF_SIZE)
        {
            /* Full sector: erase + program, then ack (sector-level flow control).
             * Interrupts are disabled inside Qflash for tens of ms here. */
            uint8_t err = ota_flush_full_sector();
            if (err != 0)
            {
                ota_respond(seq, BLE_FRAME_OTA_DATA, err, s_offset);
                app_ota_abort();
                return;
            }
            s_sector_idx++;
            s_sector_fill = 0;
            ota_respond(seq, BLE_FRAME_OTA_DATA, BLE_OTA_ST_OK, s_offset);
        }
    }
}

/**
 * @brief Program the bootsetting record and verify by read-back.
 * @return 0 on success, BLE_OTA_ST_FLASH_FAIL on error.
 */
static uint8_t ota_update_bootsetting(void)
{
    caiic_bootsetting_t bs;
    caiic_bank_t* bank;

    memcpy(&bs, (const void*)CAIIC_BOOTSETTING_ADDR, sizeof(bs));

    /* Validate the existing record; rebuild from scratch when invalid */
    if (bs.magic != CAIIC_BOOT_MAGIC ||
        bs.crc32 != caiic_crc32((const uint8_t*)CAIIC_BOOTSETTING_ADDR,
                                (uint32_t)offsetof(caiic_bootsetting_t, crc32)))
    {
        memset(&bs, 0, sizeof(bs));
    }

    bs.magic       = CAIIC_BOOT_MAGIC;
    bs.active_bank = (s_target_base == CAIIC_APP1_BASE) ? CAIIC_ACTIVE_BANK1
                                                        : CAIIC_ACTIVE_BANK2;
    bank = (s_target_base == CAIIC_APP1_BASE) ? &bs.bank1 : &bs.bank2;
    bank->size    = s_total_size;
    bank->crc32   = s_image_crc32;
    bank->version = s_version;

    bs.crc32 = caiic_crc32((const uint8_t*)&bs, (uint32_t)offsetof(caiic_bootsetting_t, crc32));

    if (Qflash_Erase_Sector(CAIIC_BOOTSETTING_ADDR) != 0)
    {
        return BLE_OTA_ST_FLASH_FAIL;
    }
    /* sizeof(caiic_bootsetting_t) = 36, already 4-byte aligned */
    if (Qflash_Write(CAIIC_BOOTSETTING_ADDR, (uint8_t*)&bs, sizeof(bs)) != 0)
    {
        return BLE_OTA_ST_FLASH_FAIL;
    }
    if (memcmp(&bs, (const void*)CAIIC_BOOTSETTING_ADDR, sizeof(bs)) != 0)
    {
        return BLE_OTA_ST_FLASH_FAIL;
    }
    return 0;
}

/**
 * @brief OTA_END: {crc32 u32}. Flush the tail, verify, switch bank, reset.
 */
static void ota_on_end(uint8_t seq, const uint8_t* p, uint16_t len)
{
    uint32_t crc_end, crc_calc;

    if (!s_active)
    {
        ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_BAD_STATE, 0);
        return;
    }
    if (len != 4u)
    {
        ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_BAD_FRAME, s_offset);
        return;
    }
    crc_end = ota_rd32(p);

    if (s_offset != s_total_size)
    {
        ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_BAD_STATE, s_offset);
        return;
    }

    /* Flush the tail: erase the partial sector, pad to a 4-byte multiple */
    if (s_sector_fill != 0)
    {
        uint32_t addr = s_target_base + s_sector_idx * OTA_SECTOR_BUF_SIZE;
        uint32_t padded = (s_sector_fill + 3u) & ~3u;

        while (s_sector_fill < padded)
        {
            s_sector_buf[s_sector_fill++] = 0xFFu;
        }
        if (Qflash_Erase_Sector(addr) != 0 ||
            Qflash_Write(addr, s_sector_buf, padded) != 0)
        {
            ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_FLASH_FAIL, s_offset);
            app_ota_abort();
            return;
        }
    }

    /* Whole-image CRC32 read back from flash */
    crc_calc = caiic_crc32((const uint8_t*)s_target_base, s_total_size);
    if (crc_calc != crc_end || crc_calc != s_image_crc32)
    {
        ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_CRC_FAIL, s_offset);
        app_ota_abort();
        return;
    }

    if (ota_update_bootsetting() != 0)
    {
        ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_FLASH_FAIL, s_offset);
        app_ota_abort();
        return;
    }

    ota_respond(seq, BLE_FRAME_OTA_END, BLE_OTA_ST_OK, s_total_size);

    /* Give the notify time to go out, then reset into the new bank */
    vTaskDelay(pdMS_TO_TICKS(OTA_RESET_DELAY_MS));
    NVIC_SystemReset();
}

void app_ota_handle_frame(uint8_t type, uint8_t seq, const uint8_t* payload, uint16_t len)
{
    switch (type)
    {
        case BLE_FRAME_OTA_BEGIN:
            ota_on_begin(seq, payload, len);
            break;
        case BLE_FRAME_OTA_DATA:
            ota_on_data(seq, payload, len);
            break;
        case BLE_FRAME_OTA_END:
            ota_on_end(seq, payload, len);
            break;
        default:
            break;
    }
}