#include "spi_nor.h"

#include <string.h>

/* Стандартные команды одинаковы для проверенных W25Q и SST25. */
#define SPI_NOR_COMMAND_WRITE_ENABLE       0x06U
#define SPI_NOR_COMMAND_READ_STATUS        0x05U
#define SPI_NOR_COMMAND_READ_DATA          0x03U
#define SPI_NOR_COMMAND_PAGE_PROGRAM       0x02U
#define SPI_NOR_COMMAND_SECTOR_ERASE       0x20U
#define SPI_NOR_COMMAND_READ_JEDEC_ID       0x9FU
#define SPI_NOR_COMMAND_RELEASE_POWER_DOWN 0xABU
#define SPI_NOR_STATUS_BUSY                0x01U
#define SPI_NOR_SECTOR_SIZE                4096UL
#define SPI_NOR_MAX_PAGE_SIZE              256U
#define SPI_NOR_VERIFY_CHUNK_SIZE          32U
#define SPI_NOR_24BIT_CAPACITY_LIMIT       0x01000000UL

/* Преобразует результат физического обмена в ошибку протокольного слоя. */
static spi_nor_status_t map_io_status(spi_nor_io_status_t status)
{
    if (status == SPI_NOR_IO_OK) {
        return SPI_NOR_OK;
    }
    if (status == SPI_NOR_IO_TIMEOUT) {
        return SPI_NOR_E_TIMEOUT;
    }
    return SPI_NOR_E_IO;
}

/* Захватывает весь составной вызов, чтобы другой поток не вклинился после WREN. */
static spi_nor_status_t begin_operation(spi_nor_t *device)
{
    spi_nor_io_status_t lock_status;

    if ((device == NULL) || (device->initialized == 0U)) {
        return SPI_NOR_E_ARGUMENT;
    }
    /* Рекурсивный вызов отклоняется до mutex, чтобы не ждать самого себя. */
    if (device->busy != 0U) {
        return SPI_NOR_E_BUSY;
    }
    if (device->io.lock != NULL) {
        lock_status = device->io.lock(device->io.context,
                                      device->config.lock_timeout_ms);
        if (lock_status != SPI_NOR_IO_OK) {
            return (lock_status == SPI_NOR_IO_TIMEOUT) ?
                SPI_NOR_E_TIMEOUT : SPI_NOR_E_BUSY;
        }
    }
    if (device->busy != 0U) {
        if (device->io.unlock != NULL) {
            device->io.unlock(device->io.context);
        }
        return SPI_NOR_E_BUSY;
    }
    device->busy = 1U;
    return SPI_NOR_OK;
}

/* Всегда освобождает как локальный флаг, так и необязательный mutex платформы. */
static void end_operation(spi_nor_t *device)
{
    device->busy = 0U;
    if (device->io.unlock != NULL) {
        device->io.unlock(device->io.context);
    }
}

/* Обмен выполняется только при активном CS, а ошибка не оставляет CS в нуле. */
static spi_nor_status_t transaction(spi_nor_t *device,
                                    const uint8_t *command,
                                    size_t command_size,
                                    const uint8_t *tx_data,
                                    uint8_t *rx_data,
                                    size_t data_size)
{
    spi_nor_io_status_t io_status;

    device->io.chip_select(device->io.context, 1U);
    io_status = device->io.transfer(device->io.context, command, NULL,
                                    command_size,
                                    device->config.io_timeout_ms);
    if ((io_status == SPI_NOR_IO_OK) && (data_size != 0U)) {
        io_status = device->io.transfer(device->io.context, tx_data, rx_data,
                                        data_size,
                                        device->config.io_timeout_ms);
    }
    device->io.chip_select(device->io.context, 0U);
    return map_io_status(io_status);
}

/* Формирует opcode и 24-битный big-endian адрес независимо от endian CPU. */
static void make_address_command(uint8_t command[4], uint8_t opcode,
                                 uint32_t address)
{
    command[0] = opcode;
    command[1] = (uint8_t)(address >> 16U);
    command[2] = (uint8_t)(address >> 8U);
    command[3] = (uint8_t)address;
}

