/**
 * @file    bootloader.c
 * @brief   Donghui ORPC OTA Bootloader
 *
 * Startup sequence:
 *   1. System clock init (HSE 72MHz)
 *   2. Check upgrade flag in param sector
 *   3. If pending upgrade:
 *      a. Read encrypted firmware from SPI Flash
 *      b. HMAC verify
 *      c. AES-CBC decrypt
 *      d. Write to inactive bank (App B if current = A, vice versa)
 *      e. Verify CRC, update boot flag
 *   4. Jump to active bank
 *
 * If no upgrade pending, jump immediately to App for fast boot.
 *
 * Vector table relocation:
 *   - Bootloader runs at 0x08000000
 *   - Each App bank has its own vector table at its base
 *   - Before jumping, SCB->VTOR is set to the bank address
 */

#include "stm32g4xx.h"
#include "firmware_map.h"
#include "crypto.h"
#include "spi_flash.h"

/* Include auto-generated keys */
#include "firmware_keys.h"

/* ==================== Flash helpers ==================== */

/* STM32G474 has 4KB sectors starting from Bank 1 */
#define FLASH_KEY1              0x45670123
#define FLASH_KEY2              0xCDEF89AB
#define FLASH_SECTOR_SIZE_BYTES 4096
#define FLASH_OPT_KEY1          0x08192A3B
#define FLASH_OPT_KEY2          0x4C5D6E7F

/* Wait for flash operation to complete */
static void flash_wait(void)
{
    while (FLASH->SR & FLASH_SR_BSY);
}

/* Unlock flash control register */
static void flash_unlock(void)
{
    if (FLASH->CR & FLASH_CR_LOCK) {
        FLASH->KEYR = FLASH_KEY1;
        FLASH->KEYR = FLASH_KEY2;
    }
}

/* Lock flash control register */
static void flash_lock(void)
{
    FLASH->CR |= FLASH_CR_LOCK;
}

/* Erase a single 4KB sector by absolute address */
static bool flash_erase_sector(uint32_t addr)
{
    uint32_t sector_num = (addr - FLASH_BASE_ADDR) / FLASH_SECTOR_SIZE_BYTES;

    flash_unlock();
    flash_wait();

    FLASH->CR = (sector_num << FLASH_CR_SNB_Pos) | FLASH_CR_SER | FLASH_CR_START;
    flash_wait();
    FLASH->CR &= ~(FLASH_CR_SER | FLASH_CR_SNB_Msk);

    flash_lock();
    return true;
}

/* Program 64-bit double word */
static bool flash_program_dword(uint32_t addr, uint64_t data)
{
    flash_unlock();
    flash_wait();

    FLASH->CR |= FLASH_CR_PG;

    *(__IO uint32_t *)addr = (uint32_t)data;
    *(__IO uint32_t *)(addr + 4) = (uint32_t)(data >> 32);

    flash_wait();
    FLASH->CR &= ~FLASH_CR_PG;

    /* Verify */
    uint64_t read = *((__IO uint64_t *)addr);
    flash_lock();

    return (read == data);
}

/* Program flash from buffer (must be 8-byte aligned) */
static bool flash_write_bank(uint32_t dst_addr, const uint8_t *data, uint32_t len)
{
    /* Erase sectors first */
    uint32_t start_sector = (dst_addr - FLASH_BASE_ADDR) / FLASH_SECTOR_SIZE_BYTES;
    uint32_t end_sector = ((dst_addr + len - 1 - FLASH_BASE_ADDR) / FLASH_SECTOR_SIZE_BYTES);

    for (uint32_t s = start_sector; s <= end_sector; s++) {
        uint32_t sector_addr = FLASH_BASE_ADDR + (s * FLASH_SECTOR_SIZE_BYTES);
        if (!flash_erase_sector(sector_addr)) {
            return false;
        }
    }

    /* Program 8 bytes at a time */
    for (uint32_t offset = 0; offset < len; offset += 8) {
        uint64_t dword = 0;
        for (int i = 0; i < 8 && (offset + i) < len; i++) {
            dword |= ((uint64_t)data[offset + i] << (i * 8));
        }
        if (!flash_program_dword(dst_addr + offset, dword)) {
            return false;
        }
    }

    return true;
}

/* ==================== Parameter sector management ==================== */

