/**
 * @file    spi_flash.c
 * @brief   W25Q64 SPI Flash driver (STM32G474, SPI2)
 */

#include "spi_flash.h"
#include "stm32g4xx.h"

/* SPI timeout */
#define SPI_TIMEOUT_MS          1000
#define FLASH_POLL_MS           10

/* ==================== Low-level SPI ==================== */

static void spi_delay(uint32_t us)
{
    /* Simple busy-wait at 72MHz: ~72 cycles per us */
    for (uint32_t i = 0; i < (us * 36); i++) {
        __NOP();
    }
}

static void spi_cs_low(void)
{
    GPIOA->BRR = GPIO_BRR_BR0;  /* PA0 = CS */
}

static void spi_cs_high(void)
{
    GPIOA->BSRR = GPIO_BSRR_BS0;  /* PA0 = CS */
}

static uint8_t spi_xfer(uint8_t data)
{
    while (!(SPI2->SR & SPI_SR_TXE));
    *((__IO uint8_t *)&SPI2->DR) = data;
    while (!(SPI2->SR & SPI_SR_RXNE));
    return *((__IO uint8_t *)&SPI2->DR);
}

static void spi_write_buf(const uint8_t *data, uint32_t len)
{
    for (uint32_t i = 0; i < len; i++) {
        while (!(SPI2->SR & SPI_SR_TXE));
        *((__IO uint8_t *)&SPI2->DR) = data[i];
    }
    /* Wait for TX to complete */
    while (SPI2->SR & SPI_SR_BSY);
}

static void spi_read_buf(uint8_t *data, uint32_t len)
{
    for (uint32_t i = 0; i < len; i++) {
        while (!(SPI2->SR & SPI_SR_TXE));
        *((__IO uint8_t *)&SPI2->DR) = 0xFF;
        while (!(SPI2->SR & SPI_SR_RXNE));
        data[i] = *((__IO uint8_t *)&SPI2->DR);
    }
}

/* ==================== API Implementation ==================== */

uint8_t FLASH_ReadStatusReg1(void)
{
    uint8_t sr;
    spi_cs_low();
    spi_xfer(FLASH_CMD_READ_STATUS1);
    sr = spi_xfer(0xFF);
    spi_cs_high();
    return sr;
}

void FLASH_Init(void)
{
    /* Enable GPIOA clock for CS (PA0) */
    RCC->AHB2ENR |= RCC_AHB2ENR_GPIOAEN;
    /* Enable GPIOB clock for SPI2 (PB13=SCK, PB14=MISO, PB15=MOSI) */
    RCC->AHB2ENR |= RCC_AHB2ENR_GPIOBEN;

    /* Configure PA0 as output push-pull */
    GPIOA->MODER &= ~(GPIO_MODER_MODE0_Msk);
    GPIOA->MODER |= (1 << GPIO_MODER_MODE0_Pos);  /* Output */
    GPIOA->OTYPER &= ~GPIO_OTYPER_OT0;             /* Push-pull */
    GPIOA->PUPDR &= ~(GPIO_PUPDR_PUPD0_Msk);       /* No pull */
    GPIOA->OSPEEDR |= (3 << GPIO_OSPEEDR_OSPEED0_Pos); /* High speed */
    spi_cs_high();

    /* Configure PB13(SCK), PB15(MOSI) as AF5, PB14(MISO) as AF5 */
    /* PB13 = AF5 */
    GPIOB->MODER &= ~(GPIO_MODER_MODE13_Msk);
    GPIOB->MODER |= (2 << GPIO_MODER_MODE13_Pos);  /* AF */
    GPIOB->AFR[1] &= ~(0xF << (13 - 8) * 4);
    GPIOB->AFR[1] |= (5 << (13 - 8) * 4);          /* AF5 = SPI2 */
    GPIOB->OSPEEDR |= (3 << GPIO_OSPEEDR_OSPEED13_Pos);

    /* PB14 = AF5 */
    GPIOB->MODER &= ~(GPIO_MODER_MODE14_Msk);
    GPIOB->MODER |= (2 << GPIO_MODER_MODE14_Pos);
    GPIOB->AFR[1] &= ~(0xF << (14 - 8) * 4);
    GPIOB->AFR[1] |= (5 << (14 - 8) * 4);
    GPIOB->OSPEEDR |= (3 << GPIO_OSPEEDR_OSPEED14_Pos);

    /* PB15 = AF5 */
    GPIOB->MODER &= ~(GPIO_MODER_MODE15_Msk);
    GPIOB->MODER |= (2 << GPIO_MODER_MODE15_Pos);
    GPIOB->AFR[1] &= ~(0xF << (15 - 8) * 4);
    GPIOB->AFR[1] |= (5 << (15 - 8) * 4);
    GPIOB->OSPEEDR |= (3 << GPIO_OSPEEDR_OSPEED15_Pos);

    /* Enable SPI2 clock */
    RCC->APB1LENR |= RCC_APB1LENR_SPI2EN;

    /* Configure SPI2: Master, CPOL=0 CPHA=0, 8-bit, MSB first */
    SPI2->CR1 = SPI_CR1_SSM | SPI_CR1_SSI |   /* Software CS */
                (2 << SPI_CR1_BR_Pos) |         /* Prescaler: /4 = 18MHz @ 72MHz */
                SPI_CR1_MSTR | SPI_CR1_SPE;

    /* Verify flash presence */
    uint32_t jedec = FLASH_ReadJEDEC();
    (void)jedec;  /* DEV: check JEDEC ID for production validation */
}

