#include <string.h>

#include "spi_flash.h"

/* Standard JEDEC SPI NOR opcodes shared by all parts. */
#define CMD_RDID  0x9F   /* read JEDEC ID            */
#define CMD_WREN  0x06   /* write enable             */
#define CMD_RDSR  0x05   /* read status register     */
#define SR_WIP    0x01   /* status: write-in-progress */

/* Largest single transport transaction for read/verify, to bound the
 * buffers the transport has to allocate. */
#define SPI_FLASH_CHUNK 4096u

/* Safety cap on status polling so a stuck WIP can't hang forever. */
#define WIP_POLL_MAX 5000000

/* --- Chip database ------------------------------------------------- */

static const spi_flash_chip chip_db[] = {
    /* Micron MT25QU256 — KCU105 config flash, 1.8 V, 32 MB, 4-byte addr.
     * 4-byte command opcodes (READ4B/PP4B/SUBSEC-ERASE4B) so the full
     * range is reachable regardless of the address-mode latch. */
    { "Micron MT25QU256", 0x20, 0xBB19, 32u * 1024 * 1024, 256, 4096, 4, 0x13, 0x12, 0x21 },
    { "Micron MT25QL256", 0x20, 0xBA19, 32u * 1024 * 1024, 256, 4096, 4, 0x13, 0x12, 0x21 },
    /* 16 MB, 3-byte parts. */
    { "Winbond W25Q128",  0xEF, 0x4018, 16u * 1024 * 1024, 256, 4096, 3, 0x03, 0x02, 0x20 },
    { "Macronix MX25L128",0xC2, 0x2018, 16u * 1024 * 1024, 256, 4096, 3, 0x03, 0x02, 0x20 },
    { "ISSI IS25LP128",   0x9D, 0x6018, 16u * 1024 * 1024, 256, 4096, 3, 0x03, 0x02, 0x20 },
};

#define CHIP_DB_LEN ((int)(sizeof(chip_db) / sizeof(chip_db[0])))

int spi_flash_chip_count(void)
{
    return CHIP_DB_LEN;
}

const spi_flash_chip *spi_flash_chip_by_index(int i)
{
    if (i < 0 || i >= CHIP_DB_LEN) return NULL;
    return &chip_db[i];
}

/* --- Helpers ------------------------------------------------------- */

static int put_addr(uint8_t *b, uint32_t addr, int addr_bytes)
{
    if (addr_bytes == 4) {
        b[0] = (uint8_t)(addr >> 24);
        b[1] = (uint8_t)(addr >> 16);
        b[2] = (uint8_t)(addr >> 8);
        b[3] = (uint8_t)(addr);
        return 4;
    }
    b[0] = (uint8_t)(addr >> 16);
    b[1] = (uint8_t)(addr >> 8);
    b[2] = (uint8_t)(addr);
    return 3;
}

static int write_enable(spi_flash *sf)
{
    uint8_t cmd = CMD_WREN;
    return sf->xfer(sf->ctx, &cmd, 1, NULL, 0);
}

/* Poll the status register until write-in-progress clears. */
static int wait_wip(spi_flash *sf)
{
    long tries;
    for (tries = 0; tries < WIP_POLL_MAX; tries++) {
        uint8_t cmd = CMD_RDSR, st = 0;
        if (sf->xfer(sf->ctx, &cmd, 1, &st, 1) < 0) return -1;
        if (!(st & SR_WIP)) return 0;
    }
    return -1; /* timeout */
}

/* --- Operations ---------------------------------------------------- */

int spi_flash_detect(spi_flash *sf)
{
    uint8_t cmd = CMD_RDID;
    uint8_t id[3] = {0, 0, 0};
    uint16_t dev;
    int i;

    if (!sf || !sf->xfer) return -1;
    if (sf->xfer(sf->ctx, &cmd, 1, id, 3) < 0) return -1;

    memcpy(sf->jedec, id, 3);
    sf->chip = NULL;

    dev = (uint16_t)((id[1] << 8) | id[2]);
    for (i = 0; i < CHIP_DB_LEN; i++) {
        if (chip_db[i].mfg_id == id[0] && chip_db[i].dev_id == dev) {
            sf->chip = &chip_db[i];
            return 0;
        }
    }
    return 1; /* read OK but unknown part */
}

