#include "w25q64.h"
#include <string.h>

static W25Q64_Handle_t w25q64 = {0};
static bool w25q64_attached = false;
static bool w25q64_ready = false;

static void w25q64_cs_low(void)
{
    if (w25q64_attached)
    {
        HAL_GPIO_WritePin(w25q64.cs_port, w25q64.cs_pin, GPIO_PIN_RESET);
    }
}

static void w25q64_cs_high(void)
{
    if (w25q64_attached)
    {
        HAL_GPIO_WritePin(w25q64.cs_port, w25q64.cs_pin, GPIO_PIN_SET);
    }
}

static W25Q64_Status_t w25q64_spi_tx(const uint8_t *data, uint16_t len)
{
    if (!w25q64_attached || w25q64.hspi == NULL)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    if (HAL_SPI_Transmit(w25q64.hspi, (uint8_t *) data, len, w25q64.timeout_ms) != HAL_OK)
    {
        return W25Q64_ERR_HAL;
    }
    return W25Q64_OK;
}

static W25Q64_Status_t w25q64_spi_rx(uint8_t *data, uint16_t len)
{
    if (!w25q64_attached || w25q64.hspi == NULL)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    if (HAL_SPI_Receive(w25q64.hspi, data, len, w25q64.timeout_ms) != HAL_OK)
    {
        return W25Q64_ERR_HAL;
    }
    return W25Q64_OK;
}

