#include "iap_upgrade.h"
#include "Callback.h"
#include "cmsis_os.h"	//osDelay
#include "uType.h"

//BootParams g_bootParams;
JsonDat_Root JsRoot;

static uint32_t upgrade_offset = 0;
uint16_t upgrade_mode_subcode = 0; // 记录使用的subcode（0x0100 或 0x0200）
uint32_t expected_file_size = 0;	// 升级文件大小值
static uint8_t expected_md5_0200[4];	// 接收的MD5码头4字节
static uint8_t expected_md5_0100[16];	// 接收的MD5码16字节
static uint32_t current_packet_index = 0;
static bool upgrade_in_progress = false; // 确保先收到起始帧再允许数据帧执行
bool erase_APP2 = false;
bool erase_APPdata_buffer = false;
bool save_JsRoot_0100_flag = false;

uint32_t data4001_last_data_time = 0;       // 上次收到数据帧的时间戳（单位：毫秒）

//static uint32_t upgrade_target_addr = IAPDataBuffer_ADDR; // 固定写入缓存区static uint32_t upgrade_target_addr = 0;

// 应答起始帧 状态码+固定6字节（00 00 01 00 00 00）
// （状态码：0=成功，其他为错误）
void send_upgrade_start_response(uint8_t func_code, uint16_t reg_addr, uint8_t status_code, uint16_t sliceCount)
{
    static uint8_t tx_buf[32] __attribute__((aligned(4)));
    uint16_t len = 0;
	
    tx_buf[len++] = JsRoot.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++] = sliceCount >> 8;
    tx_buf[len++] = sliceCount & 0xFF;
		tx_buf[len++] = 0x01;	
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
	
    uint16_t crc = modbus_crc16(tx_buf, len);
    tx_buf[len++] = crc & 0xFF;
    tx_buf[len++] = (crc >> 8) & 0xFF;

		
//	DEBUG_PRINTF("TX_BUF: ");
//    for (uint16_t i = 0; i < len; i++) {
//        DEBUG_PRINTF("%02X ", tx_buf[i]);
//    }
//    DEBUG_PRINTF("\r\n");
		
//    HAL_GPIO_WritePin(GPIOB, RS485_RE_Pin, GPIO_PIN_SET);
//    HAL_UART_Transmit_DMA(&huart1, tx_buf, len);
		uart485send(tx_buf, len);
		if(Systemmode != releasemode)
		taskprinttx(tx_buf,len);
}

// 应答数据帧  状态码 + 固定2字节（00 00）
// （状态码：0=成功，其他为错误）
void send_upgrade_data_response(uint8_t func_code, uint16_t reg_addr, uint8_t status_code, uint16_t sliceCount)
{
    static uint8_t tx_buf[32] __attribute__((aligned(4)));
    uint16_t len = 0;
	
    tx_buf[len++] = JsRoot.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++] = sliceCount >> 8;
    tx_buf[len++] = sliceCount & 0xFF;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;

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

		
//		DEBUG_PRINTF("TX_BUF: ");
//    for (uint16_t i = 0; i < len; i++) {
//        DEBUG_PRINTF("%02X ", tx_buf[i]);
//    }
//    DEBUG_PRINTF("\r\n");
		
//    HAL_GPIO_WritePin(GPIOB, RS485_RE_Pin, GPIO_PIN_SET);
//    HAL_UART_Transmit_DMA(&huart1, tx_buf, len);
		uart485send(tx_buf, len);
		if(Systemmode != releasemode)
		taskprinttx(tx_buf,len);
}

// 应答结束帧  状态码 + 固定2字节（00 00）
// （状态码：0=成功，其他为错误）
void send_upgrade_end_response(uint8_t func_code, uint16_t reg_addr, uint8_t status_code, uint16_t sliceCount)
{
    static uint8_t tx_buf[32] __attribute__((aligned(4)));
    uint16_t len = 0;
	
    tx_buf[len++] = JsRoot.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++] = sliceCount >> 8;
    tx_buf[len++] = sliceCount & 0xFF;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;
		tx_buf[len++] = 0x00;

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

		
//		DEBUG_PRINTF("TX_BUF: ");
//    for (uint16_t i = 0; i < len; i++) {
//        DEBUG_PRINTF("%02X ", tx_buf[i]);
//    }
//    DEBUG_PRINTF("\r\n");
		