static void param_read(BootParam_t *param)
{
    const BootParam_t *p = (const BootParam_t *)PARAM_SECTOR_ADDR;
    /* Copy to avoid unaligned reads */
    for (uint32_t i = 0; i < sizeof(BootParam_t); i++) {
        ((uint8_t *)param)[i] = ((const uint8_t *)p)[i];
    }
}

static void param_write(const BootParam_t *param)
{
    flash_erase_sector(PARAM_SECTOR_ADDR);
    flash_write_bank(PARAM_SECTOR_ADDR, (const uint8_t *)param, sizeof(BootParam_t));
}

static bool param_is_valid(const BootParam_t *param)
{
    return (param->magic == 0xA5A5A5A5 && param->tail_magic == 0x5A5A5A5A);
}

static void param_init(BootParam_t *param)
{
    param->magic = 0xA5A5A5A5;
    param->active_bank = 0;    /* Start from App A */
    param->upgrade_status = UPGRADE_STATUS_IDLE;
    param->boot_count = 0;
    param->app_a_version = 0x0100;
    param->app_b_version = 0x0000;
    param->app_a_crc32 = 0;
    param->app_b_crc32 = 0;
    param->ota_file_size = 0;
    param->ota_file_crc32 = 0;
    param->ota_downloaded = 0;
    for (int i = 0; i < 8; i++) param->reserve[i] = 0;
    param->tail_magic = 0x5A5A5A5A;
}

/* ==================== CRC32 (simple) ==================== */

static uint32_t crc32_calc(const uint8_t *data, uint32_t len)
{
    uint32_t crc = 0xFFFFFFFF;
    for (uint32_t i = 0; i < len; i++) {
        crc ^= data[i];
        for (int j = 0; j < 8; j++) {
            if (crc & 1)
                crc = (crc >> 1) ^ 0xEDB88320;
            else
                crc >>= 1;
        }
    }
    return ~crc;
}

/* ==================== Upgrade check and execution ==================== */

static void run_upgrade(const BootParam_t *param)
{
    uint32_t enc_size = param->ota_file_size;
    uint32_t active_bank = param->active_bank;
    uint32_t inactive_bank_addr;
    uint32_t inactive_bank_size = APP_B_SIZE;
    uint32_t fw_plain_size = 0;

    /* Determine which bank to write to (opposite of active) */
    if (active_bank == 0) {
        inactive_bank_addr = APP_B_ADDR;
    } else {
        inactive_bank_addr = APP_A_ADDR;
    }

    /* Allocate buffer on stack (or use large global buffer, beware stack size!) */
    /* For G474 with 128KB RAM, allocate a reasonable working buffer */
    /* Encrypted firmware max size: 240KB + 48 bytes overhead */
    /* We'll read and decrypt in chunks if RAM is limited */
    /* For simplicity: allocate a 32KB working buffer, process in segments */

    uint8_t work_buf[32 * 1024] __attribute__((aligned(8)));
    uint32_t total_processed = 0;

    /* Read and process chunk by chunk */
    /* Note: For production, optimize with larger RAM buffer */
    while (total_processed < enc_size) {
        uint32_t chunk = (enc_size - total_processed > sizeof(work_buf)) ?
                          sizeof(work_buf) : (enc_size - total_processed);

        /* Read encrypted data from SPI Flash */
        FLASH_Read(OTA_BUF_ADDR + total_processed, work_buf, chunk);

        /* For the first chunk (contains IV + ciphertext start + HMAC at end),
         * we need to handle the HMAC separately */
        if (total_processed == 0 && chunk < enc_size) {
            /* Can't verify HMAC until we have entire file */
            /* Just pass data through - HMAC verified on final chunk */
        }

        /* Decrypt in place */
        /* When chunk is entire file (small firmware) or last chunk with HMAC */
        if (total_processed + chunk >= enc_size) {
            /* Last chunk: need to pass from start for HMAC verify */
            break;  /* Handle remaining as single pass below */
        }

        total_processed += chunk;
    }

    /* Simple approach: if firmware fits in RAM, do one-pass decrypt */
    if (enc_size < sizeof(work_buf)) {
        FLASH_Read(OTA_BUF_ADDR, work_buf, enc_size);

        /* Decrypt (includes HMAC verify) */
        uint8_t plain_buf[240 * 1024] __attribute__((aligned(8)));
        bool ok = CRYPTO_OTA_DecryptFirmware(work_buf, enc_size,
                                             AES_KEY, HMAC_KEY,
                                             plain_buf, &fw_plain_size);
        if (!ok) {
            /* Decryption failed - abort upgrade */
            BootParam_t p = *param;
            p.upgrade_status = UPGRADE_STATUS_FAILED;
            param_write(&p);
            return;
        }

        /* Verify CRC */
        uint32_t calc_crc = crc32_calc(plain_buf, fw_plain_size);
        if (calc_crc != param->ota_file_crc32) {
            BootParam_t p = *param;
            p.upgrade_status = UPGRADE_STATUS_FAILED;
            param_write(&p);
            return;
        }

        /* Write to inactive bank */
        if (!flash_write_bank(inactive_bank_addr, plain_buf, fw_plain_size)) {
            BootParam_t p = *param;
            p.upgrade_status = UPGRADE_STATUS_FAILED;
            param_write(&p);
            return;
        }

        /* Update boot parameters */
        BootParam_t new_param = *param;
        new_param.active_bank = (active_bank == 0) ? 1 : 0;
        if (active_bank == 0) {
            new_param.app_b_version = param->ota_file_crc32;
            new_param.app_b_crc32 = calc_crc;
        } else {
            new_param.app_a_version = param->ota_file_crc32;
            new_param.app_a_crc32 = calc_crc;
        }
        new_param.upgrade_status = UPGRADE_STATUS_DONE;

        /* Clear OTA buffer flag */
        new_param.ota_file_size = 0;
        new_param.ota_downloaded = 0;

        param_write(&new_param);
    } else {
        /* Firmware too large for RAM buffer - need segmented processing */
        /* For production, implement AES-CBC streaming decryption */
        BootParam_t p = *param;
        p.upgrade_status = UPGRADE_STATUS_FAILED;
        param_write(&p);
    }
}

