/**
 *******************************************************************************
 * @file    tiny_kvdb.c
 * @brief   tiny key-value database source file.
 *******************************************************************************
 */

/* Includes ------------------------------------------------------------------*/
#include <stddef.h>
#include <string.h>
#include "tiny_kvdb.h"

/* Private Defines -----------------------------------------------------------*/
#define TINY_KVDB_MAGIC         0x54494B56  /* "TIKV" */
#define TINY_KVDB_ERASED        0xFFFFFFFF

/***************** DO NOT CHANGE !!!!! *********************/

/* 
   page tag state machine:
   erase->swap->gc->data->del->erase OR
   erase->data->del->erase
*/
#define PAGE_TAG_ERASED         0xFFFFFFFF      /* Page is just erased */
#define PAGE_TAG_SWAP           0x5041DAFF      /* Page is reserved for garbage collection. */
#define PAGE_TAG_GC             0x5041DAF0      /* Page is doing gargage collection. */
#define PAGE_TAG_DATA           0x5041DAC0      /* Page is ready for storage. */
#define PAGE_TAG_DEL            0x5041DA00      /* Page is to be deleted. */

/* record constants */
#define RECORD_KEY_UNUSED       0xFFFF      /* Record key, value after sector erase. */
#define RECORD_KEY_DELETED      0x0000      /* Record key, value after deleted. */
#define RECORD_LEN_UNUSED       0xFFFF      /* Record length, value after sector erase. */
/* record status */
#define RECORD_STATUS_UNUSED    0xFFFF
#define RECORD_STATUS_OK        0xFFF0      /* Record is ok. */
#define RECORD_STATUS_PRE_DEL   0xFF00      /* Record is to be deleted. */
#define RECORD_STATUS_DELETED   0x0000      /* Record is deleted. */

/* Private TypesDefs ---------------------------------------------------------*/
/* 8-byte aligned */
typedef struct
{
    uint32_t        magic;      /* page magic */
    uint32_t        tag;        /* page tag */
} page_header_t;

/* 8-byte aligned */
typedef struct
{
    uint16_t        key;        /* record key */
    uint16_t        len;        /* record length */
    uint16_t        status;     /* record status */
    uint16_t        check;      /* check of the header, for future use */
} record_header_t;

typedef struct 
{
    uint32_t        device_id;      /* device id */
    uint32_t        start_addr;     /* database start addr */
    uint32_t        end_addr;       /* database end addr */
} user_info_t;


/* Private Macros ------------------------------------------------------------*/
#define IS_POWER_OF_TWO(a)      ( ((a) > 1) && ((((a) - 1) & (a)) == 0) )
#define MIN(a, b)               (((a) < (b)) ? (a) : (b))

/* Private Variables ---------------------------------------------------------*/
static user_info_t m_users[TINY_KVDB_MAX_USERS];

/* Extern Variables ----------------------------------------------------------*/

/* Private FunctionPrototypes ------------------------------------------------*/

static int do_gc(tiny_kvdb_t *db, bool is_power_on);

/* Extern FunctionPrototypes -------------------------------------------------*/

/* Private Functions ---------------------------------------------------------*/

static int set_page_tag(tiny_kvdb_t *db, uint32_t addr, uint32_t tag)
{
    page_header_t page_tag = {TINY_KVDB_MAGIC, PAGE_TAG_ERASED};
    switch (tag)
    {
        case PAGE_TAG_SWAP:
        case PAGE_TAG_GC:
        case PAGE_TAG_DATA:
        case PAGE_TAG_DEL:
        {
            page_tag.tag = tag;
            break;
        }
        default:
        {
            return -1;
        }
    }

    return db->flash_param.write(addr, &page_tag, sizeof(page_header_t));
}

static inline uint32_t aligned_size(tiny_kvdb_t *db, uint32_t len)
{
    uint8_t size = db->flash_param.write_unit;

    if (size <= 1U) 
    {
        return len;
    }
    return (len + (size - 1U)) & ~(size - 1U);
}

#if TINY_KVDB_USE_CACHE
static void clear_cache(tiny_kvdb_t *db)
{
    uint32_t i;

    for (i = 0; i < TINY_KVDB_CACHE_NUM; i++)
    {
        db->cache[i].key = RECORD_KEY_UNUSED;
        db->cache[i].len = 0;
        db->cache[i].addr = 0;
    }
}
#endif

static int flag_record_pre_deleted(tiny_kvdb_t *db, uint32_t record_addr)
{
    record_header_t header;
    int ret = TINY_DB_ERR_INVALID_PARAM;

    if (db->flash_param.write_unit <= 2)
    {
        header.status = RECORD_STATUS_PRE_DEL;
        ret = db->flash_param.write(record_addr + offsetof(record_header_t, status), &header.status, 2);
    }
    else if (db->flash_param.write_unit == 4)
    {
        header.status = RECORD_STATUS_PRE_DEL;
        header.check = (uint16_t)TINY_KVDB_ERASED;
        ret = db->flash_param.write(record_addr + offsetof(record_header_t, status), &header.status, 4);
    }
    else if (db->flash_param.write_unit == 8)
    {
        ret = db->flash_param.read(record_addr, &header, 8);
        if (ret < 0)
        {
            return ret;
        }

        header.status = RECORD_STATUS_PRE_DEL;
        ret = db->flash_param.write(record_addr, &header, 8);
    }

    return ret;
}