int spi_flash_read(spi_flash *sf, uint32_t addr, uint8_t *buf, size_t len)
{
    uint8_t cmd[1 + 4];
    int n;

    if (!sf || !sf->chip || !buf) return -1;

    while (len) {
        size_t chunk = (len > SPI_FLASH_CHUNK) ? SPI_FLASH_CHUNK : len;
        cmd[0] = sf->chip->read_cmd;
        n = 1 + put_addr(&cmd[1], addr, sf->chip->addr_bytes);
        if (sf->xfer(sf->ctx, cmd, (size_t)n, buf, chunk) < 0) return -1;
        addr += (uint32_t)chunk;
        buf  += chunk;
        len  -= chunk;
    }
    return 0;
}

int spi_flash_erase_sector(spi_flash *sf, uint32_t addr)
{
    uint8_t cmd[1 + 4];
    int n;

    if (!sf || !sf->chip) return -1;

    addr -= (addr % sf->chip->sector_size); /* align to sector base */

    if (write_enable(sf) < 0) return -1;
    cmd[0] = sf->chip->erase_cmd;
    n = 1 + put_addr(&cmd[1], addr, sf->chip->addr_bytes);
    if (sf->xfer(sf->ctx, cmd, (size_t)n, NULL, 0) < 0) return -1;
    return wait_wip(sf);
}

int spi_flash_program_page(spi_flash *sf, uint32_t addr,
                           const uint8_t *data, size_t len)
{
    uint8_t cmd[1 + 4 + 256];
    int n;

    if (!sf || !sf->chip || !data) return -1;
    if (len == 0) return 0;
    if (len > sf->chip->page_size || len > 256) return -1;
    /* must not cross a page boundary */
    if ((addr % sf->chip->page_size) + len > sf->chip->page_size) return -1;

    if (write_enable(sf) < 0) return -1;
    cmd[0] = sf->chip->pp_cmd;
    n = 1 + put_addr(&cmd[1], addr, sf->chip->addr_bytes);
    memcpy(&cmd[n], data, len);
    n += (int)len;
    if (sf->xfer(sf->ctx, cmd, (size_t)n, NULL, 0) < 0) return -1;
    return wait_wip(sf);
}

int spi_flash_program(spi_flash *sf, uint32_t addr,
                      const uint8_t *data, size_t len)
{
    if (!sf || !sf->chip || !data) return -1;

    while (len) {
        uint32_t page = sf->chip->page_size;
        uint32_t off  = addr % page;
        size_t   chunk = page - off;
        if (chunk > len) chunk = len;
        if (spi_flash_program_page(sf, addr, data, chunk) < 0) return -1;
        addr += (uint32_t)chunk;
        data += chunk;
        len  -= chunk;
    }
    return 0;
}

int spi_flash_verify(spi_flash *sf, uint32_t addr,
                     const uint8_t *data, size_t len, uint32_t *mismatch_off)
{
    uint8_t tmp[SPI_FLASH_CHUNK];
    uint32_t done = 0;

    if (!sf || !sf->chip || !data) return -1;

    while (len) {
        size_t chunk = (len > sizeof(tmp)) ? sizeof(tmp) : len;
        size_t i;
        if (spi_flash_read(sf, addr, tmp, chunk) < 0) return -1;
        for (i = 0; i < chunk; i++) {
            if (tmp[i] != data[done + i]) {
                if (mismatch_off) *mismatch_off = done + (uint32_t)i;
                return 1;
            }
        }
        addr += (uint32_t)chunk;
        done += (uint32_t)chunk;
        len  -= chunk;
    }
    return 0;
}