/*******************************************************************************
 * @file    mem_flash.c
 * @brief   W25Q32RV SPI Flash Driver Implementation
 * @author  Charles KWON <charleskwon@medithings.co.kr>
 * @date    2025-02-03
 *
 * @details W25Q32RV: 4MB Serial NOR Flash
 *          - SPI Mode 0/3
 *          - Page size: 256 bytes
 *          - Sector size: 4KB
 *          - Erase before write required
 ******************************************************************************/

#include "mem_hal.h"
#include "mem_config.h"

#if MEM_USE_W25Q_FLASH

#include <string.h>
#include <stdio.h>
#include "nrf_gpio.h"
#include "nrf_delay.h"
#include "app_error.h"
#include "../spi2_bus.h"  /* Shared SPI2 bus */

/*==============================================================================
 * W25Q32RV COMMANDS
 *============================================================================*/

#define W25Q_CMD_WRITE_ENABLE       0x06
#define W25Q_CMD_WRITE_DISABLE      0x04
#define W25Q_CMD_READ_STATUS1       0x05
#define W25Q_CMD_READ_STATUS2       0x35
#define W25Q_CMD_WRITE_STATUS       0x01
#define W25Q_CMD_READ_DATA          0x03
#define W25Q_CMD_FAST_READ          0x0B
#define W25Q_CMD_PAGE_PROGRAM       0x02
#define W25Q_CMD_SECTOR_ERASE       0x20    /* 4KB */
#define W25Q_CMD_BLOCK_ERASE_32K    0x52    /* 32KB */
#define W25Q_CMD_BLOCK_ERASE_64K    0xD8    /* 64KB */
#define W25Q_CMD_CHIP_ERASE         0xC7
#define W25Q_CMD_POWER_DOWN         0xB9
#define W25Q_CMD_RELEASE_PD         0xAB
#define W25Q_CMD_DEVICE_ID          0x90
#define W25Q_CMD_JEDEC_ID           0x9F
#define W25Q_CMD_ENABLE_RESET       0x66
#define W25Q_CMD_RESET_DEVICE       0x99

/* Status Register-1 Bits */
#define W25Q_STATUS_BUSY            0x01
#define W25Q_STATUS_WEL             0x02

/* Device ID */
#define W25Q_MANUFACTURER_WINBOND   0xEF
#define W25Q_DEVICE_ID_32RV         0x4016

/*==============================================================================
 * PRIVATE VARIABLES
 *============================================================================*/

static bool m_initialized = false;

/* 4KB sector buffer for Read-Modify-Write */
static uint8_t m_sector_buffer[MEM_SECTOR_SIZE];

/*==============================================================================
 * PRIVATE FUNCTIONS
 *============================================================================*/

/**
 * @brief Assert CS (active low) - Flash CS from shared bus
 */
static inline void cs_low(void)
{
    nrf_gpio_pin_clear(SPI2_CS_FLASH);
}

/**
 * @brief Deassert CS
 */
static inline void cs_high(void)
{
    nrf_gpio_pin_set(SPI2_CS_FLASH);
}

/**
 * @brief SPI transfer with CS control (blocking mode via shared bus)
 */
static ret_code_t spi_transfer(const uint8_t *tx_buf, uint8_t tx_len,
                               uint8_t *rx_buf, uint8_t rx_len)
{
    printf("[FLASH] xfer tx=%d rx=%d...", tx_len, rx_len);
    cs_low();
    ret_code_t err = spi2_bus_transfer(tx_buf, tx_len, rx_buf, rx_len);
    cs_high();
    printf("done(%d)\r\n", err);

    return err;
}

/**
 * @brief Read Status Register-1
 */
static uint8_t w25q_read_status(void)
{
    uint8_t cmd = W25Q_CMD_READ_STATUS1;
    uint8_t rx[2];

    cs_low();
    spi2_bus_transfer(&cmd, 1, rx, 2);
    cs_high();

    return rx[1];
}

/**
 * @brief Wait for BUSY flag to clear
 */
static void w25q_wait_busy(void)
{
    uint32_t timeout = 100000;  /* ~seconds */
    while ((w25q_read_status() & W25Q_STATUS_BUSY) && timeout--)
    {
        nrf_delay_us(10);
    }
}

/**
 * @brief Send Write Enable command
 */
static void w25q_write_enable(void)
{
    uint8_t cmd = W25Q_CMD_WRITE_ENABLE;
    spi_transfer(&cmd, 1, NULL, 0);
}

/**
 * @brief Read data from flash (blocking mode via shared bus)
 */