static int flag_record_deleted(tiny_kvdb_t *db, uint32_t record_addr)
{
    record_header_t header;
    int ret = TINY_DB_ERR_INVALID_PARAM;

    if (db->flash_param.write_unit <= 2)
    {
        header.key = RECORD_KEY_DELETED;
        ret = db->flash_param.write(record_addr, &header, 2);
    }
    else if (db->flash_param.write_unit == 4)
    {
        ret = db->flash_param.read(record_addr, &header, 4);
        if (ret < 0)
        {
            return ret;
        }
        header.key = RECORD_KEY_DELETED;
        ret = db->flash_param.write(record_addr, &header, 4);
    }
    else if (db->flash_param.write_unit == 8)
    {
        ret = db->flash_param.read(record_addr, &header, 8);
        if (ret < 0)
        {
            return ret;
        }

        header.key = RECORD_KEY_DELETED;
        ret = db->flash_param.write(record_addr, &header, 8);
    }

    return ret;
}

static int delete_record(tiny_kvdb_t *db, record_info_t *record, bool direct_delete)
{
    int ret;

    if (record->key == RECORD_KEY_UNUSED || record->key == RECORD_KEY_DELETED)
    {
        return TINY_DB_ERR_NONE;
    }

    if (direct_delete)
    {
        ret = flag_record_deleted(db, record->addr);
        if (ret < 0)
        {
            return ret;
        }

        db->dirty_size += sizeof(record_header_t) + aligned_size(db, record->len);

#if TINY_KVDB_USE_CACHE
        for (uint32_t i = 0; i < TINY_KVDB_CACHE_NUM; i++)
        {
            if (db->cache[i].key == record->key)
            {
                db->cache[i].key = RECORD_KEY_DELETED;
                db->cache[i].len = 0;
                db->cache[i].addr = 0;
                break;
            }
        }
#endif
    }
    else
    {
        ret = flag_record_pre_deleted(db, record->addr);
        if (ret < 0)
        {
            return ret;
        }
    }

    return TINY_DB_ERR_NONE;
}

/**
 * @brief  Traverse the database, calculate dirty size and next write address. 
 *         If use cache, put record infomations into the cache.
 * @param  [in]db Pointer to the database
 * @return TINY_DB_ERR_NONE if success, else error code.
 */
static int iterate_and_check(tiny_kvdb_t *db)
{
    int ret;
    record_header_t header;
    uint32_t addr = db->data_addr + sizeof(page_header_t);

    /* Variables to check incomplete deleted record, CAN ONLY occur when system
       power off, and AT MOST one record */
    bool have_pre_deleted_record = false;
    record_info_t pre_deleted_record;     

#if TINY_KVDB_USE_CACHE
    uint32_t i = 0;
    clear_cache(db);
#endif

    db->dirty_size = 0;
    while (addr < db->data_addr + db->total_size - sizeof(record_header_t))
    {
        ret = db->flash_param.read(addr, &header, sizeof(record_header_t));
        if (ret < 0)
        {
            return ret;
        }
        TINY_KVDB_DEBUG("record_key:0x%04X, record_len:%d\r\n", header.key, header.len);

        if (header.key == RECORD_KEY_UNUSED)
        {
            /* data error, assume that all the rest is dirty data, return */
            if (header.len != RECORD_LEN_UNUSED)
            {
                db->dirty_size += db->data_addr + db->total_size - addr;
                db->next_addr = db->data_addr + db->total_size - 1;
                return TINY_DB_ERR_NONE;
            }
            
            /* record key and length are all 0xFFFF, should be the next write address */
            db->next_addr = addr;
            return TINY_DB_ERR_NONE;
        }
        else
        {
            /* data length error, assume that all the rest is dirty data, return */
            if (header.len == 0 || addr + sizeof(record_header_t) + header.len > db->data_addr + db->total_size)
            {
                db->dirty_size += db->data_addr + db->total_size - addr;
                db->next_addr = db->data_addr + db->total_size - 1;
                return TINY_DB_ERR_NONE;
            }
            
            /* record has been deleted */
            if (header.key == RECORD_KEY_DELETED)
            {
                db->dirty_size += sizeof(record_header_t) + aligned_size(db, header.len);
            }
            /* record status is invalid, so the record may be incomplete */
            else if (header.status != RECORD_STATUS_OK && header.status != RECORD_STATUS_PRE_DEL)
            {
                db->dirty_size += sizeof(record_header_t) + aligned_size(db, header.len);
            }
            /* record key and status are all valid */
            else
            {
                if (have_pre_deleted_record)
                {
                    /* the pre_deleted_record is updated by new value, delete it */
                    if (pre_deleted_record.key == header.key)
                    {
                        ret = flag_record_deleted(db, pre_deleted_record.addr);
                        if (ret < 0)
                        {
                            return ret;
                        }
                        have_pre_deleted_record = false;
                        db->dirty_size += sizeof(record_header_t) + aligned_size(db, pre_deleted_record.len);
                    }
                }

                if (header.status == RECORD_STATUS_PRE_DEL)
                {
                    pre_deleted_record.key = header.key;
                    pre_deleted_record.len = header.len;
                    pre_deleted_record.addr = addr;
                    have_pre_deleted_record = true;
                }

#if TINY_KVDB_USE_CACHE
                /* record status ok, put into the cache */
                else
                {
                    if (i < TINY_KVDB_CACHE_NUM)
                    {
                        db->cache[i].key = header.key;
                        db->cache[i].len = header.len;
                        db->cache[i].addr = addr;
                        i++;
                    }
                }
#endif
            }

            addr += sizeof(record_header_t) + aligned_size(db, header.len);
        }
    }

#if TINY_KVDB_USE_CACHE
    /* the pre_deleted_record is not updated, recover it */
    if (have_pre_deleted_record)
    {
        if (i < TINY_KVDB_CACHE_NUM)
        {
            db->cache[i].key = pre_deleted_record.key;
            db->cache[i].len = pre_deleted_record.len;
            db->cache[i].addr = pre_deleted_record.addr;
        }
    }
#endif

    return TINY_DB_ERR_NONE;
}

