#include "bsp_iap.h"

#include <stdio.h>
#include <string.h>

#include "app.h"

#include "bsp_uart.h"
#include "usr_config.h"


#include "os_timer.h"

/* IAP 状态定义 */
#define BSP_IAP_STATE_NOMORAL      (0)   /* 正常模式，等待跳转 */
#define BSP_IAP_STATE_UPDATEING    (1)   /* 正在升级 */
#define BSP_IAP_STATE_SUCCEED      (2)   /* 升级完成，准备跳转 */

/* IAP 协议接收帧头 */
#define BSP_IAP_PROTO_RX_HEADER             (0x5A)
/* IAP 协议接收命令码 */
#define BSP_IAP_PROTO_RX_CMD_GET_SV         (0x01)  /* 读取版本号 */
#define BSP_IAP_PROTO_RX_CMD_CODE_SIZE      (0x02)  /* 获取代码大小并擦除扇区 */
#define BSP_IAP_PROTO_RX_CMD_WRITE_DATA     (0x03)  /* 写入数据 */
#define BSP_IAP_PROTO_RX_CMD_UNKNOW_1       (0x04)  /* 未知命令1 */
#define BSP_IAP_PROTO_RX_CMD_UNKNOW_2       (0x05)  /* 未知命令2 */
#define BSP_IAP_PROTO_RX_CMD_UNKNOW_3       (0x06)  /* 未知命令3 */

/* IAP 协议发送帧头 */
#define BSP_IAP_PROTO_TX_HEADER             (0xA5)

/* IAP 协议错误码 */
#define BSP_IAP_PROTO_ERROR_CODE_SUCCEED    (0)     /* 解析成功 */
#define BSP_IAP_PROTO_ERROR_CODE_CHECK      (1)     /* 校验错误 */
#define BSP_IAP_PROTO_ERROR_CODE_CMD        (2)     /* 命令错误 */
#define BSP_IAP_PROTO_ERROR_CODE_LENGTH     (3)     /* 长度错误 */

/* 发送缓冲区大小 */
#define BSP_IAP_TX_LEN                      (32)



/* 静态变量 */
static u16 bsp_app_time_start;      /* 等待跳转的起始时间（毫秒） */
static u8 bsp_iap_tx_buf[BSP_IAP_TX_LEN];   /* 发送缓冲区 */

/* 函数声明 */
static void app_jump(void);
static void flash_page_erase(u16 page_num);
static void flash_write_u32(u32 write_addr, u32 *p_buffer, u16 num_write);
static void flash_write_u8(u32 write_addr, u8 *p_buffer, u16 num_write);
static void bsp_iap_send(u8 *p_data, u16 len);
static void bsp_iap_rx_task(u8 *p_data, u16 len, void *other_data);
static void bsp_iap_init(void);
static void bsp_iap_task(void);
static u8 sum_check(u8 *p_data, u16 len);

/******************************************
 * 结构体: bsp_iap
 * 功能: IAP 控制实例
 * 描述: 定义 IAP 的具体配置和回调函数
*******************************************/
bsp_iap_t iap =
{
    .init = bsp_iap_init,
    .task = bsp_iap_task,
};

/* 全局指针，指向 IAP 结构体 */
bsp_iap_t *p_iap = &iap;

bsp_Uart_t *p_rx_uart;

/* 应用程序起始地址 */
#define APP_START_ADDR      0x08002800
/* 扇区大小（字节） */
#define SECTOR_SIZE         FLASH_PAGE_SIZE

/******************************************
 * 函数: app_jump
 * 功能: 跳转到应用程序
 * 参数: 无
 * 返回: 无
 * 描述: 关闭所有外设和中断，设置堆栈指针，跳转到用户程序
*******************************************/
static void app_jump(void)
{
    u32 app_addr = APP_START_ADDR;

    /* 检查应用程序是否存在（堆栈指针有效） */
    if (((*(u32*)app_addr) & 0x2FFE0000) == 0x20000000)
    {
        typedef void (*iapfun)(void);
        iapfun jump_to_app;
        /* 关闭所有中断 */
        __disable_irq();

        /* 设置所有时钟到默认状态，使用HSI时钟 */
        HAL_RCC_DeInit();
        
//        LL_USART_Disable(USART1);
//        LL_USART_Disable(USART2);
//        
//        LL_USART_DisableIT_RXNE_RXFNE(USART1);
//        LL_USART_DisableIT_RXNE_RXFNE(USART2);
        
        
        SysTick->CTRL = 0;
        SysTick->LOAD = 0;
        SysTick->VAL = 0;

        /* 关闭所有中断，清除所有中断挂起标志 */
        for (u8 i = 0; i < 8; i++)
        {
            NVIC->ICER[i]=0xFFFFFFFF;
            NVIC->ICPR[i]=0xFFFFFFFF;
        }
        

    
        /* 跳转到应用程序 */
        jump_to_app = (iapfun)*(u32*)(app_addr + 4);   /* 复位向量地址 */
        __set_MSP(*(u32*)app_addr);                    /* 设置主堆栈指针 */
        __enable_irq();
        jump_to_app();                                 /* 跳转 */
        while(1);
    }
}