static ret_code_t w25q_read(uint32_t addr, uint8_t *buf, uint16_t len)
{
    uint8_t cmd[4];
    cmd[0] = W25Q_CMD_READ_DATA;
    cmd[1] = (addr >> 16) & 0xFF;
    cmd[2] = (addr >> 8) & 0xFF;
    cmd[3] = addr & 0xFF;

    cs_low();

    /* Send command */
    ret_code_t err = spi2_bus_transfer(cmd, 4, NULL, 0);
    if (err != NRF_SUCCESS) { cs_high(); return err; }

    /* Read data */
    err = spi2_bus_transfer(NULL, 0, buf, len);
    cs_high();

    return err;
}

/**
 * @brief Program a page (max 256 bytes, must not cross page boundary)
 */
static ret_code_t w25q_page_program(uint32_t addr, const uint8_t *buf, uint16_t len)
{
    if (len > MEM_PAGE_SIZE) len = MEM_PAGE_SIZE;

    w25q_write_enable();

    uint8_t cmd[4];
    cmd[0] = W25Q_CMD_PAGE_PROGRAM;
    cmd[1] = (addr >> 16) & 0xFF;
    cmd[2] = (addr >> 8) & 0xFF;
    cmd[3] = addr & 0xFF;

    cs_low();

    /* Send command */
    ret_code_t err = spi2_bus_transfer(cmd, 4, NULL, 0);
    if (err != NRF_SUCCESS) { cs_high(); return err; }

    /* Send data */
    err = spi2_bus_transfer(buf, len, NULL, 0);
    cs_high();

    if (err == NRF_SUCCESS)
    {
        w25q_wait_busy();  /* Wait for program complete (~0.25ms typ) */
    }
    return err;
}

/**
 * @brief Erase a 4KB sector
 */
static ret_code_t w25q_sector_erase(uint32_t addr)
{
    uint32_t sector_addr = addr & ~(MEM_SECTOR_SIZE - 1);

    w25q_write_enable();

    uint8_t cmd[4];
    cmd[0] = W25Q_CMD_SECTOR_ERASE;
    cmd[1] = (sector_addr >> 16) & 0xFF;
    cmd[2] = (sector_addr >> 8) & 0xFF;
    cmd[3] = sector_addr & 0xFF;

    spi_transfer(cmd, 4, NULL, 0);

    w25q_wait_busy();  /* Wait for erase complete (~30ms typ, 240ms max) */
    return NRF_SUCCESS;
}

/*==============================================================================
 * PUBLIC FUNCTIONS - mem_hal.h IMPLEMENTATION
 *============================================================================*/

ret_code_t mem_init(void)
{
    ret_code_t err;

    printf("[FLASH] Init start\r\n");

    /* Ensure shared SPI2 bus is initialized */
    if (!spi2_bus_is_initialized())
    {
        err = spi2_bus_init();
        if (err != NRF_SUCCESS)
        {
            printf("[FLASH] SPI2 init FAIL: %d\r\n", err);
            return err;
        }
    }

    /* Check MISO pin state (should be high when idle) */
    uint32_t miso_state = nrf_gpio_pin_read(SPI2_MISO_PIN);
    printf("[FLASH] MISO(P0.%d)=%d\r\n", SPI2_MISO_PIN, miso_state);

    cs_high();  /* Ensure Flash is deselected */

    /* Release from power-down - try multiple times with longer delay */
    printf("[FLASH] Wake up...\r\n");
    uint8_t cmd = W25Q_CMD_RELEASE_PD;

    for (int i = 0; i < 3; i++)
    {
        spi_transfer(&cmd, 1, NULL, 0);
        nrf_delay_ms(1);  /* Longer delay - 1ms */
    }

    /* Check MISO again after wake */
    miso_state = nrf_gpio_pin_read(SPI2_MISO_PIN);
    printf("[FLASH] MISO after wake=%d\r\n", miso_state);

    m_initialized = true;
    printf("[FLASH] Init OK\r\n");
    return NRF_SUCCESS;
}

void mem_uninit(void)
{
    if (m_initialized)
    {
        /* Enter power-down mode */
        uint8_t cmd = W25Q_CMD_POWER_DOWN;
        spi_transfer(&cmd, 1, NULL, 0);

        /* Note: Don't uninit shared bus - other devices may still use it */
        m_initialized = false;
    }
}

bool mem_is_ready(void)
{
    return m_initialized && !(w25q_read_status() & W25Q_STATUS_BUSY);
}

ret_code_t mem_read_byte(uint32_t addr, uint8_t *out)
{
    return w25q_read(addr, out, 1);
}

ret_code_t mem_write_byte(uint32_t addr, uint8_t data)
{
    return mem_write_bytes(addr, &data, 1);
}

ret_code_t mem_read_bytes(uint32_t addr, uint8_t *buf, uint16_t len)
{
    return w25q_read(addr, buf, len);
}