/**
 * @brief  find record address in the flash.
 * @param  [in]db pointer to database
 * @param  [in]key record key
 * @param  [out]p_recd pointer to save record infomation 
 * @return if found, return TINY_DB_ERR_NONE, else return error code. 
 */
static int find_record(tiny_kvdb_t *db, uint16_t key, record_info_t *p_recd)
{
#if TINY_KVDB_USE_CACHE
    uint32_t i;
    for (i = 0; i < TINY_KVDB_CACHE_NUM; i++)
    {
        if (db->cache[i].key == RECORD_KEY_UNUSED)
        {
            return TINY_DB_ERR_NOT_FOUND;
        }
        if (db->cache[i].key == key)
        {
            p_recd->key = key;
            p_recd->len = db->cache[i].len;
            p_recd->addr = db->cache[i].addr;
            return TINY_DB_ERR_NONE;
        }
    }
#endif

    int ret;
    record_header_t header;
    uint32_t addr = db->data_addr + sizeof(page_header_t);
    while (addr < db->data_addr + db->total_size - sizeof(record_header_t))
    {
        ret = db->flash_param.read(addr, &header, sizeof(record_header_t));
        if (ret < 0)
        {
            return ret;
        }

        /* reach the end of database, return false */
        if (header.key == RECORD_KEY_UNUSED)
        {
            return TINY_DB_ERR_NOT_FOUND;
        }
        else if (header.key == key && (header.status == RECORD_STATUS_OK || header.status == RECORD_STATUS_PRE_DEL))
        {
            /* data error */
            if (header.len == 0 || addr + sizeof(record_header_t) + header.len > db->data_addr + db->total_size)
            {
                return TINY_DB_ERR_NOT_FOUND;
            }
            
            /* record found */
#if TINY_KVDB_USE_CACHE

            /* put the record to the first place of the cache. */
            for (i = TINY_KVDB_CACHE_NUM - 1; i != 0; i--)
            {
                db->cache[i].key = db->cache[i - 1].key;
                db->cache[i].len = db->cache[i - 1].len;
                db->cache[i].addr = db->cache[i - 1].addr;
            }
            db->cache[0].key = header.key;
            db->cache[0].len = header.len;
            db->cache[0].addr = addr;
            
#endif

            p_recd->key = key;
            p_recd->len = header.len;
            p_recd->addr = addr;
            return TINY_DB_ERR_NONE;
        }
        else
        {
            /* data error */
            if (header.len == 0 || addr + sizeof(record_header_t) + header.len > db->data_addr + db->total_size)
            {
                return TINY_DB_ERR_NOT_FOUND;
            }
            
            addr += sizeof(record_header_t) + aligned_size(db, header.len);
        }
    }

    return TINY_DB_ERR_NOT_FOUND;
}

/*  | ! | !     or   ! | ! |     or   | ! ! |   or  ! | | !   */
static inline int space_overlap(uint32_t start_addr1, uint32_t end_addr1, uint32_t start_addr2, uint32_t end_addr2)
{
    return ((start_addr1 >= start_addr2 && start_addr1 <= end_addr2) || 
            (end_addr1 >= start_addr2 && end_addr1 <= end_addr2) ||
            (start_addr2 >= start_addr1 && start_addr2 <= end_addr1));
}