static void bsp_Flash_FLASH_ErasePage(uint32_t Address)
{
    FLASH_EraseInitTypeDef EraseInitStruct;
    uint32_t PageError = 0;
    
    EraseInitStruct.TypeErase = FLASH_TYPEERASE_PAGES;
    EraseInitStruct.PageAddress = Address;
    EraseInitStruct.NbPages = 1;
    
    HAL_FLASHEx_Erase(&EraseInitStruct, &PageError);
}

/******************************************
 * 函数: flash_page_erase
 * 功能: 擦除指定数量的 Flash 扇区
 * 参数: page_num - 要擦除的扇区数量
 * 返回: 无
 * 描述: 从 APP_START_ADDR 开始擦除连续 page_num 个扇区
*******************************************/
static void flash_page_erase(u16 page_num)
{
    u16 i;

    HAL_FLASH_Unlock();//解锁
    __HAL_FLASH_CLEAR_FLAG(FLASH_FLAG_EOP  | FLASH_FLAG_WRPERR );
    for (i = 0; i < page_num; i++)
    {
        bsp_Flash_FLASH_ErasePage(i * SECTOR_SIZE+APP_START_ADDR);
    }

    HAL_FLASH_Lock();//上锁
}


#include <string.h>   // 使用 memcpy

/**
 * @brief 将数据缓冲区写入 Flash（半字编程）
 * @param write_addr   Flash 目标地址（必须是 8 字节对齐）
 * @param code_buffer  源数据缓冲区（可以为任意对齐）
 * @param size         要写入的字节数
 */
static void flash_code_write(uint32_t write_addr, uint8_t *code_buffer, uint16_t size)
{
    HAL_FLASH_Unlock();
    __HAL_FLASH_CLEAR_FLAG(FLASH_FLAG_EOP | FLASH_FLAG_WRPERR);

    // 按半字（2字节）写入
    for (uint16_t i = 0; i < size; i += 2)
    {
        uint16_t data;
        if (i + 1 < size)
            data = ((uint16_t)code_buffer[i+1] << 8) | code_buffer[i];
        else
            data = (uint16_t)code_buffer[i];   // 最后一个字节，高8位补0

        HAL_FLASH_Program(FLASH_TYPEPROGRAM_HALFWORD, write_addr + i, data);
    }

    HAL_FLASH_Lock();
}

/******************************************
 * 函数: bsp_iap_send
 * 功能: 通过 UART1 发送数据
 * 参数: p_data - 数据缓冲区指针
 *       len    - 数据长度
 * 返回: 无
 * 描述: 调用 UART1 的发送接口
*******************************************/
static void bsp_iap_send(u8 *p_data, u16 len)
{
    if(p_rx_uart != NULL)
    {
        p_rx_uart->Send(p_rx_uart, p_data, len);
    }
}

/******************************************
 * 函数: bsp_iap_init
 * 功能: IAP 初始化
 * 参数: 无
 * 返回: 无
 * 描述: 设置初始状态，并注册接收处理函数
*******************************************/
static void bsp_iap_init(void)
{
    p_iap->state = BSP_IAP_STATE_NOMORAL;
    //COM_Uart1.Rx_DataAnalysis = bsp_iap_rx_task;
    COM_Uart2.Rx_DataAnalysis = bsp_iap_rx_task;
}

/******************************************
 * 函数: bsp_iap_task
 * 功能: IAP 任务（主循环中调用）
 * 参数: 无
 * 返回: 无
 * 描述: 根据当前状态执行相应操作
*******************************************/
static void bsp_iap_task(void)
{
    switch (p_iap->state)
    {
        case BSP_IAP_STATE_NOMORAL:         /* 正常模式，2 秒后跳转 APP */
        {
            if (TIME_TRUE == OsTimer_CheckTimeOut(bsp_app_time_start, osTime_MSecTick, 1000))
            {
                app_jump();
            }
        } break;

        case BSP_IAP_STATE_UPDATEING:       /* 升级中 */
        {
            /* 预留，可能处理超时等 */
        } break;

        case BSP_IAP_STATE_SUCCEED:         /* 升级完成，立即跳转 */
        {
            app_jump();
        } break;

        default:
            break;
    }
}

/******************************************
 * 函数: sum_check
 * 功能: 计算校验和（累加和取反）
 * 参数: p_data - 数据指针
 *       len    - 数据长度
 * 返回: 校验和字节
 * 描述: 对数据进行累加求和，然后取反（补码）
*******************************************/
static u8 sum_check(u8 *p_data, u16 len)
{
    u16 i;
    u8 sum = 0;

    for (i = 0; i < len; i++)
    {
        sum += p_data[i];
    }
    return (-sum);
}

