// Copyright 2020-2023 XMOS LIMITED.
// This Software is subject to the terms of the XMOS Public Licence: Version 1.

#define DEBUG_UNIT RTOS_QSPI_FLASH

#include <string.h>
#include <stdbool.h>
#include <stdio.h>

#include <xcore/assert.h>
#include <xcore/lock.h>

#include "rtos_qspi_flash.h"

#ifndef QSPI_FL_RETRY_DELAY_TICKS
#define QSPI_FL_RETRY_DELAY_TICKS 1000
#endif

#ifndef QSPI_FL_RETRY_ATTEMPT_CNT
#define QSPI_FL_RETRY_ATTEMPT_CNT 5
#endif

/* TODO, these will be removed once moved to the public API */
#define ERASE_CHIP 0xC7

extern void fl_int_read(
        unsigned char cmd, 
        unsigned int address, 
        unsigned char * destination, 
        unsigned int num_bytes);
extern void fl_int_write(
        unsigned char cmd,
        unsigned int pageAddress, 
        const unsigned char data[],
        unsigned int num_bytes);
extern void fl_int_sendSingleByteCommand(unsigned char cmd);
extern void fl_int_eraseSector(
        unsigned char cmd,
        unsigned int sectorAddress);
/* end TODO */

/* Library only supports 4096 sector size*/
#define QSPI_ERASE_TYPE_SIZE_LOG2   12

#define MIN(a,b) ((a) < (b) ? (a) : (b))

#define FLASH_OP_NONE  0
#define FLASH_OP_READ  1
#define FLASH_OP_WRITE 2
#define FLASH_OP_ERASE 3
#define FLASH_OP_READ_FAST_RAW 4
#define FLASH_OP_READ_FAST_NIBBLE_SWAP 5
#define FLASH_OP_LL_SETUP 6

extern unsigned __libc_hwlock;

typedef struct {
    int op;
    uint8_t *data;
    unsigned address;
    size_t len;
    unsigned priority;
} qspi_flash_op_req_t;

/*
 * Returns true if the spinlock is
 * acquired, false if not available.
 * NOT recursive.
 */
static bool spinlock_get(volatile int *lock)
{
    bool ret;

    lock_acquire(__libc_hwlock);
    {
        if (*lock == 0) {
            *lock = 1;
            ret = true;
        } else {
            ret = false;
        }
    }
    lock_release(__libc_hwlock);

    return ret;
}

/*
 * Releases the lock. It MUST be owned
 * by the caller.
 */
static void spinlock_release(volatile int *lock)
{
    *lock = 0;
}

static void rtos_qspi_flash_fl_connect_with_retry(
        rtos_qspi_flash_t *ctx)
{
    if (fl_connectToDevice(&ctx->qspi_ports, &ctx->qspi_spec, 1) == 0) {
        return;
    } else {
        int cnt = 0;
        while (fl_connectToDevice(&ctx->qspi_ports, &ctx->qspi_spec, 1) != 0) {
            delay_ticks(QSPI_FL_RETRY_DELAY_TICKS);
            if (cnt++ >= QSPI_FL_RETRY_ATTEMPT_CNT) {
                xassert(0); /* fl_connectToDevice failed too many times */
            }
        }
    }
}

int rtos_qspi_flash_read_ll(
        rtos_qspi_flash_t *ctx,
        uint8_t *data,
        unsigned address,
        size_t len)
{
    uint32_t irq_mask;
    bool lock_acquired;

//    printf("Asked to ll read %d bytes at address 0x%08x\n", len, address);

    lock_acquired = spinlock_get(&ctx->spinlock);

    while (lock_acquired && len > 0) {

        size_t read_len = MIN(len, RTOS_QSPI_FLASH_READ_CHUNK_SIZE);

        /*
         * Cap the address at the size of the flash.
         * This ensures the correction below will work if
         * address is outside the flash's address space.
         */
        if (address >= ctx->flash_size) {
            address = ctx->flash_size;
        }

        if (address + read_len > ctx->flash_size) {
            int original_len = read_len;

            /* Don't read past the end of the flash */
            read_len = ctx->flash_size - address;

            /* Return all 0xFF bytes for addresses beyond the end of the flash */
            memset(&data[read_len], 0xFF, original_len - read_len);
        }

    //    printf("Read %d bytes from flash at address 0x%x\n", read_len, address);

        fl_int_read(ctx->qspi_spec.readCommand, address, data, read_len);

        len -= read_len;
        data += read_len;
        address += read_len;
    }

    if (lock_acquired) {
        spinlock_release(&ctx->spinlock);
    }

    return lock_acquired ? 0 : -1;
}