static uint8_t check_status(page_header_t *header1, page_header_t *header2)
{
    uint8_t status;
    if (header1->magic == TINY_KVDB_ERASED)
    {
        status = 0;
    }
    else if (header1->magic == TINY_KVDB_MAGIC)
    {
        if (header1->tag == PAGE_TAG_SWAP)
        {
            status = 0x10;
        }
        else if (header1->tag == PAGE_TAG_GC)
        {
            status = 0x20;
        }
        else if (header1->tag == PAGE_TAG_DATA)
        {
            status = 0x30;
        }
        else if (header1->tag == PAGE_TAG_DEL)
        {
            status = 0x40;
        }
        else
        {
            status = 0x50;
        }
    }
    else
    {
        status = 0x50;
    }

    if (header2->magic == TINY_KVDB_ERASED)
    {
        status |= 0;
    }
    else if (header2->magic == TINY_KVDB_MAGIC)
    {
        if (header2->tag == PAGE_TAG_SWAP)
        {
            status |= 0x01;
        }
        else if (header2->tag == PAGE_TAG_GC)
        {
            status |= 0x02;
        }
        else if (header2->tag == PAGE_TAG_DATA)
        {
            status |= 0x03;
        }
        else if (header2->tag == PAGE_TAG_DEL)
        {
            status |= 0x04;
        }
        else
        {
            status |= 0x05;
        }
    }
    else
    {
        status |= 0x05;
    }

    return status;
}

int tiny_kvdb_init(tiny_kvdb_t *db, uint32_t start_addr, uint32_t sector_count, flash_param_t *flash_param)
{
    /* Parameters validation */
    if (!IS_POWER_OF_TWO(flash_param->sector_size) || flash_param->sector_size < 16)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }
    if (flash_param->write_unit != 1 && flash_param->write_unit != 2 && flash_param->write_unit != 4 && flash_param->write_unit != 8)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }
    if (flash_param->read == NULL || flash_param->write == NULL || flash_param->erase == NULL)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }
    if ((start_addr & (flash_param->sector_size - 1)) != 0)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }
    if ((sector_count & 0x01) != 0)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }


    uint32_t end_addr = start_addr + flash_param->sector_size * sector_count - 1; 
    uint32_t first_addr = start_addr;
    uint32_t second_addr = first_addr + flash_param->sector_size * sector_count / 2;

    /* Database has been initialized, check the parameters */
    if (db->init_ok)
    {
        if ((db->data_addr == first_addr && db->swap_addr == second_addr) ||
           (db->data_addr == second_addr && db->swap_addr == first_addr))
        {
            if (memcmp(&db->flash_param, flash_param, sizeof(flash_param_t)) != 0)
            {
                return TINY_DB_ERR_IN_USE;
            }
            return TINY_DB_ERR_NONE;
        }
        else
        {
            return TINY_DB_ERR_IN_USE;
        }
    }

    /* Check the max instances */
    uint32_t i, user_cnt = 0;
    for ( i = 0; i < TINY_KVDB_MAX_USERS; i++)
    {
        if (m_users[i].end_addr != 0)
        {
            user_cnt++;
            /* Storage space overlap */
            if (m_users[i].device_id == flash_param->device_id && 
                space_overlap(start_addr, end_addr, m_users[i].start_addr, m_users[i].end_addr))
            {
                return TINY_DB_ERR_IN_USE;
            }
        }
    }
    if (user_cnt == TINY_KVDB_MAX_USERS)
    {
        return TINY_DB_ERR_NO_SPACE;
    }
    /* Add to m_users */
    for (i = 0; i < TINY_KVDB_MAX_USERS; i++)
    {
        if (m_users[i].end_addr == 0)
        {
            m_users[i].device_id = flash_param->device_id;
            m_users[i].start_addr = start_addr;
            m_users[i].end_addr = end_addr;
        }
    }

    db->total_size = second_addr - first_addr;
    memcpy(&db->flash_param, flash_param, sizeof(flash_param_t));
    if (flash_param->init != NULL)
    {
        flash_param->init();
    }

    page_header_t header1, header2;

    if (flash_param->read(first_addr, &header1, sizeof(page_header_t)) < 0)
    {
        return TINY_DB_ERR_READ;
    }
    if (flash_param->read(second_addr, &header2, sizeof(page_header_t)) < 0)
    {
        return TINY_DB_ERR_READ;
    }

    uint8_t status = check_status(&header1, &header2);
    int ret = TINY_DB_ERR_NONE;
    switch (status)
    {
        /* sectors are not formatted and just erased, only need to set page tags. */
        case 0x00:
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            if (set_page_tag(db, db->data_addr, PAGE_TAG_DATA) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            db->next_addr = db->data_addr + sizeof(page_header_t);
            db->dirty_size = 0;

#if TINY_KVDB_USE_CACHE
            clear_cache(db);
#endif
            break;
        }

        /* one data page and one erased page, gc is almost done. */
        case 0x03:
        {
            db->data_addr = second_addr;
            db->swap_addr = first_addr;
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            ret = iterate_and_check(db);
            break;
        }
        case 0x30:
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            ret = iterate_and_check(db);
            break;
        }

        /* one data page and one swap page, status is OK. */
        case 0x13:
        {
            db->data_addr = second_addr;
            db->swap_addr = first_addr;
            ret = iterate_and_check(db);
            break;
        }
        case 0x31:
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            ret = iterate_and_check(db);
            break;
        }

        /******** is doing garbage collection. ********/
        /* Pre-delete data page, preparing to do GC. */
        case 0x41:          
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            ret = do_gc(db, true);
            break;
        }
        case 0x14:  
        {
            db->data_addr = second_addr;
            db->swap_addr = first_addr;
            ret = do_gc(db, true);
            break;
        }

        /* GC is in progress. */
        case 0x42:  
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            ret = do_gc(db, true);
            break;
        }
        case 0x24:  
        {
            db->data_addr = second_addr;
            db->swap_addr = first_addr;
            ret = do_gc(db, true);
            break;
        }

        /* one data page and one deleted page, GC is almost done. */
        case 0x43:  
        {
            db->data_addr = second_addr;
            db->swap_addr = first_addr;
            if (flash_param->erase(db->swap_addr, db->total_size) < 0)
            {
                return TINY_DB_ERR_ERASE;
            }
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            ret = iterate_and_check(db);
            break;
        }
        case 0x34: 
        {
            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            if (flash_param->erase(db->swap_addr, db->total_size) < 0)
            {
                return TINY_DB_ERR_ERASE;
            }
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            ret = iterate_and_check(db);
            break;
        }
        
        /* invalid status. format the database. */
        default:
        {       
            if (flash_param->erase(first_addr, db->total_size) < 0)
            {
                return TINY_DB_ERR_ERASE;
            }
            if (flash_param->erase(second_addr, db->total_size) < 0)
            {
                return TINY_DB_ERR_ERASE;
            }  

            db->data_addr = first_addr;
            db->swap_addr = second_addr;
            db->next_addr = first_addr + sizeof(page_header_t);
            db->dirty_size = 0;
            if (set_page_tag(db, db->data_addr, PAGE_TAG_DATA) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }
            if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
            {
                return TINY_DB_ERR_WRITE;
            }

#if TINY_KVDB_USE_CACHE
            clear_cache(db);
#endif
            break;
        }
    }

    if (ret != TINY_DB_ERR_NONE)
    {
        return ret;
    }

    TINY_KVDB_INFO("Tiny kvdb init OK.\r\n");
    TINY_KVDB_DEBUG("data_addr:0x%08X, swap_addr:0x%08X, next_addr:0x%08X, dirty_size:%d\r\n", db->data_addr, db->swap_addr, db->next_addr, db->dirty_size);    

