#include "ota_boot.h"
#include "main.h"
#include "../HardWare/W25Q64/w25q64.h"
#include "../HardWare/LOG/log.h"
#include <stdint.h>
#include <string.h>

extern UART_HandleTypeDef huart3;          /* main.c 中定义，跳转 APP 前需要关闭 */
extern DMA_HandleTypeDef hdma_usart3_rx;
extern IWDG_HandleTypeDef hiwdg;           /* main.c 中定义的看门狗句柄 */

#define IWDG_FEED()  HAL_IWDG_Refresh(&hiwdg)

/* 外部看门狗 TPL5010：DONE 引脚接 PC13(LED2)，电平翻转即喂狗，超时约 7 分钟。
 * BootLoader 的擦除/拷贝全程远小于 7 分钟，这里做保险性喂狗 */
#define EXT_WDOG_FEED()  HAL_GPIO_TogglePin(GPIOC, GPIO_PIN_13)




/*======================================================================*/
/*=========================Flash 驱动程序(start)=========================*/
/*======================================================================*/

static void FLASH_Unlock_Reg(void)
{
	if ((FLASH->CR & FLASH_CR_LOCK) == 0) return;
	FLASH->KEYR = 0x45670123;
	FLASH->KEYR = 0xCDEF89AB;
}

static void FLASH_Lock_Reg(void)
{
	FLASH->CR |= FLASH_CR_LOCK;
}

/* 等待内部 Flash 空闲，带超时；返回 1=就绪，0=超时 */
static int FLASH_WaitBusy(uint32_t timeout_ms)
{
	uint32_t start = HAL_GetTick();
	while (FLASH->SR & FLASH_SR_BSY)
	{
		if ((HAL_GetTick() - start) > timeout_ms) return 0;
	}
	return 1;
}

static void FLASH_ClearFlags(void)
{
	FLASH->SR = (FLASH_SR_EOP | FLASH_SR_PGERR | FLASH_SR_WRPRTERR);
}

static int Erase_page(unsigned int pageaddr, unsigned int num)
{
	log_info("> Erasing %u pages from 0x%X...", num, pageaddr);

	for (unsigned int i = 0; i < num; i++)
	{
		unsigned int pg_addr = pageaddr + i * PageSize;

		IWDG_FEED();
		EXT_WDOG_FEED();

		FLASH_Unlock_Reg();
		FLASH_ClearFlags();
		if (!FLASH_WaitBusy(500))
		{
			log_error("> Erase FAIL at 0x%X: busy timeout", pg_addr);
			FLASH_Lock_Reg();
			return 0;
		}

		FLASH->CR |= FLASH_CR_PER;
		FLASH->AR = pg_addr;
		FLASH->CR |= FLASH_CR_STRT;
		if (!FLASH_WaitBusy(500))
		{
			log_error("> Erase FAIL at 0x%X: busy timeout", pg_addr);
			FLASH->CR &= ~FLASH_CR_PER;
			FLASH_Lock_Reg();
			return 0;
		}

		uint32_t sr = FLASH->SR;
		FLASH->CR &= ~FLASH_CR_PER;

		if (sr & FLASH_SR_WRPRTERR)
		{
			log_error("> Erase FAIL at 0x%X: SR=0x%X", pg_addr, (unsigned int)sr);
			FLASH_ClearFlags();
			FLASH_Lock_Reg();
			return 0;
		}
		FLASH_Lock_Reg();
	}

	log_info("> Erase OK (%u pages)", num);
	return 1;
}

static int FLASH_WriteHalfWord(unsigned int addr, uint16_t data)
{
	FLASH_Unlock_Reg();
	FLASH_ClearFlags();
	if (!FLASH_WaitBusy(100))
	{
		FLASH_Lock_Reg();
		return 0;
	}

	FLASH->CR |= FLASH_CR_PG;
	*(__IO uint16_t *)addr = data;
	if (!FLASH_WaitBusy(100))
	{
		FLASH->CR &= ~FLASH_CR_PG;
		FLASH_Lock_Reg();
		return 0;
	}

	uint32_t sr = FLASH->SR;
	FLASH->CR &= ~FLASH_CR_PG;
	FLASH_ClearFlags();

	if (sr & FLASH_SR_WRPRTERR)
	{
		FLASH_Lock_Reg();
		return 0;
	}

	FLASH_Lock_Reg();
	return 1;
}