static void read_op(
        rtos_qspi_flash_t *ctx,
        uint8_t *data,
        unsigned address,
        size_t len)
{
//    printf("Asked to read %d bytes at address 0x%08x\n", len, address);

    while (len > 0) {

        size_t read_len = MIN(len, RTOS_QSPI_FLASH_READ_CHUNK_SIZE);

        /*
         * Cap the address at the size of the flash.
         * This ensures the correction below will work if
         * address is outside the flash's address space.
         */
        if (address >= ctx->flash_size) {
            address = ctx->flash_size;
        }

        if (address + read_len > ctx->flash_size) {
            int original_len = read_len;

            /* Don't read past the end of the flash */
            read_len = ctx->flash_size - address;

            /* Return all 0xFF bytes for addresses beyond the end of the flash */
            memset(&data[read_len], 0xFF, original_len - read_len);
        }

     //   printf("Read %d bytes from flash at address 0x%x\n", read_len, address);

        fl_int_read(ctx->qspi_spec.readCommand, address, data, read_len);

        len -= read_len;
        data += read_len;
        address += read_len;
    }
}


static void while_busy(void)
{
    bool busy;

    do {
        busy = fl_getBusyStatus();
    } while (busy);
}

static void write_op(
        rtos_qspi_flash_t *ctx,
        const uint8_t *data,
        unsigned address,
        size_t len)
{
    size_t bytes_left_to_write = len;
    unsigned address_to_write = address;
    const uint8_t *write_buf = data;

   // printf("Asked to write %d bytes at address 0x%08x\n", bytes_left_to_write, address_to_write);

    while (bytes_left_to_write > 0) {
        /* compute the maximum number of bytes that can be written to the current page. */
        size_t max_bytes_to_write = fl_getPageSize() - (address_to_write & (fl_getPageSize() - 1));
        size_t bytes_to_write = bytes_left_to_write <= max_bytes_to_write ? bytes_left_to_write : max_bytes_to_write;

        if (address_to_write >= ctx->flash_size) {
            break; /* do not write past the end of the flash */
        }

     //   printf("Write %d bytes from flash at address 0x%x\n", bytes_to_write, address_to_write);
            fl_int_sendSingleByteCommand(ctx->qspi_spec.writeEnableCommand);
            fl_int_write(ctx->qspi_spec.programPageCommand, address_to_write, write_buf, bytes_to_write);
        while_busy();
            fl_int_sendSingleByteCommand(ctx->qspi_spec.writeDisableCommand);

        bytes_left_to_write -= bytes_to_write;
        write_buf += bytes_to_write;
        address_to_write += bytes_to_write;
    }
}

#define SECTORS_TO_BYTES(s, ss_log2) ((s) << (ss_log2))
#define BYTES_TO_SECTORS(b, ss_log2) (((b) + (1 << ss_log2) - 1) >> (ss_log2))

#define SECTOR_TO_BYTE_ADDRESS(s, ss_log2) SECTORS_TO_BYTES(s, ss_log2)
#define BYTE_TO_SECTOR_ADDRESS(b, ss_log2) ((b) >> (ss_log2))