#if TINY_KVDB_USE_CACHE
    for (i = 0; i < TINY_KVDB_CACHE_NUM; i++)
    {
        TINY_KVDB_DEBUG("cache %d: key=0x%04X, len=%d, addr=0x%08X\r\n", i, db->cache[i].key, db->cache[i].len, db->cache[i].addr);
    }
#endif    

    db->init_ok = 1;
    return TINY_DB_ERR_NONE;
}
        
int tiny_kvdb_read(tiny_kvdb_t *db, uint16_t key, void *p_dst, uint16_t len)
{
    if (!db->init_ok)
    {
        TINY_KVDB_ERROR("Database not initialized!\r\n");
        return TINY_DB_ERR_NOT_INIT;
    }
    if (key == RECORD_KEY_UNUSED || key == RECORD_KEY_DELETED)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }

    record_info_t recd;
    int ret = find_record(db, key, &recd);
    if (ret != TINY_DB_ERR_NONE)
    {
        TINY_KVDB_INFO("Record key %04x not found!\r\n", key);
        return ret;
    }
    if (p_dst != NULL)
    {
        len = len > recd.len ? recd.len : len;
        return db->flash_param.read(recd.addr + sizeof(record_header_t), p_dst, len);
    }

    return len;
}


int tiny_kvdb_write(tiny_kvdb_t *db, uint16_t key, void *p_src, uint16_t len)
{
    if (!db->init_ok)
    {
        TINY_KVDB_ERROR("Database not initialized!\r\n");
        return TINY_DB_ERR_NOT_INIT;
    }
    if (key == RECORD_KEY_UNUSED || key == RECORD_KEY_DELETED)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }

    int ret;
    bool old_found = false;
    bool need_gc = false;
    record_info_t recd = 
    {
        .key = RECORD_KEY_UNUSED,
        .len = 0,
        .addr = 0
    };
    ret = find_record(db, key, &recd);
    if (ret == TINY_DB_ERR_NONE)
    {
        old_found = true;
    }
    else if (ret != TINY_DB_ERR_NOT_FOUND)
    {
        return ret;
    }

    /* if p_src is NULL, delete record */
    if (p_src == NULL)
    {        
        if (old_found)
        {
            return delete_record(db, &recd, true);       
        }

        return TINY_DB_ERR_NONE;
    }

    if (len == 0)
    {
        return TINY_DB_ERR_INVALID_LENGTH;
    }

    /* data sector do not have enough space */
    if (db->next_addr + sizeof(record_header_t) + aligned_size(db, len) >= db->data_addr + db->total_size)
    {
        /* no enough space to save record */
        if (db->data_addr + db->total_size + db->dirty_size + aligned_size(db, recd.len) - db->next_addr < sizeof(record_header_t) + aligned_size(db, len))
        {
            TINY_KVDB_ERROR("No enough space to save record!\r\n");
            return TINY_DB_ERR_NO_SPACE;
        }

        /* need to run garbage collection */
        need_gc = true;
        ret = tiny_kvdb_gc(db);
        if (ret < 0)
        {
            return ret;
        }
    }

    TINY_KVDB_DEBUG("Writing record... \r\nrecord_key:0x%04X, record_len:%d\r\n", key, len);

    /* pre delete old record */
    if (old_found)
    {
        if (need_gc)
        {
            recd.key = RECORD_KEY_UNUSED;
            ret = find_record(db, key, &recd);
            if (ret != TINY_DB_ERR_NONE && ret != TINY_DB_ERR_NOT_FOUND)
            {
                return ret;
            }
        }
        
        ret = delete_record(db, &recd, false);
        if (ret != TINY_DB_ERR_NONE)
        {
            return ret;
        }        
    }

    record_header_t header = 
    {
        .key = key,
        .len = len,
        .status = 0xFFFF,
        .check = 0xFFFF
    };

    /* write record key and length */
    if (db->flash_param.write_unit <= 4)
    {
        ret = db->flash_param.write(db->next_addr, (uint8_t *)&header, 4);
    }
    else
    {
        ret = db->flash_param.write(db->next_addr, (uint8_t *)&header, 8);
    }
    if (ret < 0)
    {
        return ret;
    }

    /* write record content */
    ret = db->flash_param.write(db->next_addr + sizeof(record_header_t), p_src, len);
    if (ret < 0)
    {
        return ret;
    }

    /* write record status and check */
    header.status = RECORD_STATUS_OK;
    if (db->flash_param.write_unit <= 4)
    {
        ret = db->flash_param.write(db->next_addr + offsetof(record_header_t, status), (uint8_t *)&header.status, 4);
    }
    else
    {
        ret = db->flash_param.write(db->next_addr, (uint8_t *)&header, 8);
    }
    if (ret < 0)
    {
        return ret;
    }

    /* complete delete old record */
    if (old_found)
    {
        ret = delete_record(db, &recd, true);
        if (ret != TINY_DB_ERR_NONE)
        {
            return ret;
        }        
    }

