#include "w25q64.h"
#include "log.h"
#include "app_loop.h"   /* storage_lock/unlock：跨任务访问互斥 */
#include <string.h>

/* *******************************************
 * 软件 SPI 位操作宏（依赖 main.h 中定义的引脚）
 * ******************************************* */
#ifndef W25Q64_CS_L
#define W25Q64_CS_L()     HAL_GPIO_WritePin(W25Q64_CS_PORT,  W25Q64_CS_PIN,  GPIO_PIN_RESET)
#define W25Q64_CS_H()     HAL_GPIO_WritePin(W25Q64_CS_PORT,  W25Q64_CS_PIN,  GPIO_PIN_SET)
#define W25Q64_CLK_L()    HAL_GPIO_WritePin(W25Q64_CLK_PORT, W25Q64_CLK_PIN, GPIO_PIN_RESET)
#define W25Q64_CLK_H()    HAL_GPIO_WritePin(W25Q64_CLK_PORT, W25Q64_CLK_PIN, GPIO_PIN_SET)
#define W25Q64_DI_L()     HAL_GPIO_WritePin(W25Q64_DI_PORT,  W25Q64_DI_PIN,  GPIO_PIN_RESET)
#define W25Q64_DI_H()     HAL_GPIO_WritePin(W25Q64_DI_PORT,  W25Q64_DI_PIN,  GPIO_PIN_SET)
#define W25Q64_DO_READ()  HAL_GPIO_ReadPin (W25Q64_DO_PORT,  W25Q64_DO_PIN)
#endif

static void W25Q64_Delay(void)
{
    /* 72MHz 下 GPIO 翻转本身已有足够延时，这里加少量空操作稳定时序 */
    __NOP(); __NOP(); __NOP(); __NOP();
}

static uint8_t W25Q64_SPI_TransferByte(uint8_t tx)
{
    uint8_t rx = 0;
    for (int i = 0; i < 8; i++) {
        if (tx & 0x80) W25Q64_DI_H(); else W25Q64_DI_L();
        tx <<= 1;
        W25Q64_Delay();
        W25Q64_CLK_H();
        W25Q64_Delay();
        rx <<= 1;
        if (W25Q64_DO_READ() == GPIO_PIN_SET) rx |= 0x01;
        W25Q64_CLK_L();
    }
    return rx;
}

static void W25Q64_SPI_WriteByte(uint8_t tx)
{
    (void)W25Q64_SPI_TransferByte(tx);
}

static uint8_t W25Q64_SPI_ReadByte(void)
{
    return W25Q64_SPI_TransferByte(0xFF);
}

/* *******************************************
 * 初始化
 * ******************************************* */
uint8_t W25Q64_Init(void)
{
    GPIO_InitTypeDef gpio = {0};

    __HAL_RCC_GPIOA_CLK_ENABLE();
    __HAL_RCC_GPIOB_CLK_ENABLE();
    __HAL_RCC_AFIO_CLK_ENABLE();

    /* 释放 PA15、PB3、PB4 作为 GPIO（关闭 JTAG，保留 SWD） */
    __HAL_AFIO_REMAP_SWJ_NOJTAG();

    gpio.Mode  = GPIO_MODE_OUTPUT_PP;
    gpio.Pull  = GPIO_NOPULL;
    gpio.Speed = GPIO_SPEED_FREQ_HIGH;

    gpio.Pin = W25Q64_CS_PIN;  HAL_GPIO_Init(W25Q64_CS_PORT,  &gpio);
    gpio.Pin = W25Q64_CLK_PIN; HAL_GPIO_Init(W25Q64_CLK_PORT, &gpio);
    gpio.Pin = W25Q64_DI_PIN;  HAL_GPIO_Init(W25Q64_DI_PORT,  &gpio);

    gpio.Mode = GPIO_MODE_INPUT;
    gpio.Pull = GPIO_PULLUP;
    gpio.Pin  = W25Q64_DO_PIN; HAL_GPIO_Init(W25Q64_DO_PORT,  &gpio);

    W25Q64_CS_H();
    W25Q64_CLK_L();
    W25Q64_DI_L();

    uint8_t id[3] = {0};
    W25Q64_ReadID(id);
    uint32_t id32 = ((uint32_t)id[0] << 16) | ((uint32_t)id[1] << 8) | id[2];
    if (id32 != W25Q64_JEDEC_ID) {
        log_error("> W25Q64 init FAIL! ID=0x%06X", (unsigned)id32);
        return 1;   /* ID 不匹配 */
    }
    log_info("> W25Q64 init OK, ID=0x%06X", (unsigned)id32);
    return 0;
}