/* ==================== Jump to application ==================== */

typedef void (*app_func_t)(void);

static void jump_to_app(uint32_t app_addr)
{
    /* Disable interrupts */
    __disable_irq();

    /* Reset all peripherals (SysTick, NVIC) */
    SysTick->CTRL = 0;
    for (int i = 0; i < 8; i++) {
        NVIC->ICER[i] = 0xFFFFFFFF;
        NVIC->ICPR[i] = 0xFFFFFFFF;
    }

    /* Set vector table to app address */
    SCB->VTOR = app_addr;

    /* Set MSP to first word in vector table (stack pointer) */
    __set_MSP(*(volatile uint32_t *)app_addr);

    /* Jump to reset handler (second word) */
    app_func_t app_entry = (app_func_t)(*(volatile uint32_t *)(app_addr + 4));
    app_entry();
}

/* ==================== Main bootloader entry ==================== */

void bootloader_main(void)
{
    BootParam_t param;

    /* Read boot parameters */
    param_read(&param);

    /* If param sector is invalid (first boot or corrupted), initialize */
    if (!param_is_valid(&param)) {
        param_init(&param);
        param_write(&param);
        param_read(&param);
    }

    /* Increment boot count */
    param.boot_count++;
    param_write(&param);

    /* Check for pending upgrade */
    if (param.upgrade_status == UPGRADE_STATUS_DOWNLOADED ||
        param.upgrade_status == UPGRADE_STATUS_VERIFIED) {

        /* Update status to flashing */
        param.upgrade_status = UPGRADE_STATUS_FLASHING;
        param_write(&param);

        /* Execute upgrade */
        run_upgrade(&param);
    }

    /* Determine which bank to boot */
    uint32_t boot_addr;
    if (param.active_bank == 0) {
        boot_addr = APP_A_ADDR;
    } else {
        boot_addr = APP_B_ADDR;
    }

    /* Validate app vector table before jumping */
    uint32_t msp = *(volatile uint32_t *)boot_addr;
    uint32_t reset_vec = *(volatile uint32_t *)(boot_addr + 4);

    if (msp < 0x20000000 || msp > 0x20020000 ||  /* MSP within SRAM range */
        reset_vec == 0xFFFFFFFF ||
        reset_vec == 0x00000000) {
        /* Invalid vector table - stay in bootloader */
        while (1) {
            /* Blink error LED */
        }
    }

    /* Jump to application */
    jump_to_app(boot_addr);
}