#if TINY_KVDB_USE_CACHE
    for (uint32_t i = 0; i < TINY_KVDB_CACHE_NUM; i++)
    {
        if (db->cache[i].key == RECORD_KEY_DELETED || db->cache[i].key == RECORD_KEY_UNUSED)
        {
            db->cache[i].key = key;
            db->cache[i].len = len;
            db->cache[i].addr = db->next_addr;
            break;
        }
    }
#endif

    db->next_addr += sizeof(record_header_t) + aligned_size(db, len);

    TINY_KVDB_INFO("Record key 0x%04X write done! \r\n", key);
    return len;
}

int tiny_kvdb_delete(tiny_kvdb_t *db, uint16_t key)
{
    if (!db->init_ok)
    {
        return TINY_DB_ERR_NOT_INIT;
    }
    if (key == RECORD_KEY_UNUSED || key == RECORD_KEY_DELETED)
    {
        return TINY_DB_ERR_INVALID_PARAM;
    }

    TINY_KVDB_DEBUG("Deleting record... \r\nrecord_key:0x%04X\r\n", key);

    record_info_t recd;
    int ret = find_record(db, key, &recd);
    if (ret != TINY_DB_ERR_NONE)
    {
        return ret;
    }

    ret = delete_record(db, &recd, true);
    if (ret != TINY_DB_ERR_NONE)
    {
        return ret;
    } 

    TINY_KVDB_INFO("Record key 0x%04X deleted! \r\n", key);
    return TINY_DB_ERR_NONE;
}

