/* 
 * (c) danielinux 2019
 *
 *      GPLv.2
 *
 * See LICENSE for details
 */
#include "system.h"
#include "spi_drv.h"

#define SPI_FLASH_SECTOR_SIZE (4096)
#define SPI_FLASH_PAGE_SIZE   (256)

#define MDID            0x90
#define RDSR            0x05
#define WRSR            0x01
#   define ST_BUSY (1 << 0)
#   define ST_WEL  (1 << 1)
#   define ST_BP0  (1 << 2)
#   define ST_BP1  (1 << 3)
#   define ST_BP2  (1 << 4)
#   define ST_BP3  (1 << 5)
#   define ST_AAI  (1 << 6)
#   define ST_BRO  (1 << 7)
#define WREN            0x06
#define WRDI            0x04
#define SECTOR_ERASE    0x20
#define BYTE_READ       0x03
#define BYTE_WRITE      0x02
#define AUTOINC         0xAD
#define EWSR            0x50
#define EBSY            0x70
#define DBSY            0x80


static enum write_mode {
    WB_WRITEPAGE = 0x00,
    SST_AAI      = 0x01
} chip_write_mode = WB_WRITEPAGE;

static void write_address(uint32_t address)
{
    spi_write((address & 0xFF00) >> 8);
    spi_read();
    spi_write((address & 0xFF0000) >> 16);
    spi_read();
    spi_write((address & 0xFF000000) >> 24);
    spi_read();
}

static uint8_t read_status(void)
{
    uint8_t status;
    int i;
    spi_cs_on();
    spi_write(RDSR);
    spi_read();
    spi_write(0xFF);
    status = spi_read();
    spi_cs_off();
    return status;
}

static void spi_cmd(uint8_t cmd)
{
    spi_cs_on();
    spi_write(cmd);
    spi_read();
    spi_cs_off();
}

static inline void flash_aai_enable(void)
{
    spi_cmd(EBSY);
}

static inline void flash_aai_disable(void)
{
    spi_cmd(DBSY);
}

static void flash_write_enable(void)
{
    uint8_t status;
    do {
        spi_cmd(WREN);
        status = read_status();
    } while ((status & ST_WEL) == 0);
}

static void flash_write_disable(void)
{
    uint8_t status;
    spi_cmd(WRDI);
}
static void wait_busy(void)
{
    uint8_t status;
    do {
        status = read_status();
    } while(status & ST_BUSY);
}

static int spi_flash_write_page(uint32_t address, const void *data, int len)
{
    const uint8_t *buf = data;
    int j = 0;
    while (len > 0) {
        wait_busy();
        flash_write_enable();
        spi_cs_on();
        spi_write(BYTE_WRITE);
        spi_read();
        write_address(address);
        do {
            spi_write(buf[j++]);
            address++;
            spi_read();
            len--;
        } while ((address % SPI_FLASH_PAGE_SIZE) != 0);
        spi_cs_off();
    }
    wait_busy();
    return j;
}

static int spi_flash_write_aai(uint32_t address, const void *data, int len)
{
    const uint8_t *buf = data;
    int j = 0;
    int cont = 0;
    wait_busy();
    if (len < 1)
        return -1;
    while (len > 0) {
        if ((address & 0x01) || (len < 2)) {
            flash_write_enable();
            spi_cs_on();
            spi_write(BYTE_WRITE);
            spi_read();
            write_address(address);
            spi_write(buf[j++]);
            spi_read();
            spi_cs_off();
            len--;
            address++;
        } else {
            if (!cont) {
                flash_aai_enable();
                flash_write_enable();
            }
            spi_cs_on();
            spi_write(AUTOINC);
            spi_read();
            if (!cont) {
                /* First AAI transaction, send address. */
                write_address(address);
                cont = 1;
            }
            spi_write(buf[j++]);
            spi_read();
            spi_write(buf[j++]);
            spi_read();
            spi_cs_off();
            len -= 2;
            address += 2;
            read_status();
        }
    }
    if (cont) {
        flash_write_disable();
        flash_aai_disable();
    }
    wait_busy();
    return j;
}

/* --- */

uint16_t spi_flash_probe(void)
{
    uint8_t manuf, product, b0;
    int i;
    wait_busy();
    spi_cs_on();
    spi_write(MDID);
    b0 = spi_read();

    write_address(0);
    spi_write(0xFF);
    manuf = spi_read();
    spi_write(0xFF);
    product = spi_read();
    spi_cs_off();
    if (manuf == 0xBF)
        chip_write_mode = SST_AAI;
    if (manuf == 0xEF)
        chip_write_mode = WB_WRITEPAGE;

#ifndef READONLY
    spi_cmd(EWSR);
    spi_cs_on();
    spi_write(WRSR);
    spi_read();
    spi_write(0x00);
    spi_read();
    spi_cs_off();
#endif
    return (uint16_t)(manuf << 8 | product);
}


void spi_flash_sector_erase(uint32_t address)
{
    uint8_t status;
    address &= (~(SPI_FLASH_SECTOR_SIZE - 1));

    wait_busy();
    flash_write_enable();
    spi_cs_on();
    spi_write(SECTOR_ERASE);
    spi_read();
    write_address(address);
    spi_cs_off();
    wait_busy();
}

int spi_flash_read(uint32_t address, void *data, int len)
{
    uint8_t *buf = data;
    int i = 0;
    wait_busy();
    spi_cs_on();
    spi_write(BYTE_READ);
    spi_read();
    write_address(address);
    while (len > 0) {
        spi_write(0xFF);
        buf[i++] = spi_read();
        len--;
    }
    spi_cs_off();
    return i;
}

int spi_flash_write(uint32_t address, const void *data, int len)
{
    if (chip_write_mode == SST_AAI)
        return spi_flash_write_aai(address, data, len);
    if (chip_write_mode == WB_WRITEPAGE)
        return spi_flash_write_page(address, data, len);
    return -1;
}