//    HAL_GPIO_WritePin(GPIOB, RS485_RE_Pin, GPIO_PIN_SET);
//    HAL_UART_Transmit_DMA(&huart1, tx_buf, len);
		uart485send(tx_buf, len);
		if(Systemmode != releasemode)
		taskprinttx(tx_buf,len);
}


/*
状态码值：
		00：正常
		01：busy
		02：升级文件大小超限
		03：单帧切片大小超限
		04：切片索引错误
		05：校验错误
		06：参数错误
		07：命令码错误
		08：执行错误（如无起始帧却先收到了数据帧）
*/

// 处理起始帧 (0x4000)
void handle_upgrade_start_frame(uint8_t *data)
{	
		DEBUG_PRINTF("\r\n============================\r\n");
		DEBUG_PRINTF("=== Upgrade start ===\r\n");
		DEBUG_PRINTF("0x4000 start\r\n");
		
    uint16_t subcode = (data[6] << 8 | data[7]);
		upgrade_mode_subcode = subcode;  // 记录本次升级的 subcode
	
		data4001_last_data_time = HAL_GetTick();  // 收到起始帧数据更新时间戳
	
		if (upgrade_mode_subcode == 0x0100) { // md5校验码16个字节签名在bin文件最后
				DEBUG_PRINTF("subcode: 0x0100\r\n");
        expected_file_size = (data[23] << 24) | (data[24] << 16) | (data[25] << 8) | data[26];
				DEBUG_PRINTF("expected_file_size[0x%04x]\r\n",expected_file_size);
			
        if (expected_file_size > APP_SIZE) {
            send_upgrade_start_response(0x42, 0xAABB, 0x02,0x0000); // 文件过大
            return;
        }
				erase_APPdata_buffer = true; // 擦除APP2缓存区

        upgrade_offset = 0;
        current_packet_index = 0;
				
        upgrade_in_progress = true; // 证明收到起始帧，数据帧允许接收
				
				save_JsRoot_0100_flag = true; // 写入升级subcode以及文件大小
				
				
				send_upgrade_start_response(0x42, 0xAABB, 0x00,0x0000); // 特殊起始帧应答
				
				Systemmode = IAPbootloader;
				DEBUG_PRINTF("0x4000 end\r\n");
				DEBUG_PRINTF("============================\r\n");
    } else if (upgrade_mode_subcode == 0x0200) { // 无签名，md5校验码直接在起始帧负载里
			
				DEBUG_PRINTF("subcode: 0x0200\r\n");
        expected_file_size = (data[8] << 24) | (data[9] << 16) | (data[10] << 8) | data[11];
				memcpy(expected_md5_0200, &data[12], 4); // 支持MD5的头4字节
					
				save_JsRoot_0100_flag = true; // 写入升级subcode以及文件大小
			
				DEBUG_PRINTF("expected_file_size[0x%04x]\r\n",expected_file_size); 
				DEBUG_PRINTF("expected_md5_0200:");
				for (int i = 0; i < 4; i++)
				{
						DEBUG_PRINTF(" %02X", expected_md5_0200[i]);
						if ((i + 1) % 4 == 0) 
						{
								DEBUG_PRINTF("\n");
						}
				}
				DEBUG_PRINTF("\n");
				DEBUG_PRINTF("============================\r\n");
				
				
        send_upgrade_start_response(0x42, 0xAABB, 0x00,0x0000); // 特殊起始帧应答
				Flash_ErasePages(IAPDataBuffer_ADDR, APP_SIZE / 4); // 擦除APP缓存区数据等待写入
				Systemmode = IAPbootloader;
				
    }else{
				send_upgrade_start_response(0x42, 0xAABB, 0x06, 0x0000);	// 参数错误
				
		}
}

