#include "iap_upgrade.h"


static uint32_t upgrade_offset = 0;
static uint32_t expected_file_size = 0;	// 升级文件大小值
static uint8_t expected_md5[16];	// 接收的MD5码头4字节
static uint32_t current_packet_index = 0;
static bool upgrade_in_progress = false;
static uint32_t upgrade_target_addr = 0;

// 发送应答帧（状态码：0=成功，其他为错误）
void send_upgrade_response(uint8_t func_code, uint16_t reg_addr, uint8_t status_code, uint16_t sliceCount)
{
    uint8_t tx_buf[32];
    uint16_t len = 0;
	
    tx_buf[len++] = g_deviceParams.SlaveDevice_addr;
    tx_buf[len++] = func_code;
    tx_buf[len++] = reg_addr >> 8;
    tx_buf[len++] = reg_addr & 0xFF;
    tx_buf[len++] = 0x41; // 自定义应答命令
    tx_buf[len++] = 0x00;
    tx_buf[len++] = status_code; // 状态码
    tx_buf[len++] = 0x00;
    tx_buf[len++] = sliceCount >> 8;
    tx_buf[len++] = sliceCount & 0xFF;

    uint16_t crc = modbus_crc16(tx_buf, len);
    tx_buf[len++] = crc & 0xFF;
    tx_buf[len++] = (crc >> 8) & 0xFF;

    HAL_GPIO_WritePin(GPIOB, RS485_RE_Pin, GPIO_PIN_SET);
    HAL_UART_Transmit_DMA(&huart1, tx_buf, len);
}

// 默认写入到 APP2 区，如果当前运行的是 APP2，则写 APP1 区
uint32_t get_backup_app_addr(void)
{
    return (APP1_ADDR == BOOTLOADER_ADDR) ? APP2_ADDR : APP1_ADDR;
}

// 清除升级状态（在失败或完成后调用）
void reset_upgrade_state(void)
{
    upgrade_offset = 0;
    expected_file_size = 0;
    memset(expected_md5, 0, sizeof(expected_md5));
    current_packet_index = 0;
    upgrade_in_progress = false;
}
/*
状态码值：
		00：正常
		01：busy
		02：升级文件大小超限
		03：单帧切片大小超限
		04：切片索引错误
		05：crc校验错误
		06：参数错误
		07：命令码错误
		08：执行错误（如无起始帧却先收到了数据帧）
*/

// 处理起始帧 (0x4000)
void handle_upgrade_start_frame(uint8_t *data)
{
    uint16_t subcode = (data[6] << 8 | data[7]);

    if (subcode == 0x0100) {
        expected_file_size = (data[23] << 24) | (data[24] << 16) | (data[25] << 8) | data[26];
        if (expected_file_size > MAX_UPGRADE_SIZE) {
            send_upgrade_response(0x42, 0xAABB, 0x02,0x0000); // 文件过大
            return;
        }
        upgrade_target_addr = get_backup_app_addr();
        Flash_ErasePages(upgrade_target_addr, expected_file_size / PAGE_SIZE + 1);
        upgrade_offset = 0;
        current_packet_index = 0;
        upgrade_in_progress = true;
        send_upgrade_response(0x42, 0xAABB, 0x00,0x0000);
				
    } else if (subcode == 0x0200) {
        expected_file_size = (data[8] << 24) | (data[9] << 16) | (data[10] << 8) | data[11];
        memcpy(expected_md5, &data[12], 4); // 支持MD5的头4字节
        send_upgrade_response(0x42, 0x4000, 0x00, 0x0000);
			
    }else{
				send_upgrade_response(0x42, 0xAABB, 0x06, 0x0000);	// 参数错误
		}
}

// 处理数据帧 (0x4001)
void handle_upgrade_data_frame(uint8_t *data)
{
    if (!upgrade_in_progress) {
        send_upgrade_response(0x42, 0x4001, 0x08, 0x0000); // 执行错误08（无起始帧却先收到数据帧）
        return;
    }
    if ((data[2] << 8 | data[3]) != 0x4001) return;
    
    uint16_t index = data[6] << 8 | data[7];
    if (index != current_packet_index) {
        send_upgrade_response(0x42, 0x4001, 0x04, 0x0000);
        reset_upgrade_state();
        return;
    }

    Flash_Write64(upgrade_target_addr + upgrade_offset, (uint64_t*)&data[8], PACKET_SIZE / 8);

    upgrade_offset += PACKET_SIZE;
    current_packet_index++;

    send_upgrade_response(0x42, 0x4001, 0x00, 0x0000);
}

// 处理结束帧 (0x4002)
void handle_upgrade_end_frame(void)
{
    uint32_t addr = get_backup_app_addr();
    uint8_t result_md5[16];
    MD5_CTX context;

    MD5Init(&context);
    MD5Update(&context, (uint8_t*)addr, expected_file_size);
    MD5Final(result_md5, &context);

    if (memcmp(result_md5, expected_md5, 4) == 0) {
        BootParams new_params;
        SaveBootloaderParams();
        memcpy(&new_params, &params, sizeof(BootParams));
        if (addr == APP1_ADDR) {
            memcpy(new_params.app1_md5, result_md5, 16);
            new_params.app1_valid = 0xA5;
            new_params.app1_version += 1;
        } else {
            memcpy(new_params.app2_md5, result_md5, 16);
            new_params.app2_valid = 0xA5;
            new_params.app2_version += 1;
        }
        new_params.upgrade_flag = 0xAABBCCDD;
        memcpy(&params, &new_params, sizeof(BootParams));
        SaveBootloaderParams();
        NVIC_SystemReset();
    } else {
        Flash_ErasePages(addr, expected_file_size / PAGE_SIZE + 1);
    }

    reset_upgrade_state();
}