static int WriteFlash(unsigned int addr, unsigned int * buff, int word_size)
{
	for (int i = 0; i < word_size; i++)
	{
		unsigned int val = buff[i];
		uint16_t low  = (uint16_t)(val & 0xFFFF);
		uint16_t high = (uint16_t)(val >> 16);

		if (!FLASH_WriteHalfWord(addr + 4 * i, low))
		{
			log_error("> Write FAIL at 0x%X (low=0x%04X)", addr + 4 * i, low);
			return 0;
		}
		if (!FLASH_WriteHalfWord(addr + 4 * i + 2, high))
		{
			log_error("> Write FAIL at 0x%X (high=0x%04X)", addr + 4 * i + 2, high);
			return 0;
		}
	}
	return 1;
}

static void ReadFlash(unsigned int addr, unsigned int * buff, uint16_t word_size)
{
	for(int i =0; i < word_size; i++)
	{
		buff[i] = *(__IO unsigned int*)(addr + 4 * i);
	}
}
/*====================================================================*/
/*=========================Flash 驱动程序(end)=========================*/
/*====================================================================*/

/* MODBUS CRC16 */
static uint16_t crc16_modbus_update(uint16_t crc, uint8_t data)
{
    crc ^= data;
    for (int i = 0; i < 8; i++) {
        if (crc & 1) crc = (crc >> 1) ^ 0xA001;
        else crc >>= 1;
    }
    return crc;
}
static uint16_t crc16_modbus(const uint8_t *data, uint32_t len)
{
    uint16_t crc = 0xFFFF;
    for (uint32_t i = 0; i < len; i++) crc = crc16_modbus_update(crc, data[i]);
    return crc;
}

/*======================================================================*/
/* 外部 Flash 与内部 Flash 之间拷贝 */
/*======================================================================*/
#define COPY_CHUNK_WORDS  256   /* 1024 bytes */
#define COPY_CHUNK_BYTES  (COPY_CHUNK_WORDS * 4)

static int Copy_External_To_App1(uint32_t ext_addr, uint32_t byte_size)
{
	static uint32_t words[COPY_CHUNK_WORDS];
	static uint8_t  tmp[COPY_CHUNK_BYTES];
	uint32_t offset = 0;

	log_info("[BL] Copy external 0x%06X -> APP1 0x%08X, size=%u",
		(unsigned)ext_addr, Application_1_Addr, (unsigned)byte_size);

	while (offset < byte_size)
	{
		uint32_t chunk = (byte_size - offset > COPY_CHUNK_BYTES) ? COPY_CHUNK_BYTES : (byte_size - offset);
		IWDG_FEED();
		EXT_WDOG_FEED();
		W25Q64_Read(ext_addr + offset, tmp, chunk);

		for (uint32_t i = 0; i < chunk / 4; i++) {
			words[i] = ((uint32_t)tmp[i * 4]) |
				       ((uint32_t)tmp[i * 4 + 1] << 8) |
				       ((uint32_t)tmp[i * 4 + 2] << 16) |
				       ((uint32_t)tmp[i * 4 + 3] << 24);
		}

		if (!WriteFlash(Application_1_Addr + offset, words, chunk / 4))
		{
			log_error("[FAIL] Write APP1 at 0x%X failed", Application_1_Addr + offset);
			return 0;
		}

		for (uint32_t i = 0; i < chunk / 4; i++) {
			uint32_t r = *(__IO uint32_t *)(Application_1_Addr + offset + i * 4);
			if (r != words[i]) {
				log_error("[FAIL] Verify APP1 at 0x%X: wrote=0x%08X read=0x%08X",
					Application_1_Addr + offset + i * 4, (unsigned)words[i], (unsigned)r);
				return 0;
			}
		}

		offset += chunk;
		if ((offset & 0x3FFF) == 0) {
			log_info("> Progress: %u/%u KB (%u%%)", offset / 1024, byte_size / 1024, offset * 100 / byte_size);
		}
	}
	log_info("[ OK ] Copy done");
	return 1;
}