/* Проверяет весь полуоткрытый диапазон без переполнения address + size. */
static spi_nor_status_t validate_range(const spi_nor_t *device,
                                       uint32_t address,
                                       size_t size)
{
    if ((device->detected == 0U) || (size == 0U)) {
        return SPI_NOR_E_ARGUMENT;
    }
    if ((address >= device->info.capacity_bytes) ||
        (size > (size_t)(device->info.capacity_bytes - address)) ||
        (device->info.capacity_bytes > SPI_NOR_24BIT_CAPACITY_LIMIT)) {
        return SPI_NOR_E_BOUNDS;
    }
    return SPI_NOR_OK;
}

/* Читает регистр состояния в отдельной законченной SPI-транзакции. */
static spi_nor_status_t read_status(spi_nor_t *device, uint8_t *status)
{
    uint8_t command = SPI_NOR_COMMAND_READ_STATUS;

    return transaction(device, &command, 1U, NULL, status, 1U);
}

/* Публичное чтение статуса не допускает пересечения с program/erase/read. */
spi_nor_status_t spi_nor_read_status_register(spi_nor_t *device,
                                              uint8_t *status)
{
    spi_nor_status_t result;

    if ((device == NULL) || (status == NULL) || (device->detected == 0U)) {
        return SPI_NOR_E_ARGUMENT;
    }
    result = begin_operation(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    result = read_status(device, status);
    end_operation(device);
    return result;
}

/* Ожидает снятия BUSY с wrap-safe арифметикой tick и конечным тайм-аутом. */
static spi_nor_status_t wait_ready(spi_nor_t *device, uint32_t timeout_ms)
{
    uint32_t started_ms = device->io.tick_ms(device->io.context);
    uint8_t status = SPI_NOR_STATUS_BUSY;
    spi_nor_status_t result;

    while ((status & SPI_NOR_STATUS_BUSY) != 0U) {
        result = read_status(device, &status);
        if (result != SPI_NOR_OK) {
            return result;
        }
        if ((status & SPI_NOR_STATUS_BUSY) == 0U) {
            return SPI_NOR_OK;
        }
        if ((uint32_t)(device->io.tick_ms(device->io.context) - started_ms) >=
            timeout_ms) {
            return SPI_NOR_E_TIMEOUT;
        }
        if (device->config.ready_poll_delay_ms != 0U) {
            device->io.delay_ms(device->io.context,
                                device->config.ready_poll_delay_ms);
        }
    }
    return SPI_NOR_OK;
}

/* Перед каждой изменяющей командой устанавливает volatile Write Enable Latch. */
static spi_nor_status_t write_enable(spi_nor_t *device)
{
    uint8_t command = SPI_NOR_COMMAND_WRITE_ENABLE;

    return transaction(device, &command, 1U, NULL, NULL, 0U);
}

/* Внутреннее чтение не захватывает mutex повторно и используется verify-кодом. */
static spi_nor_status_t raw_read(spi_nor_t *device, uint32_t address,
                                 void *data, size_t size)
{
    uint8_t command[4];

    make_address_command(command, SPI_NOR_COMMAND_READ_DATA, address);
    return transaction(device, command, sizeof(command), NULL, data, size);
}

/* Программирует один page/byte chunk и немедленно проверяет его чтением. */
static spi_nor_status_t program_chunk(spi_nor_t *device,
                                      uint32_t address,
                                      const uint8_t *data,
                                      size_t size)
{
    uint8_t command[4];
    uint8_t verify[SPI_NOR_MAX_PAGE_SIZE];
    spi_nor_status_t result;

    result = write_enable(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    make_address_command(command, SPI_NOR_COMMAND_PAGE_PROGRAM, address);
    result = transaction(device, command, sizeof(command), data, NULL, size);
    if (result != SPI_NOR_OK) {
        return result;
    }
    result = wait_ready(device, device->config.program_timeout_ms);
    if (result != SPI_NOR_OK) {
        return result;
    }
    result = raw_read(device, address, verify, size);
    if (result != SPI_NOR_OK) {
        return result;
    }
    return (memcmp(verify, data, size) == 0) ?
        SPI_NOR_OK : SPI_NOR_E_VERIFY;
}

/* Проверяет каждый байт сектора после erase, а не только первый word. */
static spi_nor_status_t verify_erased_sector(spi_nor_t *device,
                                              uint32_t address)
{
    uint8_t verify[SPI_NOR_VERIFY_CHUNK_SIZE];
    uint32_t offset;
    size_t index;
    spi_nor_status_t result;

    for (offset = 0U; offset < SPI_NOR_SECTOR_SIZE;
         offset += sizeof(verify)) {
        result = raw_read(device, address + offset, verify, sizeof(verify));
        if (result != SPI_NOR_OK) {
            return result;
        }
        for (index = 0U; index < sizeof(verify); ++index) {
            if (verify[index] != 0xFFU) {
                return SPI_NOR_E_VERIFY;
            }
        }
    }
    return SPI_NOR_OK;
}

/* Сопоставляет только проверенные JEDEC; неизвестную ёмкость не угадывает. */
static spi_nor_status_t identify_device(spi_nor_t *device)
{
    const uint8_t manufacturer = device->info.manufacturer_id;
    const uint8_t memory_type = device->info.memory_type;
    const uint8_t capacity = device->info.capacity_code;

    device->info.sector_size = SPI_NOR_SECTOR_SIZE;
    if ((manufacturer == 0xEFU) && (memory_type == 0x40U) &&
        (capacity == 0x17U)) {
        device->info.model = SPI_NOR_MODEL_W25Q64;
        device->info.capacity_bytes = 8UL * 1024UL * 1024UL;
        device->info.page_size = 256U;
        return SPI_NOR_OK;
    }
    if ((manufacturer == 0xEFU) && (memory_type == 0x40U) &&
        (capacity == 0x18U)) {
        device->info.model = SPI_NOR_MODEL_W25Q128;
        device->info.capacity_bytes = 16UL * 1024UL * 1024UL;
        device->info.page_size = 256U;
        return SPI_NOR_OK;
    }
    if ((manufacturer == 0xBFU) && (memory_type == 0x25U) &&
        (capacity == 0x41U)) {
        device->info.model = SPI_NOR_MODEL_SST25VF016B;
        device->info.capacity_bytes = 2UL * 1024UL * 1024UL;
        /* Byte Program 0x02 не требует SST AAI sequencing. */
        device->info.page_size = 1U;
        return SPI_NOR_OK;
    }
    if (((manufacturer == 0x00U) && (memory_type == 0x00U) &&
         (capacity == 0x00U)) ||
        ((manufacturer == 0xFFU) && (memory_type == 0xFFU) &&
         (capacity == 0xFFU))) {
        return SPI_NOR_E_NOT_FOUND;
    }
    return SPI_NOR_E_UNSUPPORTED;
}

/* Публичные defaults сохраняют тайм-ауты исходного HAL-драйвера. */
spi_nor_config_t spi_nor_default_config(void)
{
    spi_nor_config_t config;

    config.io_timeout_ms = 100U;
    config.lock_timeout_ms = 100U;
    config.program_timeout_ms = 1000U;
    config.erase_timeout_ms = 5000U;
    config.ready_poll_delay_ms = 1U;
    return config;
}

/* Проверяет обязательные callbacks и сохраняет независимую копию контракта. */
spi_nor_status_t spi_nor_init(spi_nor_t *device, const spi_nor_io_t *io,
                              const spi_nor_config_t *config)
{
    if ((device == NULL) || (io == NULL) || (io->transfer == NULL) ||
        (io->chip_select == NULL) || (io->tick_ms == NULL) ||
        (io->delay_ms == NULL) ||
        ((io->lock == NULL) != (io->unlock == NULL))) {
        return SPI_NOR_E_ARGUMENT;
    }
    memset(device, 0, sizeof(*device));
    device->io = *io;
    device->config = (config != NULL) ? *config : spi_nor_default_config();
    if ((device->config.io_timeout_ms == 0U) ||
        (device->config.program_timeout_ms == 0U) ||
        (device->config.erase_timeout_ms == 0U)) {
        memset(device, 0, sizeof(*device));
        return SPI_NOR_E_ARGUMENT;
    }
    device->io.chip_select(device->io.context, 0U);
    device->initialized = 1U;
    return SPI_NOR_OK;
}

/* Будит устройство и сохраняет raw JEDEC даже при неизвестной модели. */
spi_nor_status_t spi_nor_probe(spi_nor_t *device)
{
    uint8_t command = SPI_NOR_COMMAND_RELEASE_POWER_DOWN;
    uint8_t jedec[3] = {0U, 0U, 0U};
    spi_nor_status_t result;

    result = begin_operation(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    device->detected = 0U;
    memset(&device->info, 0, sizeof(device->info));
    result = transaction(device, &command, 1U, NULL, NULL, 0U);
    if (result == SPI_NOR_OK) {
        device->io.delay_ms(device->io.context, 1U);
        command = SPI_NOR_COMMAND_READ_JEDEC_ID;
        result = transaction(device, &command, 1U, NULL, jedec,
                             sizeof(jedec));
    }
    if (result == SPI_NOR_OK) {
        device->info.manufacturer_id = jedec[0];
        device->info.memory_type = jedec[1];
        device->info.capacity_code = jedec[2];
        result = identify_device(device);
        if (result == SPI_NOR_OK) {
            device->detected = 1U;
        }
    }
    end_operation(device);
    return result;
}

/* Чтение проверяет адрес до активации CS и не использует внутренний RAM-кэш. */
spi_nor_status_t spi_nor_read(spi_nor_t *device, uint32_t address,
                              void *data, size_t size)
{
    spi_nor_status_t result;

    if (data == NULL) {
        return SPI_NOR_E_ARGUMENT;
    }
    result = begin_operation(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    result = validate_range(device, address, size);
    if (result == SPI_NOR_OK) {
        result = raw_read(device, address, data, size);
    }
    end_operation(device);
    return result;
}

/* Запись разбивается по границам страниц; SST получает byte-program chunks. */
spi_nor_status_t spi_nor_program(spi_nor_t *device, uint32_t address,
                                 const void *data, size_t size)
{
    const uint8_t *source = data;
    size_t remaining = size;
    spi_nor_status_t result;

    if (data == NULL) {
        return SPI_NOR_E_ARGUMENT;
    }
    result = begin_operation(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    result = validate_range(device, address, size);
    while ((result == SPI_NOR_OK) && (remaining != 0U)) {
        size_t page_remaining = device->info.page_size -
            (address % device->info.page_size);
        size_t chunk = (remaining < page_remaining) ?
            remaining : page_remaining;

        result = program_chunk(device, address, source, chunk);
        address += (uint32_t)chunk;
        source += chunk;
        remaining -= chunk;
    }
    end_operation(device);
    return result;
}

/* Erase принимает только целые выровненные 4-КиБ секторы в пределах памяти. */
spi_nor_status_t spi_nor_erase(spi_nor_t *device, uint32_t address,
                               size_t size)
{
    size_t remaining = size;
    spi_nor_status_t result;

    result = begin_operation(device);
    if (result != SPI_NOR_OK) {
        return result;
    }
    if (((address % SPI_NOR_SECTOR_SIZE) != 0U) ||
        ((size % SPI_NOR_SECTOR_SIZE) != 0U)) {
        result = SPI_NOR_E_ALIGNMENT;
    } else {
        result = validate_range(device, address, size);
    }
    while ((result == SPI_NOR_OK) && (remaining != 0U)) {
        uint8_t command[4];

        result = write_enable(device);
        if (result == SPI_NOR_OK) {
            make_address_command(command, SPI_NOR_COMMAND_SECTOR_ERASE,
                                 address);
            result = transaction(device, command, sizeof(command), NULL,
                                 NULL, 0U);
        }
        if (result == SPI_NOR_OK) {
            result = wait_ready(device, device->config.erase_timeout_ms);
        }
        if (result == SPI_NOR_OK) {
            result = verify_erased_sector(device, address);
        }
        address += SPI_NOR_SECTOR_SIZE;
        remaining -= SPI_NOR_SECTOR_SIZE;
    }
    end_operation(device);
    return result;
}

/* Копия info не позволяет внешнему коду менять геометрию активного контекста. */
spi_nor_status_t spi_nor_get_info(const spi_nor_t *device,
                                  spi_nor_info_t *info)
{
    if ((device == NULL) || (info == NULL)) {
        return SPI_NOR_E_ARGUMENT;
    }
    if (device->detected == 0U) {
        return SPI_NOR_E_NOT_FOUND;
    }
    *info = device->info;
    return SPI_NOR_OK;
}

/* Raw JEDEC остаётся диагностически доступен после NOT_FOUND/UNSUPPORTED. */
spi_nor_status_t spi_nor_get_last_jedec(const spi_nor_t *device,
                                        uint8_t jedec[3])
{
    if ((device == NULL) || (jedec == NULL) ||
        (device->initialized == 0U)) {
        return SPI_NOR_E_ARGUMENT;
    }
    jedec[0] = device->info.manufacturer_id;
    jedec[1] = device->info.memory_type;
    jedec[2] = device->info.capacity_code;
    return SPI_NOR_OK;
}