void W25Q64_ReadID(uint8_t *id3)
{
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_READ_ID);
    id3[0] = W25Q64_SPI_ReadByte();
    id3[1] = W25Q64_SPI_ReadByte();
    id3[2] = W25Q64_SPI_ReadByte();
    W25Q64_CS_H();
}

uint8_t W25Q64_ReadStatus(void)
{
    uint8_t status;
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_READ_STATUS1);
    status = W25Q64_SPI_ReadByte();
    W25Q64_CS_H();
    return status;
}

void W25Q64_WaitBusy(void)
{
    while (W25Q64_ReadStatus() & 0x01) {
        __NOP();
    }
}

void W25Q64_WriteEnable(void)
{
    W25Q64_WaitBusy();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_WRITE_ENABLE);
    W25Q64_CS_H();
}

static void W25Q64_WriteAddr(uint32_t addr)
{
    W25Q64_SPI_WriteByte((uint8_t)(addr >> 16));
    W25Q64_SPI_WriteByte((uint8_t)(addr >> 8));
    W25Q64_SPI_WriteByte((uint8_t)(addr));
}

void W25Q64_SectorErase(uint32_t addr)
{
    W25Q64_WriteEnable();
    W25Q64_WaitBusy();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_SECTOR_ERASE);
    W25Q64_WriteAddr(addr);
    W25Q64_CS_H();
    W25Q64_WaitBusy();
}

void W25Q64_BlockErase(uint32_t addr)
{
    W25Q64_WriteEnable();
    W25Q64_WaitBusy();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_BLOCK_ERASE);
    W25Q64_WriteAddr(addr);
    W25Q64_CS_H();
    W25Q64_WaitBusy();
}

void W25Q64_ChipErase(void)
{
    W25Q64_WriteEnable();
    W25Q64_WaitBusy();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_CHIP_ERASE);
    W25Q64_CS_H();
    W25Q64_WaitBusy();
}

void W25Q64_PageProgram(uint32_t addr, const uint8_t *data, uint16_t len)
{
    if (len == 0 || len > W25Q64_PAGE_SIZE) return;

    W25Q64_WriteEnable();
    W25Q64_WaitBusy();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_PAGE_PROGRAM);
    W25Q64_WriteAddr(addr);
    for (uint16_t i = 0; i < len; i++) {
        W25Q64_SPI_WriteByte(data[i]);
    }
    W25Q64_CS_H();
    W25Q64_WaitBusy();
}

void W25Q64_Read(uint32_t addr, uint8_t *buf, uint32_t len)
{
    storage_lock();
    W25Q64_CS_L();
    W25Q64_SPI_WriteByte(W25Q64_CMD_READ_DATA);
    W25Q64_WriteAddr(addr);
    for (uint32_t i = 0; i < len; i++) {
        buf[i] = W25Q64_SPI_ReadByte();
    }
    W25Q64_CS_H();
    storage_unlock();
}

void W25Q64_EraseRegion(uint32_t addr, uint32_t size)
{
    if (size == 0) return;
    storage_lock();
    uint32_t end = addr + size - 1;
    uint32_t pos = addr;

    /* 对齐到 64KB Block 边界 */
    while ((pos & (W25Q64_BLOCK_SIZE - 1)) != 0 && pos <= end) {
        W25Q64_SectorErase(pos);
        pos += W25Q64_SECTOR_SIZE;
    }
    /* 64KB Block 擦除 */
    while ((pos + W25Q64_BLOCK_SIZE - 1) <= end) {
        W25Q64_BlockErase(pos);
        pos += W25Q64_BLOCK_SIZE;
    }
    /* 剩余 4KB Sector */
    while (pos <= end) {
        W25Q64_SectorErase(pos);
        pos += W25Q64_SECTOR_SIZE;
    }
    storage_unlock();
}

int W25Q64_WriteBuffer(uint32_t addr, const uint8_t *data, uint32_t len)
{
    storage_lock();
    while (len) {
        uint32_t page_offset = addr & (W25Q64_PAGE_SIZE - 1);
        uint32_t page_remain = W25Q64_PAGE_SIZE - page_offset;
        uint16_t chunk = (len > page_remain) ? (uint16_t)page_remain : (uint16_t)len;
        W25Q64_PageProgram(addr, data, chunk);
        addr += chunk;
        data += chunk;
        len  -= chunk;
    }
    storage_unlock();
    return 1;
}

int W25Q64_Verify(uint32_t addr, const uint8_t *data, uint32_t len)
{
    storage_lock();
    while (len) {
        uint16_t chunk = (len > 256) ? 256 : (uint16_t)len;
        uint8_t tmp[256];
        W25Q64_Read(addr, tmp, chunk);
        if (memcmp(tmp, data, chunk) != 0) { storage_unlock(); return 0; }
        addr += chunk;
        data += chunk;
        len  -= chunk;
    }
    storage_unlock();
    return 1;
}