static int Copy_App1_To_External(uint32_t ext_addr, uint32_t byte_size)
{
	static uint32_t words[COPY_CHUNK_WORDS];
	static uint8_t  tmp[COPY_CHUNK_BYTES];
	uint32_t offset = 0;

	log_info("[BL] Backup APP1 -> external 0x%06X, size=%u", (unsigned)ext_addr, (unsigned)byte_size);

	while (offset < byte_size)
	{
		uint32_t chunk = (byte_size - offset > COPY_CHUNK_BYTES) ? COPY_CHUNK_BYTES : (byte_size - offset);
		IWDG_FEED();
		EXT_WDOG_FEED();
		ReadFlash(Application_1_Addr + offset, words, chunk / 4);
		for (uint32_t i = 0; i < chunk / 4; i++) {
			tmp[i * 4]     = (uint8_t)(words[i]);
			tmp[i * 4 + 1] = (uint8_t)(words[i] >> 8);
			tmp[i * 4 + 2] = (uint8_t)(words[i] >> 16);
			tmp[i * 4 + 3] = (uint8_t)(words[i] >> 24);
		}
		W25Q64_WriteBuffer(ext_addr + offset, tmp, chunk);
		if (!W25Q64_Verify(ext_addr + offset, tmp, chunk)) {
			log_error("[FAIL] W25Q64 backup verify at 0x%06X", (unsigned)(ext_addr + offset));
			return 0;
		}
		offset += chunk;
	}
	return 1;
}


/* 校验固件完整性：检查栈顶地址和复位向量是否合法 */
static int Verify_Firmware(unsigned int addr, unsigned int size)
{
	/* 检查栈顶地址是否在合法RAM范围 (0x20000000) */
	unsigned int sp = *(__IO unsigned int *)addr;
	if ((sp & 0x2FFE0000) != 0x20000000)
	{
		log_error("> Verify fail: Invalid SP=0x%X", sp);
		return 0;
	}

	/* 检查复位向量是否在APP1固件范围内（固件链接地址是APP1） */
	unsigned int reset_vector = *(__IO unsigned int *)(addr + 4);
	if (reset_vector < Application_1_Addr || reset_vector > (Application_1_Addr + size))
	{
		log_error("> Verify fail: Invalid reset vector=0x%X", reset_vector);
		return 0;
	}

	log_info("> Firmware verify OK");
	return 1;
}



/* 设置主堆栈的栈值 */
__asm void MSR_MSP (unsigned int ulAddr)
{
    MSR MSP, r0 			                   //set Main Stack value
    BX r14
}




/* 程序跳转函数: 检查栈顶合法性后跳转到目标 APP */
typedef void (*Jump_Fun)(void);
void IAP_ExecuteApp (unsigned int App_Addr)
{
	Jump_Fun JumpToApp;

	if ( ( ( * ( __IO unsigned int * ) App_Addr ) & 0x2FFE0000 ) == 0x20000000 )
	{
		unsigned int sp = *(__IO unsigned int *)App_Addr;
		unsigned int pc = *(__IO unsigned int *)(App_Addr + 4);
		log_info("[BL] Jumping to APP1 @0x%08X...", App_Addr);

		/* 跳转前关闭 BootLoader 使用的中断和外设，避免干扰 APP。
		 * 注意不能用 __disable_irq()：跳转不是复位，PRIMASK 会带进 APP。 */
		HAL_UART_DMAStop(&huart3);
		__HAL_UART_DISABLE_IT(&huart3, UART_IT_IDLE);
		HAL_NVIC_DisableIRQ(USART3_IRQn);
		HAL_NVIC_DisableIRQ(DMA1_Channel3_IRQn);
		NVIC_ClearPendingIRQ(USART3_IRQn);
		NVIC_ClearPendingIRQ(DMA1_Channel3_IRQn);
		SysTick->CTRL = 0;                      /* 停止 SysTick */
		SCB->ICSR = SCB_ICSR_PENDSTCLR_Msk;     /* 清 SysTick 挂起位 */

		JumpToApp = (Jump_Fun)pc;
		MSR_MSP(sp);
		JumpToApp();
	}
	else
	{
		log_error("[BL] Error: APP @0x%08X SP invalid!", App_Addr);
	}
}