static W25Q64_Status_t w25q64_write_enable(void)
{
    uint8_t cmd = W25Q64_CMD_WRITE_ENABLE;

    w25q64_cs_low();
    if (w25q64_spi_tx(&cmd, 1U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    w25q64_cs_high();
    return W25Q64_OK;
}

static W25Q64_Status_t w25q64_read_status(uint8_t *status)
{
    uint8_t cmd = W25Q64_CMD_READ_STATUS_1;

    if (status == NULL)
    {
        return W25Q64_ERR_PARAM;
    }

    w25q64_cs_low();
    if (w25q64_spi_tx(&cmd, 1U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    if (w25q64_spi_rx(status, 1U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    w25q64_cs_high();
    return W25Q64_OK;
}

static W25Q64_Status_t w25q64_wait_ready(uint32_t timeout_ms)
{
    uint8_t status = 0;
    uint32_t start = HAL_GetTick();

    do
    {
        if (w25q64_read_status(&status) != W25Q64_OK)
        {
            return W25Q64_ERR_HAL;
        }
        if ((status & 0x01U) == 0U)
        {
            return W25Q64_OK;
        }
    } while ((HAL_GetTick() - start) < timeout_ms);

    return W25Q64_ERR_BUSY;
}

static W25Q64_Status_t w25q64_command_addr(uint8_t cmd, uint32_t addr)
{
    uint8_t header[4] = {
        cmd,
        (uint8_t) (addr >> 16),
        (uint8_t) (addr >> 8),
        (uint8_t) addr
    };

    w25q64_cs_low();
    if (w25q64_spi_tx(header, sizeof(header)) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    return W25Q64_OK;
}

void W25Q64_Attach(W25Q64_Handle_t *handle)
{
    if (handle == NULL)
    {
        W25Q64_Detach();
        return;
    }

    w25q64 = *handle;
    w25q64_attached = (w25q64.hspi != NULL) && (w25q64.cs_port != NULL);
    w25q64.timeout_ms = (w25q64.timeout_ms == 0U) ? 100U : w25q64.timeout_ms;
    if (w25q64_attached)
    {
        HAL_GPIO_WritePin(w25q64.cs_port, w25q64.cs_pin, GPIO_PIN_SET);
    }
    w25q64_ready = false;
}

void W25Q64_Detach(void)
{
    memset(&w25q64, 0, sizeof(w25q64));
    w25q64_attached = false;
    w25q64_ready = false;
}

W25Q64_Status_t W25Q64_ReadID(uint32_t *jedec_id)
{
    uint8_t cmd = W25Q64_CMD_READ_ID;
    uint8_t id[3] = {0};

    if (jedec_id == NULL)
    {
        return W25Q64_ERR_PARAM;
    }
    if (!w25q64_attached)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    w25q64_cs_low();
    if (w25q64_spi_tx(&cmd, 1U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    if (w25q64_spi_rx(id, 3U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    w25q64_cs_high();

    *jedec_id = ((uint32_t) id[0] << 16) | ((uint32_t) id[1] << 8) | id[2];
    return W25Q64_OK;
}

W25Q64_Status_t W25Q64_Init(void)
{
    uint32_t jedec_id = 0;

    if (!w25q64_attached)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    if (W25Q64_ReadID(&jedec_id) != W25Q64_OK)
    {
        w25q64_ready = false;
        return W25Q64_ERR_HAL;
    }

    w25q64_ready = (jedec_id == W25Q64_JEDEC_ID);
    return w25q64_ready ? W25Q64_OK : W25Q64_ERR_ID;
}

W25Q64_Status_t W25Q64_Read(uint32_t addr, uint8_t *data, uint32_t len)
{
    W25Q64_Status_t ret;
    uint32_t offset = 0U;

    if (data == NULL || len == 0U)
    {
        return W25Q64_ERR_PARAM;
    }
    if (!w25q64_ready)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    ret = w25q64_command_addr(W25Q64_CMD_READ_DATA, addr);
    if (ret != W25Q64_OK)
    {
        return ret;
    }

    while (offset < len)
    {
        uint16_t chunk = (uint16_t) ((len - offset) > UINT16_MAX ? UINT16_MAX : (len - offset));
        ret = w25q64_spi_rx(data + offset, chunk);
        if (ret != W25Q64_OK)
        {
            w25q64_cs_high();
            return ret;
        }
        offset += chunk;
    }
    w25q64_cs_high();
    return W25Q64_OK;
}

static W25Q64_Status_t w25q64_page_program(uint32_t addr, const uint8_t *data, uint16_t len)
{
    W25Q64_Status_t ret;

    ret = w25q64_write_enable();
    if (ret != W25Q64_OK)
    {
        return ret;
    }

    ret = w25q64_command_addr(W25Q64_CMD_PAGE_PROGRAM, addr);
    if (ret != W25Q64_OK)
    {
        return ret;
    }

    if (w25q64_spi_tx(data, len) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    w25q64_cs_high();

    return w25q64_wait_ready(w25q64.timeout_ms);
}

W25Q64_Status_t W25Q64_Write(uint32_t addr, const uint8_t *data, uint32_t len)
{
    uint32_t offset = 0U;

    if (data == NULL || len == 0U)
    {
        return W25Q64_ERR_PARAM;
    }
    if (!w25q64_ready)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    while (offset < len)
    {
        uint32_t page_offset = (addr + offset) % W25Q64_PAGE_SIZE;
        uint32_t room = W25Q64_PAGE_SIZE - page_offset;
        uint16_t chunk = (uint16_t) ((len - offset) < room ? (len - offset) : room);
        W25Q64_Status_t ret = w25q64_page_program(addr + offset, data + offset, chunk);
        if (ret != W25Q64_OK)
        {
            return ret;
        }
        offset += chunk;
    }

    return W25Q64_OK;
}

W25Q64_Status_t W25Q64_EraseSector(uint32_t addr)
{
    W25Q64_Status_t ret;

    if (!w25q64_ready)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    ret = w25q64_write_enable();
    if (ret != W25Q64_OK)
    {
        return ret;
    }

    ret = w25q64_command_addr(W25Q64_CMD_SECTOR_ERASE, addr);
    if (ret != W25Q64_OK)
    {
        return ret;
    }
    w25q64_cs_high();

    return w25q64_wait_ready(400U);
}

W25Q64_Status_t W25Q64_Erase64K(uint32_t addr)
{
    W25Q64_Status_t ret;

    if (!w25q64_ready)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    ret = w25q64_write_enable();
    if (ret != W25Q64_OK)
    {
        return ret;
    }

    ret = w25q64_command_addr(W25Q64_CMD_BLOCK_ERASE_64K, addr);
    if (ret != W25Q64_OK)
    {
        return ret;
    }
    w25q64_cs_high();

    return w25q64_wait_ready(2000U);
}

W25Q64_Status_t W25Q64_ChipErase(void)
{
    uint8_t cmd = W25Q64_CMD_CHIP_ERASE;

    if (!w25q64_ready)
    {
        return W25Q64_ERR_NOT_ATTACHED;
    }

    if (w25q64_write_enable() != W25Q64_OK)
    {
        return W25Q64_ERR_HAL;
    }

    w25q64_cs_low();
    if (w25q64_spi_tx(&cmd, 1U) != W25Q64_OK)
    {
        w25q64_cs_high();
        return W25Q64_ERR_HAL;
    }
    w25q64_cs_high();

    return w25q64_wait_ready(100000U);
}