/* *******************************************
 * OTA 信息区（双副本校验和）
 * ******************************************* */
static uint16_t W25Q64_InfoChecksum(const W25Q64_OtaInfo *info)
{
    const uint8_t *p = (const uint8_t *)info;
    uint16_t sum = 0;
    for (uint32_t i = 0; i < (sizeof(W25Q64_OtaInfo) - 2); i++) {
        sum += p[i];
    }
    return sum;
}

static int W25Q64_InfoValid(const W25Q64_OtaInfo *info)
{
    if (info->magic != W25Q64_INFO_MAGIC) return 0;
    return (info->checksum == W25Q64_InfoChecksum(info)) ? 1 : 0;
}

int W25Q64_OTA_ReadInfo(W25Q64_OtaInfo *info)
{
    W25Q64_OtaInfo a, b;
    int va = 0, vb = 0;

    storage_lock();
    W25Q64_Read(W25Q64_INFO_SECTOR_A, (uint8_t *)&a, sizeof(a));
    W25Q64_Read(W25Q64_INFO_SECTOR_B, (uint8_t *)&b, sizeof(b));

    va = W25Q64_InfoValid(&a);
    vb = W25Q64_InfoValid(&b);

    if (!va && !vb) {
        memset(info, 0, sizeof(*info));
        storage_unlock();
        return -1;
    }
    if (va && vb) {
        *info = (a.sequence >= b.sequence) ? a : b;
    } else if (va) {
        *info = a;
    } else {
        *info = b;
    }
    storage_unlock();
    return 0;
}

int W25Q64_OTA_WriteInfo(const W25Q64_OtaInfo *info)
{
    W25Q64_OtaInfo tmp = *info;
    tmp.checksum = W25Q64_InfoChecksum(&tmp);

    storage_lock();
    /* 先写入 A 区，再写入 B 区。掉电时至少保留一个有效副本 */
    W25Q64_SectorErase(W25Q64_INFO_SECTOR_A);
    W25Q64_WriteBuffer(W25Q64_INFO_SECTOR_A, (const uint8_t *)&tmp, sizeof(tmp));

    W25Q64_SectorErase(W25Q64_INFO_SECTOR_B);
    W25Q64_WriteBuffer(W25Q64_INFO_SECTOR_B, (const uint8_t *)&tmp, sizeof(tmp));
    storage_unlock();
    return 0;
}

int W25Q64_OTA_SetNewInfo(uint16_t version, uint32_t size, uint16_t crc16)
{
    int ret;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) {
        memset(&info, 0, sizeof(info));
        info.magic = W25Q64_INFO_MAGIC;
    }
    info.sequence++;
    info.new_version  = version;
    info.new_size     = size;
    info.new_crc16    = crc16;
    info.new_valid    = 1;
    info.update_state = W25Q64_STATE_DOWNLOADED;
    info.boot_trial_count = 0;
    info.app_confirmed = 0;
    info.backup_valid  = 0;    /* 新下载后，下一次更新需要重新备份当前 APP */
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_SetBackupInfo(uint32_t size, uint16_t crc16)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.backup_size   = size;
    info.backup_crc16  = crc16;
    info.backup_valid  = 1;
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_ClearNew(void)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.new_valid = 0;
    info.update_state = W25Q64_STATE_IDLE;
    info.app_confirmed = 0;
    info.backup_valid = 0;     /* 更新已确认，旧备份不再有效 */
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_SetState(uint8_t state)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.update_state = state;
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_IncTrial(void)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.boot_trial_count++;
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_SetConfirmed(void)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.app_confirmed = 1;
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

int W25Q64_OTA_ClearConfirmed(void)
{
    int ret = 0;
    storage_lock();
    W25Q64_OtaInfo info;
    if (W25Q64_OTA_ReadInfo(&info) != 0) { storage_unlock(); return -1; }
    info.sequence++;
    info.app_confirmed = 0;
    ret = W25Q64_OTA_WriteInfo(&info);
    storage_unlock();
    return ret;
}

uint16_t W25Q64_CRC16(const uint8_t *data, uint32_t len)
{
    uint16_t crc = 0xFFFF;
    for (uint32_t i = 0; i < len; i++) {
        crc ^= data[i];
        for (int b = 0; b < 8; b++) {
            if (crc & 0x0001) crc = (crc >> 1) ^ 0xA001;
            else crc >>= 1;
        }
    }
    return crc;
}