// 处理数据帧 (0x4001)
void handle_upgrade_data_frame(uint8_t *data)
{
		
    if (!upgrade_in_progress) {
        send_upgrade_data_response(0x42, 0xAABB, 0x08, 0x0000); // 执行错误08（无起始帧却先收到数据帧）
        return;
    }
    if ((data[4] << 8 | data[5]) != 0x4001) return;
    
		data4001_last_data_time = HAL_GetTick();  // 每次收到数据更新时间戳
    uint16_t index = data[6] << 8 | data[7];
    if (index != current_packet_index) {
        send_upgrade_data_response(0x42, 0xAABB, 0x04, 0x0000 );
        reset_upgrade_state();
        return;
    }
    Flash_Write64(IAPDataBuffer_ADDR + upgrade_offset, (uint64_t*)&data[8], PACKET_SIZE / 8);
    upgrade_offset += PACKET_SIZE;
    current_packet_index++;
    send_upgrade_data_response(0x42, 0xAABB, 0x00, current_packet_index);
}

// 处理结束帧 (0x4002)
void handle_upgrade_end_frame(void)
{
    uint8_t result_md5[16];
    MD5_CTX context;
	
		send_upgrade_end_response(0x42, 0xAABB, 0x00, 0x0000); // 应答表示收到结束帧
		
		bool md5_pass = false;
		
		if(upgrade_mode_subcode == 0x0100)
		{
				MD5Init(&context);
				MD5Update(&context, (uint8_t*)IAPDataBuffer_ADDR, expected_file_size - 16); // 0100 -16
				MD5Final(result_md5, &context); // 计算收到的整体数据的MD5
			
				// 从 IAPDataBuffer 最后 16 字节读取 MD5 签名
				Flash_ReadBytes(IAPDataBuffer_ADDR + expected_file_size - 16, expected_md5_0100, 16); // 从 IAPDataBuffer 最后16字节读取 MD5 校验码
				
				DEBUG_PRINTF("result_md5:"); 
				for (int i = 0; i < 16; i++)
				{
						DEBUG_PRINTF(" %02X", result_md5[i]);
						if ((i + 1) % 16 == 0) // 
						{
								DEBUG_PRINTF("\n");
						}
				}
				DEBUG_PRINTF("\n");
				
				DEBUG_PRINTF("expected_md5_0100:"); // 收到MD5码16字节
				for (uint16_t i = 0; i < 16; i++)
				{
						DEBUG_PRINTF(" %02X", expected_md5_0100[i]);
						if ((i + 1) % 16 == 0) // 每行打印16字节
						{
								DEBUG_PRINTF("\n");
						}
				}
				DEBUG_PRINTF("\n");
				
				md5_pass = (memcmp(result_md5, expected_md5_0100, 16) == 0); 
				// 比较APP数据缓存区的自计算的MD5和收到的MD5码是否一致，一致则表示APP缓存区接收和写入的没问题。
		} else if(upgrade_mode_subcode == 0x0200){

				MD5Init(&context);
				MD5Update(&context, (uint8_t*)IAPDataBuffer_ADDR, expected_file_size); // 0200 -0
				MD5Final(result_md5, &context); // 计算APP缓存区的MD5
			
				DEBUG_PRINTF("expected_md5_0200: "); // 收到的MD5码
        for (uint16_t i = 0; i < 4; i++) {
            DEBUG_PRINTF(" %02X", expected_md5_0200[i]);
        }
        DEBUG_PRINTF("\n");
				DEBUG_PRINTF("result_md5 (first 4 bytes): "); // 计算的MD5码
        for (uint16_t i = 0; i < 4; i++) {
            DEBUG_PRINTF(" %02X", result_md5[i]);
        }
        DEBUG_PRINTF("\n");
				md5_pass = (memcmp(result_md5, expected_md5_0200, 4) == 0); // 因为收到的只有4字节MD5码，所以只验证前四字节对不对
		}
		
    if (md5_pass) //如果验证APP数据缓存区的MD5码通过
		{
			DEBUG_PRINTF(" MD5 check passed \r\n");       

			JsRoot.iapLoadStatus = 300;					
			DB_SaveUInt(DB_ID_SNAPSHOT, KV_KEY_IAP_LOAD , JsRoot.iapLoadStatus ); // 存入到SNAPSHOT数据库
			
			// 将每个字节转换为两位16进制字符
			for(int i = 0; i < 4; i++) {
					sprintf(&JsRoot.iapMd5[i*2], "%02X", result_md5[i]);
			}
			JsRoot.iapMd5[8] = '\0';
			DB_SaveChar(DB_ID_SNAPSHOT, KV_KEY_IAP_MD5 , JsRoot.iapMd5 ); // 存入MD5到SNAPSHOT数据库
			
			osDelay(500);
			NVIC_SystemReset();
			
    } else {
			// 收到的APP数据缓存区的数据不完整，就擦除掉缓存区的数据
				DEBUG_PRINTF (" MD5 check fails, Erase IAPDataBuffer_ADDR Flash \r\n");
				Flash_ErasePages(IAPDataBuffer_ADDR, APP_SIZE / 4);  // 擦除整个 IAPDataBuffer 区域（80KB）擦除 40 页（2048字节/页）
				send_upgrade_end_response(0x42, 0xAABB, 0x05, 0x0000); //MD5校验错误
    }

    reset_upgrade_state(); // 清除相关升级状态标志
// 4.09.06
}
// 清除升级状态（在失败或完成后调用）
void reset_upgrade_state(void)
{
    upgrade_offset = 0;
    expected_file_size = 0;
    memset(expected_md5_0100, 0, sizeof(expected_md5_0100));
		memset(expected_md5_0200, 0, sizeof(expected_md5_0200));
    current_packet_index = 0;
    upgrade_in_progress = false;
}
void check_upgrade_data4001_timeout(void)
{
    if (upgrade_in_progress)
    {
        if (HAL_GetTick() - data4001_last_data_time > UPGRADE_TIMEOUT_MS)
        {
            DEBUG_PRINTF("Upgrade timeout\r\n");
            reset_upgrade_state();       // 清除升级状态
        }
    }
}