void Start_BootLoader(void)
{
	log_info("");
	log_info("* * * * * * * * * * * * * * * * * *");
	log_info("*           BootLoader            *");
	log_info("*           Flash: %dKB           *", *(volatile uint16_t *)0x1FFFF7E0);
	log_info("* * * * *    By Helei   * * * * * *");

	EXT_WDOG_FEED();    /* 进 BootLoader 先喂一次外部看门狗 */

	/* 读取并清除复位原因：用于区分"固件卡死(看门狗复位)"和"现场断电/按键/正常流程" */
	uint32_t reset_flags = RCC->CSR;
	RCC->CSR |= RCC_CSR_RMVF;

	W25Q64_OtaInfo info;
	int info_ok = (W25Q64_OTA_ReadInfo(&info) == 0);

	/* 固件大小合法性检查：信息区 checksum 较弱，防止损坏的 size 导致越界写内部 Flash */
	if (info_ok && info.new_valid
		&& (info.update_state == W25Q64_STATE_DOWNLOADED || info.update_state == W25Q64_STATE_COPYING)
		&& (info.new_size == 0 || info.new_size > Application_Size || (info.new_size & 0x3) != 0))
	{
		log_warn("[BL] Invalid new_size=%u, discard new image", (unsigned)info.new_size);
		W25Q64_OTA_ClearNew();
		info.new_valid = 0;    /* 阻止进入下面的更新分支 */
	}

	/* 1. 新固件已下载到外部 Flash：备份当前 APP1，然后写入 APP1。
	 * COPYING 也按下载完成处理：拷贝/擦除中途掉电或看门狗复位后，
	 * 外部新固件仍在，重新执行更新流程，否则会卡死在"无有效 APP"状态 */
	if (info_ok && info.new_valid
		&& (info.update_state == W25Q64_STATE_DOWNLOADED || info.update_state == W25Q64_STATE_COPYING))
	{
		log_info("[BL] New firmware in external: ver=%u size=%u crc=0x%04X",
			(unsigned)info.new_version, (unsigned)info.new_size, (unsigned)info.new_crc16);

		W25Q64_OTA_SetState(W25Q64_STATE_COPYING);

		/* 备份策略：当前 APP1 有效时，必须先做出有效备份才允许更新，
		 * 否则新固件不健康时无版本可回退（变砖）；备份失败则放弃本次更新，
		 * 外部新固件保留，下次重启重试（APP1 本身已损坏时无物可备，直接更新自救） */
		int app1_valid = Verify_Firmware(Application_1_Addr, Application_Size);
		if (!info.backup_valid && app1_valid)
		{
			/* 备份前必须先擦除备份槽：W25Q64 只能把位从 1 写 0，不擦除直接写会得到错误数据 */
			log_info("[BL] Erase backup slot...");
			W25Q64_EraseRegion(W25Q64_SLOT_BACKUP_ADDR, Application_Size);
			uint16_t backup_crc = crc16_modbus((const uint8_t *)Application_1_Addr, Application_Size);
			if (Copy_App1_To_External(W25Q64_SLOT_BACKUP_ADDR, Application_Size)
				&& W25Q64_OTA_SetBackupInfo(Application_Size, backup_crc) == 0)
			{
				info.backup_valid = 1;    /* 同步本地副本，下面的更新分支要用 */
				log_info("[BL] Backup OK, crc=0x%04X", (unsigned)backup_crc);
			}
			else
			{
				log_error("[BL] Backup FAILED, abort update (keep old APP1, retry next boot)");
				W25Q64_OTA_SetState(W25Q64_STATE_DOWNLOADED);
			}
		}
		else if (!app1_valid)
		{
			log_warn("[BL] APP1 invalid, nothing to backup; update is the only rescue");
		}

		if (info.backup_valid || !app1_valid)
		{
			log_info("[BL] Updating APP1...");
			if (Erase_page(Application_1_Addr, Application_Size / PageSize))
			{
				if (Copy_External_To_App1(W25Q64_SLOT_NEW_ADDR, info.new_size))
				{
					uint16_t app_crc = crc16_modbus((const uint8_t *)Application_1_Addr, info.new_size);
					if (app_crc == info.new_crc16 && Verify_Firmware(Application_1_Addr, info.new_size))
					{
						log_info("[BL] APP1 update OK, crc=0x%04X, jump", (unsigned)app_crc);
						W25Q64_OTA_SetState(W25Q64_STATE_INSTALLED);
						W25Q64_OTA_ClearConfirmed();
						IAP_ExecuteApp(Application_1_Addr);
						return;
					}
					log_error("[BL] APP1 CRC or header mismatch! app_crc=0x%04X", (unsigned)app_crc);
				}
			}
			log_warn("[BL] APP1 update failed, keep external new image for retry");
			W25Q64_OTA_SetState(W25Q64_STATE_DOWNLOADED);
		}
	}

	/* 2. 已安装但未确认：启动次数超过阈值则回退 */
	if (info_ok && info.update_state == W25Q64_STATE_INSTALLED)
	{
		if (info.app_confirmed)
		{
			log_info("[BL] Update confirmed, clean OTA flags");
			W25Q64_OTA_ClearNew();
		}
		else
		{
			/* 只有看门狗复位才消耗试错次数（说明新固件真的卡死）；
			   上电/按键/软件复位不计，避免现场断电或正常流程误回退。
			   APP 健康超时主动复位前会自行调用 W25Q64_OTA_IncTrial() 计一次。 */
			if (reset_flags & (RCC_CSR_IWDGRSTF | RCC_CSR_WWDGRSTF))
			{
				W25Q64_OTA_IncTrial();
				info.boot_trial_count++;
				log_warn("[BL] APP hung (watchdog reset), trial count=%u", (unsigned)info.boot_trial_count);
			}
			else
			{
				log_info("[BL] Trial boot, reset flags=0x%08X (not counted)", (unsigned)reset_flags);
			}

			if (info.boot_trial_count >= W25Q64_MAX_TRIALS)
			{
				log_warn("[BL] Max trials reached, rollback to backup");
				/* 备份大小合法性检查，防止越界写内部 Flash */
				if (info.backup_valid && info.backup_size > 0
					&& info.backup_size <= Application_Size && (info.backup_size & 0x3) == 0
					&& Erase_page(Application_1_Addr, Application_Size / PageSize))
				{
					if (Copy_External_To_App1(W25Q64_SLOT_BACKUP_ADDR, info.backup_size))
					{
						uint16_t app_crc = crc16_modbus((const uint8_t *)Application_1_Addr, info.backup_size);
						if (app_crc == info.backup_crc16 && Verify_Firmware(Application_1_Addr, Application_Size))
						{
							W25Q64_OTA_SetState(W25Q64_STATE_ROLLBACK);
							IAP_ExecuteApp(Application_1_Addr);
							return;
						}
					}
				}
				log_error("[BL] Rollback failed!");
			}
		}
	}

	/* 3. 回退状态：直接跳转当前 APP1 */
	if (info_ok && info.update_state == W25Q64_STATE_ROLLBACK)
	{
		IAP_ExecuteApp(Application_1_Addr);
		return;
	}

	/* 4. 没有 OTA 任务：直接启动 APP1 */
	if (Verify_Firmware(Application_1_Addr, Application_Size))
	{
		IAP_ExecuteApp(Application_1_Addr);
	}
	else
	{
		log_error("[BL] No valid APP1 found!");
	}
}