uint32_t FLASH_ReadJEDEC(void)
{
    uint32_t id = 0;
    spi_cs_low();
    spi_xfer(FLASH_CMD_JEDEC_ID);
    id |= (uint32_t)spi_xfer(0xFF) << 16;
    id |= (uint32_t)spi_xfer(0xFF) << 8;
    id |= (uint32_t)spi_xfer(0xFF);
    spi_cs_high();
    return id;
}

bool FLASH_WaitBusy(uint32_t timeout_ms)
{
    uint32_t elapsed = 0;
    while (FLASH_ReadStatusReg1() & FLASH_SR_BUSY) {
        spi_delay(FLASH_POLL_MS * 1000);
        elapsed += FLASH_POLL_MS;
        if (timeout_ms > 0 && elapsed >= timeout_ms) {
            return false;
        }
    }
    return true;
}

static void flash_write_enable(void)
{
    spi_cs_low();
    spi_xfer(FLASH_CMD_WRITE_ENABLE);
    spi_cs_high();
}

bool FLASH_EraseSector(uint32_t addr)
{
    if (addr >= FLASH_TOTAL_SIZE) return false;
    if (!FLASH_WaitBusy(5000)) return false;

    flash_write_enable();
    spi_cs_low();
    spi_xfer(FLASH_CMD_SECTOR_ERASE);
    spi_xfer((uint8_t)(addr >> 16));
    spi_xfer((uint8_t)(addr >> 8));
    spi_xfer((uint8_t)(addr));
    spi_cs_high();
    return true;
}

bool FLASH_EraseBlock64(uint32_t addr)
{
    if (addr >= FLASH_TOTAL_SIZE) return false;
    if (!FLASH_WaitBusy(10000)) return false;

    flash_write_enable();
    spi_cs_low();
    spi_xfer(FLASH_CMD_BLOCK_ERASE_64K);
    spi_xfer((uint8_t)(addr >> 16));
    spi_xfer((uint8_t)(addr >> 8));
    spi_xfer((uint8_t)(addr));
    spi_cs_high();
    return true;
}

bool FLASH_ChipErase(void)
{
    if (!FLASH_WaitBusy(10000)) return false;

    flash_write_enable();
    spi_cs_low();
    spi_xfer(FLASH_CMD_CHIP_ERASE);
    spi_cs_high();
    return true;
}

bool FLASH_Write(uint32_t addr, const uint8_t *data, uint32_t len)
{
    if (addr >= FLASH_TOTAL_SIZE) return false;
    if (addr + len > FLASH_TOTAL_SIZE) return false;

    while (len > 0) {
        if (!FLASH_WaitBusy(5000)) return false;

        uint32_t page_offset = addr & (FLASH_PAGE_SIZE - 1);
        uint32_t chunk = (len > (FLASH_PAGE_SIZE - page_offset)) ?
                         (FLASH_PAGE_SIZE - page_offset) : len;

        flash_write_enable();
        spi_cs_low();
        spi_xfer(FLASH_CMD_PAGE_PROGRAM);
        spi_xfer((uint8_t)(addr >> 16));
        spi_xfer((uint8_t)(addr >> 8));
        spi_xfer((uint8_t)(addr));
        spi_write_buf(data, chunk);
        spi_cs_high();

        addr += chunk;
        data += chunk;
        len -= chunk;
    }
    return true;
}

void FLASH_Read(uint32_t addr, uint8_t *data, uint32_t len)
{
    if (addr >= FLASH_TOTAL_SIZE) return;
    if (addr + len > FLASH_TOTAL_SIZE) {
        len = FLASH_TOTAL_SIZE - addr;
    }

    spi_cs_low();
    spi_xfer(FLASH_CMD_READ_DATA);
    spi_xfer((uint8_t)(addr >> 16));
    spi_xfer((uint8_t)(addr >> 8));
    spi_xfer((uint8_t)(addr));
    spi_read_buf(data, len);
    spi_cs_high();
}

void FLASH_FastRead(uint32_t addr, uint8_t *data, uint32_t len)
{
    if (addr >= FLASH_TOTAL_SIZE) return;
    if (addr + len > FLASH_TOTAL_SIZE) {
        len = FLASH_TOTAL_SIZE - addr;
    }

    spi_cs_low();
    spi_xfer(FLASH_CMD_FAST_READ);
    spi_xfer((uint8_t)(addr >> 16));
    spi_xfer((uint8_t)(addr >> 8));
    spi_xfer((uint8_t)(addr));
    spi_xfer(0x00);  /* Dummy cycle */
    spi_read_buf(data, len);
    spi_cs_high();
}

void FLASH_PowerDown(void)
{
    spi_cs_low();
    spi_xfer(FLASH_CMD_POWER_DOWN);
    spi_cs_high();
}

void FLASH_ReleasePowerDown(void)
{
    spi_cs_low();
    spi_xfer(FLASH_CMD_RELEASE_PD_ID);
    spi_cs_high();
    spi_delay(5);  /* tRES = 3us typical */
}