static int copy_records(tiny_kvdb_t *db, bool is_power_on)
{
    uint32_t addr_from, addr_to, addr;
    uint32_t len, bytes_to_copy;
    uint8_t buff[TINY_KVDB_BUFFER_SIZE];
    record_header_t header;
    int ret;

    addr_from = db->data_addr + sizeof(page_header_t);
    addr_to = db->swap_addr + sizeof(page_header_t);

    /* Variables to check incomplete deleted record, CAN ONLY apper when power off, and AT MOST one record */
    bool have_pre_deleted_record = false;
    record_info_t pre_deleted_record;     

#if TINY_KVDB_USE_CACHE

    uint32_t i = 0;

    /* if cache is not full, just copy records in the cache. */
    if (!is_power_on)
    {
        for (i = 0; i < TINY_KVDB_CACHE_NUM; i++)
        {            
            if (db->cache[i].key == RECORD_KEY_UNUSED || db->cache[i].key == RECORD_KEY_DELETED)
            {
                break;
            }
        }

        if (i != TINY_KVDB_CACHE_NUM)
        {
            for (i = 0; i < TINY_KVDB_CACHE_NUM; i++)
            {
                if (db->cache[i].key == RECORD_KEY_UNUSED)
                {
                    break;
                }
                if (db->cache[i].key == RECORD_KEY_DELETED)
                {
                    continue;
                }

                addr = db->cache[i].addr;
                len = aligned_size(db, sizeof(record_header_t) + db->cache[i].len);
                addr = addr_from;
                while (len) 
                {
                    bytes_to_copy = MIN(TINY_KVDB_BUFFER_SIZE, len);
                    ret = db->flash_param.read(addr, buff, bytes_to_copy);
                    if (ret < 0)
                    {
                        return ret;
                    }

                    ret = db->flash_param.write(addr_to, buff, bytes_to_copy);
                    if (ret < 0)
                    {
                        return ret;
                    }

                    len -= bytes_to_copy;
                    addr += bytes_to_copy;
                    addr_to += bytes_to_copy;
                }
            } 

            db->dirty_size = 0;
            db->next_addr = addr_to;
            return TINY_DB_ERR_NONE;           
        }
        else
        {
            clear_cache(db);
        }
    }
    else
    {
        clear_cache(db);
    }
    
#endif

    db->dirty_size = 0;
    while (addr_from < db->data_addr + db->total_size - sizeof(record_header_t))
    {
        ret = db->flash_param.read(addr_from, &header, sizeof(record_header_t));
        if (ret < 0)
        {
            return ret;
        }
        TINY_KVDB_DEBUG("record_key:0x%04X, record_len:%d\r\n", header.key, header.len);

        if (header.key == RECORD_KEY_UNUSED)
        {
            db->next_addr = addr_to;
            return TINY_DB_ERR_NONE;
        }
        else
        {
            /* data length error, assume that all the rest is dirty data, return */
            if (header.len == 0 || addr_from + sizeof(record_header_t) + header.len > db->data_addr + db->total_size)
            {
                db->next_addr = addr_to;
                return TINY_DB_ERR_NONE;
            }
            
            /* record has been deleted */
            if (header.key == RECORD_KEY_DELETED)
            {
                /* do nothing */
            }
            /* record status is invalid, so the record may be incomplete */
            else if (header.status != RECORD_STATUS_OK && header.status != RECORD_STATUS_PRE_DEL)
            {
                /* do nothing */
            }
            /* record key and status are all valid */
            else
            {
                if (have_pre_deleted_record)
                {
                    if (pre_deleted_record.key == header.key)
                    {
                        ret = flag_record_deleted(db, pre_deleted_record.addr);
                        if (ret < 0)
                        {
                            return ret;
                        }
                        have_pre_deleted_record = false;
                        db->dirty_size += sizeof(record_header_t) + aligned_size(db, pre_deleted_record.len);
                    }
                }

                if (header.status == RECORD_STATUS_PRE_DEL)
                {
                    pre_deleted_record.key = header.key;
                    pre_deleted_record.len = header.len;
                    pre_deleted_record.addr = addr_to;
                    have_pre_deleted_record = true;
                }
#if TINY_KVDB_USE_CACHE
                /* record status ok, put into the cache */
                else
                {
                    if (i < TINY_KVDB_CACHE_NUM)
                    {
                        db->cache[i].key = header.key;
                        db->cache[i].len = header.len;
                        db->cache[i].addr = addr_to;
                        i++;
                    }
                }
#endif

                /* copy record */
                ret = db->flash_param.write(addr_to, (uint8_t *)&header, sizeof(record_header_t));
                if (ret < 0)
                {
                    return ret;
                }
                addr_to += sizeof(record_header_t);

                len = aligned_size(db, header.len);
                addr = addr_from + sizeof(record_header_t);
                while (len) 
                {
                    bytes_to_copy = MIN(TINY_KVDB_BUFFER_SIZE, len);
                    ret = db->flash_param.read(addr, buff, bytes_to_copy);
                    if (ret < 0)
                    {
                        return ret;
                    }

                    ret = db->flash_param.write(addr_to, buff, bytes_to_copy);
                    if (ret < 0)
                    {
                        return ret;
                    }

                    len -= bytes_to_copy;
                    addr += bytes_to_copy;
                    addr_to += bytes_to_copy;
                }

            }

            addr_from += sizeof(record_header_t) + aligned_size(db, header.len);
        }
    }

#if TINY_KVDB_USE_CACHE
    /* the pre_deleted_record is not updated, so recover it */
    if (have_pre_deleted_record)
    {
        if (i < TINY_KVDB_CACHE_NUM)
        {
            db->cache[i].key = pre_deleted_record.key;
            db->cache[i].len = pre_deleted_record.len;
            db->cache[i].addr = pre_deleted_record.addr;
        }
    }
#endif

    return TINY_DB_ERR_NONE;
}