static void erase_op(
        rtos_qspi_flash_t *ctx,
        unsigned address,
        size_t len)
{
    size_t bytes_left_to_erase = len;
    unsigned address_to_erase = address;

//    printf("Asked to erase %d bytes at address 0x%08x\n", bytes_left_to_erase, address_to_erase);

    if (address_to_erase == 0 && bytes_left_to_erase >= ctx->flash_size) {
        /* Use chip erase when being asked to erase the entire address range */
 //       printf("Erasing entire chip\n");
            fl_int_sendSingleByteCommand(ctx->qspi_spec.writeEnableCommand);
            fl_int_sendSingleByteCommand(ERASE_CHIP);
        while_busy();
            fl_int_sendSingleByteCommand(ctx->qspi_spec.writeDisableCommand);
    } else {
        if (SECTOR_TO_BYTE_ADDRESS(BYTE_TO_SECTOR_ADDRESS(address_to_erase, QSPI_ERASE_TYPE_SIZE_LOG2), QSPI_ERASE_TYPE_SIZE_LOG2) != address_to_erase) {
            /*
             * If the provided starting erase address does not begin on the smallest
             * sector boundary, then update the starting address and number of bytes
             * to erase so that it does.
             */
            unsigned sector_address;
            sector_address = BYTE_TO_SECTOR_ADDRESS(address_to_erase, QSPI_ERASE_TYPE_SIZE_LOG2);
            bytes_left_to_erase += address_to_erase - SECTOR_TO_BYTE_ADDRESS(sector_address, QSPI_ERASE_TYPE_SIZE_LOG2);
            address_to_erase = SECTOR_TO_BYTE_ADDRESS(sector_address, QSPI_ERASE_TYPE_SIZE_LOG2);
         //   printf("adjusted starting erase address to %d\n", address_to_erase);
        }

        while (bytes_left_to_erase > 0) {
            int erase_length;
            int erase_length_log2 = QSPI_ERASE_TYPE_SIZE_LOG2;

            if (address_to_erase >= ctx->flash_size) {
                break; /* do not erase past the end of the flash */
            }

            erase_length = 1 << erase_length_log2;

            xassert(address_to_erase == SECTOR_TO_BYTE_ADDRESS(BYTE_TO_SECTOR_ADDRESS(address_to_erase, erase_length_log2), erase_length_log2));

         //   printf("Erasing %d bytes (%d) at byte address %d, sector %d\n", erase_length, bytes_left_to_erase, address_to_erase, BYTE_TO_SECTOR_ADDRESS(address_to_erase, erase_length_log2));

                fl_int_sendSingleByteCommand(ctx->qspi_spec.writeEnableCommand);
                fl_int_eraseSector(ctx->qspi_spec.sectorEraseCommand, address_to_erase);
            while_busy();
                fl_int_sendSingleByteCommand(ctx->qspi_spec.writeDisableCommand);

            address_to_erase += erase_length;
            bytes_left_to_erase -= erase_length < bytes_left_to_erase ? erase_length : bytes_left_to_erase;
        }
    }

//    printf("Erasing complete\n");
}

__attribute__((fptrgroup("rtos_qspi_flash_read_fptr_grp")))
static void qspi_flash_local_read(
        rtos_qspi_flash_t *ctx,
        uint8_t *data,
        unsigned address,
        size_t len)
{
    qspi_flash_op_req_t op = {
            .op = FLASH_OP_READ,
            .data = data,
            .address = address,
            .len = len
    };
    read_op(ctx, op.data, op.address, op.len);

}

__attribute__((fptrgroup("rtos_qspi_flash_write_fptr_grp")))
static void qspi_flash_local_write(
        rtos_qspi_flash_t *ctx,
        const uint8_t *data,
        unsigned address,
        size_t len)
{
    qspi_flash_op_req_t op = {
            .op = FLASH_OP_WRITE,
            .address = address,
            .len = len
    };

    write_op(ctx, data, op.address, op.len);
}

__attribute__((fptrgroup("rtos_qspi_flash_erase_fptr_grp")))
static void qspi_flash_local_erase(
        rtos_qspi_flash_t *ctx,
        unsigned address,
        size_t len)
{
    qspi_flash_op_req_t op = {
            .op = FLASH_OP_ERASE,
            .address = address,
            .len = len
    };
    erase_op(ctx, op.address, op.len);
}

void rtos_qspi_flash_init(
        rtos_qspi_flash_t *ctx,
        xclock_t clock_block,
        port_t cs_port,
        port_t sclk_port,
        port_t sio_port,
        fl_QuadDeviceSpec *spec)
{
    ctx->qspi_ports.qspiCS = cs_port;
    ctx->qspi_ports.qspiSCLK = sclk_port;
    ctx->qspi_ports.qspiSIO = sio_port;
    ctx->qspi_ports.qspiClkblk = clock_block;

    fl_QuadDeviceSpec default_spec = FL_QUADDEVICE_DEFAULT;

    if (spec == NULL) {
        memcpy(&ctx->qspi_spec, &default_spec, sizeof(fl_QuadDeviceSpec));
    } else {
        memcpy(&ctx->qspi_spec, spec, sizeof(fl_QuadDeviceSpec));
    }

    xassert(fl_connectToDevice(&ctx->qspi_ports, &ctx->qspi_spec, 1) == 0);
    /* Copy the spec back in case one was provided which has params populated by sfdp */
    xassert(fl_copySpec(&ctx->qspi_spec) == 0);

    ctx->flash_size = fl_getFlashSize();

    /* Driver currently only supports 4096 sector size */
    xassert(rtos_qspi_flash_sector_size_get(ctx) == (1 << QSPI_ERASE_TYPE_SIZE_LOG2));

    /* Enable quad flash */
    xassert(fl_quadEnable() == 0);

    ctx->calibration_valid = 0;
    ctx->last_op = FLASH_OP_NONE;
    ctx->read = qspi_flash_local_read;
    ctx->write = qspi_flash_local_write;
    ctx->erase = qspi_flash_local_erase;
}