/******************************************
 * 函数: bsp_iap_rx_task
 * 功能: IAP 协议接收处理函数
 * 参数: p_data      - 接收数据缓冲区
 *       len         - 数据长度
 *       other_data  - 附加数据（未使用）
 * 返回: 无
 * 描述: 解析 IAP 协议帧，执行相应命令并回复
*******************************************/
static void bsp_iap_rx_task(u8 *p_data, u16 len, void *other_data)
{
    u16 i, cmd, data_len, tx_len = 0;
    u8 check_data, error_code = BSP_IAP_PROTO_ERROR_CODE_SUCCEED;
    u8 *p_data_offset;

    if(BSP_IAP_STATE_SUCCEED == p_iap->state)/*升级成功后，不再进行检测*/
    {
        return;
    }
    /* 帧头检测和长度检查 */
    if (p_data[0] != BSP_IAP_PROTO_RX_HEADER || len < 4)
    {
        return;
    }

    data_len = p_data[1] << 8 | p_data[2];
    if (len - 4 != data_len)                    /* 长度错误 */
    {
        error_code = BSP_IAP_PROTO_ERROR_CODE_LENGTH;
    }
    else
    {
        check_data = sum_check(p_data, len - 1);
        if (check_data != p_data[len - 1])       /* 校验错误 */
        {
            error_code = BSP_IAP_PROTO_ERROR_CODE_CHECK;
        }
    }
    
    p_rx_uart = (bsp_Uart_t*)(other_data);
    cmd = p_data[3];
    p_data_offset = &p_data[4];

    /* 根据命令处理 */
    switch (cmd)
    {
        case BSP_IAP_PROTO_RX_CMD_GET_SV:           /* 获取软件版本号 */
        {
            u8 str_len, tx_data_len;

            p_iap->state = BSP_IAP_STATE_UPDATEING;   /* 进入升级状态 */

            str_len = strlen(SwVersion) + 1;
            memcpy(&bsp_iap_tx_buf[3], SwVersion, str_len);
            tx_data_len = str_len;

            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        case BSP_IAP_PROTO_RX_CMD_CODE_SIZE:        /* 获取代码大小并擦除扇区 */
        {
            u8 tx_data_len;
            u16 erase_size, erase_page, page_size;

            erase_page = p_data_offset[0] << 8 | p_data_offset[1];
            page_size = p_data_offset[2] << 8 | p_data_offset[3];
            p_iap->page_size = page_size;

            /* 计算需要擦除的扇区数 */
            erase_page /= ((float)SECTOR_SIZE / (float)page_size);
            flash_page_erase(erase_page);

            tx_data_len = 0;

            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        case BSP_IAP_PROTO_RX_CMD_WRITE_DATA:       /* 写入数据 */
        {
            u8 tx_data_len;
            u16 index, code_size, addr_offset, i;
            u32 addr; 
            u8 *p_code_data;

            index = p_data_offset[0] << 8 | p_data_offset[1];
            code_size = data_len - 3;

            addr = index * p_iap->page_size;
            p_code_data = &p_data_offset[2];

            
            flash_code_write(APP_START_ADDR+addr, p_code_data, code_size);
            

            /* 构造应答帧，返回当前索引 */
            bsp_iap_tx_buf[3] = index >> 8;
            bsp_iap_tx_buf[4] = index & 0x00ff;
            tx_data_len = 2;

            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        case BSP_IAP_PROTO_RX_CMD_UNKNOW_1:        /* 未知命令1（升级完成） */
        {
            u8 tx_data_len;
            tx_data_len = 0;
            p_iap->state = BSP_IAP_STATE_SUCCEED;   /* 升级完成 */
            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        case BSP_IAP_PROTO_RX_CMD_UNKNOW_3:        /* 未知命令3（回传两个字节） */
        {
            u8 tx_data_len;
            tx_data_len = 2;
            bsp_iap_tx_buf[3] = p_data_offset[0];
            bsp_iap_tx_buf[4] = p_data_offset[1];
            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        case BSP_IAP_PROTO_RX_CMD_UNKNOW_2:        /* 未知命令2（简单应答） */
        {
            u8 tx_data_len;
            tx_data_len = 0;
            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = tx_data_len + 1;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[tx_data_len + 3] = sum_check(bsp_iap_tx_buf, bsp_iap_tx_buf[1] + 2);
            tx_len = bsp_iap_tx_buf[1] + 3;
        } break;

        default:
        {
            error_code = BSP_IAP_PROTO_ERROR_CODE_CMD;
        } break;
    }

    /* 错误处理 */
    switch (error_code)
    {
        case BSP_IAP_PROTO_ERROR_CODE_SUCCEED:
        {
            bsp_iap_send(bsp_iap_tx_buf, tx_len);
        } break;

        case BSP_IAP_PROTO_ERROR_CODE_CHECK:
        case BSP_IAP_PROTO_ERROR_CODE_CMD:
        case BSP_IAP_PROTO_ERROR_CODE_LENGTH:
        {
            bsp_iap_tx_buf[0] = BSP_IAP_PROTO_TX_HEADER;
            bsp_iap_tx_buf[1] = 0x01;
            bsp_iap_tx_buf[2] = cmd;
            bsp_iap_tx_buf[3] = error_code;
            bsp_iap_tx_buf[4] = sum_check(bsp_iap_tx_buf, 4);
            bsp_iap_send(bsp_iap_tx_buf, 5);
        } break;

        default:
            break;
    }
}