static int do_gc(tiny_kvdb_t *db, bool is_power_on)
{
    if (!is_power_on)
    {
        if (db->dirty_size == 0)
        {
            return TINY_DB_ERR_NO_SPACE;
        }
    }

    int ret;

    /* First step, tag data page as deleted */
    ret = set_page_tag(db, db->data_addr, PAGE_TAG_DEL);
    if (ret < 0)
    {
        return ret;
    }

    /* Second step, tag swap page as garbage collecting */
    if (is_power_on)
    {
        page_header_t hdr;
        ret = db->flash_param.read(db->swap_addr, &hdr, sizeof(page_header_t));
        if (ret < 0)
        {
            return ret;
        }

        /* is copying record when last power off, so may be some unreliable data
           in the page, so erase page first */
        if (hdr.tag == PAGE_TAG_GC)
        {
            ret = db->flash_param.erase(db->swap_addr, db->total_size);
            if (ret < 0)
            {
                return ret;
            }
        }
    }
    ret = set_page_tag(db, db->swap_addr, PAGE_TAG_GC);
    if (ret < 0)
    {
        return ret;
    }

    /* Third step, copy records from data page to swap page */
    ret = copy_records(db, is_power_on);
    if (ret < 0)
    {
        return ret;
    }

    /* Fourth step, tag swap page as data */
    ret = set_page_tag(db, db->swap_addr, PAGE_TAG_DATA);
    if (ret < 0)
    {
        return ret;
    }

    /* Fifth step, erase data page */
    ret = db->flash_param.erase(db->data_addr, db->total_size);
    if (ret < 0)
    {
        return ret;
    }

    /* Sixth step, tag erased page as swap */
    ret = set_page_tag(db, db->data_addr, PAGE_TAG_SWAP);
    if (ret < 0)
    {
        return ret;
    }

    /* Last step, exchange the page tags */
    uint32_t addr = db->data_addr;
    db->data_addr = db->swap_addr;
    db->swap_addr = addr;

    return TINY_DB_ERR_NONE;
}

int tiny_kvdb_gc(tiny_kvdb_t *db)
{
    int ret;

    TINY_KVDB_INFO("garbage collecting... \r\n");

    ret = do_gc(db, false);
    if (ret < 0)
    {
        return ret;
    }

    TINY_KVDB_INFO("garbage collection done!\r\n");
    TINY_KVDB_DEBUG("data_addr:0x%08X, swap_addr:0x%08X, next_addr:0x%08X, dirty_size:%d\r\n", db->data_addr, db->swap_addr, db->next_addr, db->dirty_size);    
#if TINY_KVDB_USE_CACHE
    for (uint8_t i = 0; i < TINY_KVDB_CACHE_NUM; i++)
    {
        TINY_KVDB_DEBUG("cache %d: key=0x%04X, len=%d, addr=0x%08X\r\n", i, db->cache[i].key, db->cache[i].len, db->cache[i].addr);
    }
#endif 

    return TINY_DB_ERR_NONE;
}


int tiny_kvdb_format(tiny_kvdb_t *db)
{
    if (!db->init_ok)
    {
        return TINY_DB_ERR_NOT_INIT;
    }

    if (db->flash_param.erase(db->data_addr, db->total_size) < 0)
    {
        return TINY_DB_ERR_ERASE;
    }
    if (db->flash_param.erase(db->swap_addr, db->total_size) < 0)
    {
        return TINY_DB_ERR_ERASE;
    }

    if (set_page_tag(db, db->data_addr, PAGE_TAG_DATA) < 0)
    {
        return TINY_DB_ERR_WRITE;
    }
    if (set_page_tag(db, db->swap_addr, PAGE_TAG_SWAP) < 0)
    {
        return TINY_DB_ERR_WRITE;
    }
    db->next_addr = db->data_addr + sizeof(page_header_t);
    db->dirty_size = 0;
    db->init_ok = 1;

#if TINY_KVDB_USE_CACHE
    clear_cache(db);
#endif 

    return TINY_DB_ERR_NONE;
} 

int tiny_kvdb_get_free_space(tiny_kvdb_t *db)
{
    if (!db->init_ok)
    {
        return TINY_DB_ERR_NOT_INIT;
    }
    
    int free_size = (int32_t)db->data_addr + (int32_t)db->total_size - (int32_t)db->next_addr - (int32_t)sizeof(record_header_t);
    if (free_size < 0)
    {
        free_size = 0;
    }

    return free_size;
}


/******************************** END OF FILE *********************************/