ret_code_t mem_write_bytes(uint32_t addr, const uint8_t *buf, uint16_t len)
{
    /*
     * Flash write requires Read-Modify-Write:
     * 1. Read entire sector
     * 2. Modify data in buffer
     * 3. Erase sector
     * 4. Write back entire sector
     */

    uint32_t sector_addr = addr & ~(MEM_SECTOR_SIZE - 1);
    uint32_t offset = addr - sector_addr;

    /* Check if data fits in one sector */
    if (offset + len > MEM_SECTOR_SIZE)
    {
        /* TODO: Handle cross-sector writes */
        return NRF_ERROR_INVALID_LENGTH;
    }

    /* 1. Read entire sector */
    ret_code_t err = w25q_read(sector_addr, m_sector_buffer, MEM_SECTOR_SIZE);
    if (err != NRF_SUCCESS) return err;

    /* 2. Modify data */
    memcpy(&m_sector_buffer[offset], buf, len);

    /* 3. Erase sector */
    err = w25q_sector_erase(sector_addr);
    if (err != NRF_SUCCESS) return err;

    /* 4. Write back page by page */
    for (uint16_t i = 0; i < MEM_SECTOR_SIZE; i += MEM_PAGE_SIZE)
    {
        err = w25q_page_program(sector_addr + i, &m_sector_buffer[i], MEM_PAGE_SIZE);
        if (err != NRF_SUCCESS) return err;
    }

    return NRF_SUCCESS;
}

ret_code_t mem_read_word(uint32_t addr, uint16_t *out)
{
    uint8_t buf[2];
    ret_code_t err = w25q_read(addr, buf, 2);
    if (err == NRF_SUCCESS)
    {
        *out = buf[0] | (buf[1] << 8);
    }
    return err;
}

ret_code_t mem_write_word(uint32_t addr, uint16_t data)
{
    uint8_t buf[2] = { data & 0xFF, (data >> 8) & 0xFF };
    return mem_write_bytes(addr, buf, 2);
}

ret_code_t mem_read_uint32(uint32_t addr, uint32_t *out)
{
    uint8_t buf[4];
    ret_code_t err = w25q_read(addr, buf, 4);
    if (err == NRF_SUCCESS)
    {
        *out = buf[0] | (buf[1] << 8) | (buf[2] << 16) | (buf[3] << 24);
    }
    return err;
}

ret_code_t mem_write_uint32(uint32_t addr, uint32_t data)
{
    uint8_t buf[4] = {
        data & 0xFF,
        (data >> 8) & 0xFF,
        (data >> 16) & 0xFF,
        (data >> 24) & 0xFF
    };
    return mem_write_bytes(addr, buf, 4);
}

ret_code_t mem_read_uint16_array(uint32_t addr, uint16_t *arr, uint16_t count)
{
    return w25q_read(addr, (uint8_t *)arr, count * 2);
}

ret_code_t mem_write_uint16_array(uint32_t addr, const uint16_t *arr, uint16_t count)
{
    return mem_write_bytes(addr, (const uint8_t *)arr, count * 2);
}

ret_code_t mem_read_decrypted(uint32_t addr, uint8_t *buf, uint16_t len)
{
    /* TODO: Implement AES-128 CBC decryption */
    /* For now, just read raw data */
    return w25q_read(addr, buf, len);
}

ret_code_t mem_write_encrypted(uint32_t addr, const uint8_t *buf, uint16_t len)
{
    /* TODO: Implement AES-128 CBC encryption */
    /* For now, just write raw data */
    return mem_write_bytes(addr, buf, len);
}

ret_code_t mem_erase_sector(uint32_t addr)
{
    return w25q_sector_erase(addr);
}

ret_code_t mem_get_device_id(uint8_t *manufacturer_id, uint16_t *device_id)
{
    uint8_t cmd = W25Q_CMD_JEDEC_ID;
    uint8_t rx[4] = {0};

    printf("[FLASH] Reading JEDEC ID...\r\n");

    cs_low();
    ret_code_t err = spi2_bus_transfer(&cmd, 1, rx, 4);
    cs_high();

    if (err != NRF_SUCCESS)
    {
        printf("[FLASH] JEDEC xfer err: %d\r\n", err);
        return err;
    }

    /* rx[1] = Manufacturer ID (0xEF for Winbond) */
    /* rx[2:3] = Device ID (0x4016 for W25Q32RV) */
    *manufacturer_id = rx[1];
    *device_id = (rx[2] << 8) | rx[3];

    printf("[FLASH] JEDEC: MFR=0x%02X DEV=0x%04X\r\n", rx[1], (rx[2]<<8)|rx[3]);

    return NRF_SUCCESS;
}

#endif /* MEM_USE_W25Q_FLASH */