void handle_rollback_request(void) // 回退
{
    JsRoot.iapLoadStatus = 0;
    JsRoot.ubootback = 1;
		DB_SaveUInt(DB_ID_SNAPSHOT, KV_KEY_IAP_LOAD , JsRoot.iapLoadStatus ); // 存入到SNAPSHOT数据库
		DB_SaveUInt(DB_ID_SNAPSHOT, KV_KEY_UBOOTBACK , JsRoot.ubootback ); // 存入到SNAPSHOT数据库
		osDelay(300);
    NVIC_SystemReset();
}
void app_jump_check_and_jump(void) // 初始化完成之后，检查状态是否跳转 
{

	if (__HAL_RCC_GET_FLAG(RCC_FLAG_IWDGRST)){ // 是否有看门狗导致复位
			printf ("\r\n IWDG_RST flag");
			// 增加计数
		}else{
			printf ("\r\n No IWDG_RST flag");
		}

    if (JsRoot.iapLoadStatus == 2) {
				printf ("\r\n valid APP2. Run APP2\r\n");
				if (!JumpToApp(APP2_ADDR)) {

						printf("\r\n Jump failed. Run APP1" );
        } 
					
    } else if(JsRoot.iapLoadStatus == 0){
      // 运行当前 APP1 程序
			printf ("\r\n No valid APP2. Run APP1\r\n");
    }
}

bool JumpToApp(uint32_t addr) 
{
		// APP1: sp = 0x100037A8 ; pc = 0x08014281;
		// APP2: sp = 0x100037A8 ; pc = 0x08000281;
    uint32_t sp = *(volatile uint32_t*)addr;// *(volatile uint32_t*)addr;// 0x100037A8;// 
    uint32_t pc = *(volatile uint32_t*)(addr + 4);// *(volatile uint32_t*)(addr + 4);// 0x08014281;// 
				
		if( sp ==0x100037A8 && pc == 0x08014281) // 跳转到APP2区域
		{
			DEBUG_PRINTF("\r\n Jumping app_address. SP=0x%08X, PC=0x%08X \r\n", sp, pc);
			DEBUG_PRINTF("\r\n Jumping app_address success");
			HAL_RCC_DeInit();		// 复位时钟
			HAL_DeInit();				// 复位外设
			SCB->VTOR = addr;	  // 设置向量表基地址		
			__set_MSP(sp);		  // 设置主堆栈指针
			__disable_irq();    // 关闭总中断
			void (*jump)(void) = (void (*)(void))pc; 
			jump(); // 跳转到APP
			return true;
		}else{
			printf("\r\n Jumping app_address. SP=0x%08X, PC=0x%08X \r\n", sp, pc);
			printf("\r\n Jumping app_address fail");
			return false;
		}
    
}