//--------------------------------- INCLUDES ----------------------------------
#include "nvs_storage.h"
#include "fds.h"
#include "bond3_watchdog.h"
#include "nrf_pwr_mgmt.h"
#include "main.h"

#define NRF_LOG_MODULE_NAME "nvs_storage"
#define NRF_LOG_LEVEL 3
#include "nrf_log.h"
#include "nrf_log_ctrl.h"

//---------------------------------- MACROS -----------------------------------
#define FILE_ID     0x1111
#define REC_KEY     0x2222
#define OUID_FILE_ID        (0x0020)
#define OUID_BACKUP_FILE_ID (0x0030)
#define OUID_REC_KEY        (0x0010)

//---------------------- PRIVATE FUNCTION PROTOTYPES --------------------------
static void run_garbage_collector(void);
static void fds_evt_handler(fds_evt_t const *p_evt);
static void wait_for_fds_ready(void);
static void wait_for_gc_to_finish(void);
static void wait_for_fds_update(void);
static void wait_for_fds_write(void);
static void wait_for_fds_delete(void);
static uint32_t init_fds_and_wait_until_ready(void);
static uint32_t check_for_fds_corruption(fds_stat_t *p_fds_stat);
static uint32_t create_blank_ouid_records(void);
static uint32_t clear_ouid_data(void);
static uint32_t write_ouid_data_to_flash(void);
static uint32_t search_for_and_init_ouid_data(FDSSysDataComponents_t *p_fds_sys_data);
static uint32_t open_ouid_data_record(FDSSysDataComponents_t *p_fds_sys_data);
static uint32_t open_ouid_data_backup_record(FDSSysDataComponents_t *p_fds_sys_data);
static uint32_t handle_possible_ouid_data_corruptions(FDSSysDataComponents_t *p_fds_sys_data);
static void dump_all_system_data_to_console(void);
static void dump_ouid_data(void);

//------------------------- STATIC DATA & CONSTANTS ---------------------------
/* Flag to check fds initialization. */
static bool volatile m_fds_initialized = false;
static bool volatile m_gc_is_running   = false;
static bool volatile m_fds_is_writing  = false;
static bool volatile m_fds_is_updating = false;
static bool volatile m_fds_is_deleting = false;

static uint32_t ouid_ = 0;
static uint32_t ouid_backup_ = 0;
static bool ouid_is_stored_ = false;
static bool _is_ouid_data_valid = false;
static bool _is_ouid_data_backup_valid = false;

/* Array to map FDS return values to strings. */
//NOTE(bojankoce): this serves for debug only. Could be deleted to save some space, if needed. 
char const *fds_err_str[] = {
    "FDS_SUCCESS",
    "FDS_ERR_OPERATION_TIMEOUT",
    "FDS_ERR_NOT_INITIALIZED",
    "FDS_ERR_UNALIGNED_ADDR",
    "FDS_ERR_INVALID_ARG",
    "FDS_ERR_NULL_ARG",
    "FDS_ERR_NO_OPEN_RECORDS",
    "FDS_ERR_NO_SPACE_IN_FLASH",
    "FDS_ERR_NO_SPACE_IN_QUEUES",
    "FDS_ERR_RECORD_TOO_LARGE",
    "FDS_ERR_NOT_FOUND",
    "FDS_ERR_NO_PAGES",
    "FDS_ERR_USER_LIMIT_REACHED",
    "FDS_ERR_CRC_CHECK_FAILED",
    "FDS_ERR_BUSY",
    "FDS_ERR_INTERNAL",
};

/* Array to map FDS events to strings. */
static char const *fds_evt_str[] = {
    "FDS_EVT_INIT",       "FDS_EVT_WRITE",    "FDS_EVT_UPDATE",
    "FDS_EVT_DEL_RECORD", "FDS_EVT_DEL_FILE", "FDS_EVT_GC",
};

static fds_record_desc_t  ouid_fds_desc_ = {0};  // Descriptor for our ouid
static fds_record_desc_t  ouid_backup_fds_desc_ = {0};  // Descriptor for our backup ouid

static const fds_record_chunk_t ouid_data_chunk = {
    .p_data = &ouid_,
    .length_words = 1
};

static const fds_record_chunk_t ouid_backup_data_chunk = {
    .p_data = &ouid_backup_,
    .length_words = 1
};


static fds_record_t const ouid_record = {   // Record to hold our ouid
    .file_id = OUID_FILE_ID,
    .key = OUID_REC_KEY,
    .data.p_chunks = &ouid_data_chunk,
    .data.num_chunks = 1
};

static fds_record_t const ouid_backup_record = {   // Record to hold our backup ouid
    .file_id = OUID_BACKUP_FILE_ID,
    .key = OUID_REC_KEY,
    .data.p_chunks = &ouid_backup_data_chunk,
    .data.num_chunks = 1    
};

/* Keep track of the progress of a delete_all operation. */
static struct {
    bool delete_next; //!< Delete next record.
    bool pending; //!< Waiting for an fds FDS_EVT_DEL_RECORD event, to delete
                  //! the next record.
} m_delete_all;

//---------------------------- PRIVATE FUNCTIONS ------------------------------

static void fds_evt_handler(fds_evt_t const *p_evt) {
    //NRF_LOG_DEBUG("Event: %s received (%s)", fds_evt_str[p_evt->id],
    //              fds_err_str[p_evt->result]);

    switch (p_evt->id) {
    case FDS_EVT_INIT:
        NRF_LOG_INFO("FDS_EVT_INIT! 0x%X\r\n", p_evt->result);
        if (p_evt->result == FDS_SUCCESS) {
            m_fds_initialized = true;
            NRF_LOG_INFO("FDS initialized!\r\n");
        }
        break;
    case FDS_EVT_GC:
        NRF_LOG_INFO("FDS_EVT_GC!\r\n");
        m_gc_is_running = false;
        break;
    case FDS_EVT_WRITE: {
        if (p_evt->result == FDS_SUCCESS) {
            NRF_LOG_DEBUG("Record ID:\t0x%04x", p_evt->write.record_id);
            NRF_LOG_DEBUG("File ID:\t0x%04x", p_evt->write.file_id);
            NRF_LOG_DEBUG("Record key:\t0x%04x", p_evt->write.record_key);
        }
        NRF_LOG_INFO("FDS_EVT_WRITE! 0x%X\r\n", p_evt->result);
        m_fds_is_writing = false;
    } break;

    case FDS_EVT_DEL_FILE:
    case FDS_EVT_DEL_RECORD: {
        if (p_evt->result == FDS_SUCCESS) {
            NRF_LOG_DEBUG("Record ID:\t0x%04x", p_evt->del.record_id);
            NRF_LOG_DEBUG("File ID:\t0x%04x", p_evt->del.file_id);
            NRF_LOG_DEBUG("Record key:\t0x%04x", p_evt->del.record_key);
        }
        NRF_LOG_INFO("FDS_EVT_DEL_FILE/RECORD! 0x%X\r\n", p_evt->result);
        m_fds_is_deleting    = false;
        m_delete_all.pending = false;
    } break;

    case FDS_EVT_UPDATE: {
        if (p_evt->result == FDS_SUCCESS) {
            NRF_LOG_DEBUG("Record ID:\t0x%04x", p_evt->write.record_id);
            NRF_LOG_DEBUG("File ID:\t0x%04x", p_evt->write.file_id);
            NRF_LOG_DEBUG("Record key:\t0x%04x", p_evt->write.record_key);
        }
        NRF_LOG_INFO("FDS_EVT_UPDATE!: 0x%X\r\n", p_evt->result);
        m_fds_is_updating = false;
    } break;

    default:
        break;
    }
}

/**@brief   Wait for fds to initialize. */
static void wait_for_fds_ready(void) {
    while (!m_fds_initialized) {
        //nrf_pwr_mgmt_run();
    }
}

static void wait_for_gc_to_finish(void) {
    while (m_gc_is_running)
        ;
}

static void wait_for_fds_update(void) {
    while (m_fds_is_updating)
        ;
}

static void wait_for_fds_write(void) {
    while (m_fds_is_writing)
        ;
}

static void wait_for_fds_delete(void) {
    while (m_fds_is_deleting)
        ;
}

static void run_garbage_collector(void) {
    uint32_t err_code;

    NRF_LOG_DEBUG("Running garbage collector.\r\n");
    m_gc_is_running = true;
    feed_wdt();
    err_code = fds_gc();
    feed_wdt();
    if (err_code) {
        //NRF_LOG_WARNING("GC failure: %s", fds_err_str[err_code]);
        NRF_LOG_WARNING("GC failure: %d", err_code);
        return;
    }
    wait_for_gc_to_finish();
}

static uint32_t init_fds_and_wait_until_ready(void) {
    NRF_LOG_INFO("init_fds_and_wait_until_ready::Starting Process\r\n");
    uint32_t err_code;

    /* Register first to receive an event when initialization is complete. */
    err_code = fds_register(fds_evt_handler);
    NRF_LOG_INFO("fds_register: 0x%X\r\n", err_code);
    if (err_code) {
        //APP_ERROR_CHECK(err_code);
        return err_code;
    }

    err_code = fds_init();
    NRF_LOG_INFO("fds_init: 0x%X\r\n", err_code);
    if (err_code == FDS_ERR_NO_PAGES) {
        NRF_LOG_ERROR("FDS found NO Valid pages. Attempting to erase ALL FDS pages.\r\n");
        //TODO(bojankoce): check if it is needed to erase all FDS pages!
        feed_wdt();        
        perform_system_reset();

    } else if (err_code) {
        //APP_ERROR_CHECK(err_code);
        return err_code;
    }

    feed_wdt();

    wait_for_fds_ready();
    return err_code;
}

static uint32_t check_for_fds_corruption(fds_stat_t *p_fds_stat) {
    NRF_LOG_INFO("check_for_fds_corruption::Starting Process\r\n");
    uint32_t err_code;
    bool     gc_required = false;

    if (!p_fds_stat) {
        NRF_LOG_ERROR("Null Pointer\r\n");
        return NRF_ERROR_NULL;
    }

    err_code = fds_stat(p_fds_stat);
    if (err_code) {
        //APP_ERROR_CHECK(err_code);
        return err_code;
    }

    NRF_LOG_INFO("Found %d valid records.\r\n", p_fds_stat->valid_records);
    NRF_LOG_INFO("Found %d dirty records.\r\n", p_fds_stat->dirty_records);

    if (p_fds_stat->dirty_records >=
        MAX_DIRTY_RECORDS_BEFORE_GARBAGE_COLLECTION) {
        NRF_LOG_WARNING("Sufficient dirty records to run GC.\r\n");
        gc_required = true;
    }

    if (gc_required) {
        feed_wdt();
        run_garbage_collector();
    }

    return err_code;
}

static uint32_t search_for_and_init_ouid_data(FDSSysDataComponents_t *p_fds_sys_data) {
    uint32_t err_code;

    if (!p_fds_sys_data) {
        NRF_LOG_ERROR("Null Pointer\r\n");
        return NRF_ERROR_NULL;
    }

    //APP_ERROR_CHECK(open_ouid_data_record(p_fds_sys_data));
    //APP_ERROR_CHECK(open_ouid_data_backup_record(p_fds_sys_data));
    err_code = open_ouid_data_record(p_fds_sys_data);   
    NRF_LOG_INFO("open_ouid_data_record: 0x%X\r\n", err_code); 
    
    err_code = open_ouid_data_backup_record(p_fds_sys_data);
    NRF_LOG_ERROR("open_ouid_data_backup_record: 0x%X\r\n", err_code); 
    
    //err_code = (handle_possible_ouid_data_corruptions(p_fds_sys_data));
    //APP_ERROR_CHECK(err_code);
    err_code = handle_possible_ouid_data_corruptions(p_fds_sys_data);
    NRF_LOG_INFO("handle_possible_ouid_data_corruptions: 0x%X\r\n", err_code);

    return err_code;
}

static uint32_t open_ouid_data_record(FDSSysDataComponents_t *p_fds_sys_data) {
    uint32_t err_code;

    if (!p_fds_sys_data) {
        NRF_LOG_ERROR("Null Pointer\r\n");
        return NRF_ERROR_NULL;
    }

    err_code = fds_record_find(ouid_record.file_id,
                               ouid_record.key, &ouid_fds_desc_,
                               &(p_fds_sys_data->sys_data_ouid_token));
    NRF_LOG_INFO("fds_record_find: 0x%X\r\n", err_code);
    if (!err_code) {
        NRF_LOG_INFO("System Data OUID found\r\n");
        err_code = fds_record_open(&ouid_fds_desc_,
                                   &(p_fds_sys_data->sys_data_ouid_record));
        if (!err_code) {
            NRF_LOG_WARNING("System Data OUID valid\r\n");
            _is_ouid_data_valid = true;            
        } else {
            NRF_LOG_WARNING("System Data OUID invalid\r\n");
            _is_ouid_data_valid = false;
        }
    } else {
        NRF_LOG_INFO("System Data OUID not found\r\n");
    }

    return err_code;
}

static uint32_t open_ouid_data_backup_record(FDSSysDataComponents_t *p_fds_sys_data) {
    uint32_t err_code;

    if (!p_fds_sys_data) {
        NRF_LOG_ERROR("Null Pointer\r\n");
        return NRF_ERROR_NULL;
    }
    
    err_code = fds_record_find(
        ouid_backup_record.file_id, ouid_backup_record.key,
        &ouid_backup_fds_desc_,
        &(p_fds_sys_data->sys_data_ouid_backup_token));
    if (!err_code) {
        NRF_LOG_INFO("System Data OUID Backup found\r\n");
        err_code =
            fds_record_open(&ouid_backup_fds_desc_,
                            &(p_fds_sys_data->sys_data_ouid_backup_record));
        if (!err_code) {
            NRF_LOG_WARNING("System Data OUID Backup valid\r\n");
            _is_ouid_data_backup_valid = true;            
        } else {
            NRF_LOG_WARNING("System Data OUID Backup invalid\r\n");
            _is_ouid_data_backup_valid = false;
        }
    } else {
        NRF_LOG_INFO("System Data OUID Backup not found\r\n");
    }
    return err_code;
}

static uint32_t handle_possible_ouid_data_corruptions(FDSSysDataComponents_t *p_fds_sys_data) {
    uint32_t err_code;

    if (!p_fds_sys_data) {
        NRF_LOG_ERROR("Null pointer\r\n");
        return NRF_ERROR_NULL;
    }

    if (_is_ouid_data_valid && _is_ouid_data_backup_valid) {
        // All OUID data is valid
        NRF_LOG_INFO("All system data ouid records are intact.\r\n");
        memcpy(&ouid_,
               p_fds_sys_data->sys_data_ouid_record.p_data,
               sizeof(ouid_));
        memcpy(&ouid_backup_,
               p_fds_sys_data->sys_data_ouid_backup_record.p_data,
               sizeof(ouid_));

        /* Close the records when done reading. */
        err_code = fds_record_close(&ouid_fds_desc_);
        APP_ERROR_CHECK(err_code);
        err_code = fds_record_close(&ouid_backup_fds_desc_);
        APP_ERROR_CHECK(err_code);

    } else if (!_is_ouid_data_valid && _is_ouid_data_backup_valid) {
        NRF_LOG_WARNING("Using system data OUID backup.\r\n");
        memcpy(&ouid_,
               p_fds_sys_data->sys_data_ouid_backup_record.p_data,
               sizeof(ouid_));
        memcpy(&ouid_backup_,
               p_fds_sys_data->sys_data_ouid_backup_record.p_data,
               sizeof(ouid_));

        err_code = fds_record_close(&ouid_backup_fds_desc_);
        APP_ERROR_CHECK(err_code);

    } else if (_is_ouid_data_valid && !_is_ouid_data_backup_valid) {
        NRF_LOG_WARNING("System data OUID backup invalid. Using system data OUID data.\r\n");
        memcpy(&ouid_,
               p_fds_sys_data->sys_data_ouid_record.p_data,
               sizeof(ouid_));
        memcpy(&ouid_backup_,
               p_fds_sys_data->sys_data_ouid_record.p_data,
               sizeof(ouid_));

        err_code = fds_record_close(&ouid_fds_desc_);
        APP_ERROR_CHECK(err_code);

    } else {
        NRF_LOG_WARNING("Neither system data OUID nor system data OUID backup were valid! \r\n");
        err_code = NRF_ERROR_NOT_FOUND;
    }

    return err_code;
}

static uint32_t create_blank_ouid_records(void){
    uint32_t err_code = NRF_SUCCESS;

    // OUID data
    NRF_LOG_INFO("Creating new ouid data file\r\n");

    m_fds_is_writing = true;

    err_code = fds_record_write(&ouid_fds_desc_, &ouid_record); 
    NRF_LOG_INFO("fds_record_write: 0x%X\r\n", err_code);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    } else {
        NRF_LOG_INFO("fds_record_write OK!\r\n");
    }

    wait_for_fds_write();    

    NRF_LOG_INFO("Creating new ouid backup file\r\n");
    m_fds_is_writing = true;
    err_code         = fds_record_write(&ouid_backup_fds_desc_,
                                &ouid_backup_record);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    }
    wait_for_fds_write();

    return err_code;
}

static uint32_t clear_ouid_data(void){
    uint32_t virgin_ouid = 0;
    return set_ouid(&virgin_ouid);
}

static void dump_all_system_data_to_console(void){
    dump_ouid_data();
}

static void dump_ouid_data(void){
    NRF_LOG_INFO("OUID: 0x%X\r\n", ouid_);
    NRF_LOG_INFO("Backup OUID: 0x%X\r\n", ouid_backup_);
}

//------------------------------ PUBLIC FUNCTIONS -----------------------------

uint32_t init_nvs_storage(void) {
    uint32_t               err_code;
    FDSSysDataComponents_t fds_sys_data;

    memset(&fds_sys_data, 0, sizeof(fds_sys_data));

    err_code = init_fds_and_wait_until_ready();
    //NRF_LOG_INFO("init_fds_and_wait_until_ready: 0x%X\r\n", err_code);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    }

    err_code = check_for_fds_corruption(&(fds_sys_data.fds_stat));
    //NRF_LOG_INFO("check_for_fds_corruption: 0x%X\r\n", err_code);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    }

    err_code = search_for_and_init_ouid_data(&fds_sys_data);
    NRF_LOG_INFO("search_for_and_init_ouid_data: 0x%X\r\n", err_code);

    if (err_code) {
        NRF_LOG_WARNING("Unable to load ouid from internal flash. Will "
                        "create virgin records: 0x%X\r\n",
                        err_code);
        err_code = create_blank_ouid_records();
        if (err_code) {
            NRF_LOG_ERROR("Unable to create new OUID records: 0x%X\r\n",
                          err_code);
        } else {
            NRF_LOG_INFO("OUID records newly created\r\n");
        }
    } else {
        NRF_LOG_INFO("OUID records loaded from internal flash\r\n");
    }    

    //TODO(bojankoce): We can use non-volatile flash to store some System Status/Settings data

    dump_all_system_data_to_console();
    NRF_LOG_INFO("Internal flash initialized\r\n");

    return err_code;
}

uint32_t get_ouid(uint32_t *ouid) {
    int ret = NRF_SUCCESS;
#if STORE_IN_NVS    
    *ouid = ouid_;
#else 
    *ouid = ouid_;
    ret = 4;
#endif
    return ret;
}

uint32_t set_ouid(uint32_t *new_ouid) {
    uint32_t err_code = NRF_SUCCESS;
#if STORE_IN_NVS    
    
    fds_stat_t stat = {0};

    if (!new_ouid) {
        NRF_LOG_ERROR("Null pointer\r\n");
        return NRF_ERROR_NULL;
    }

    ouid_ = *new_ouid;
    ouid_backup_ = *new_ouid;
    
    feed_wdt();

    m_fds_is_updating = true;
    err_code = fds_record_update(&ouid_fds_desc_, &ouid_record);
    NRF_LOG_INFO("fds_record_update: 0x%X\r\n", err_code);
    if (err_code) {
        //APP_ERROR_CHECK(err_code);
        return err_code;
    }
    wait_for_fds_update();

    m_fds_is_updating = true;
    err_code          = fds_record_update(&ouid_backup_fds_desc_,
                                 &ouid_backup_record);
    NRF_LOG_INFO("fds_record_update: 0x%X\r\n", err_code);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    }
    wait_for_fds_update();

    err_code = fds_stat(&stat);
    if (err_code) {
        APP_ERROR_CHECK(err_code);
        return err_code;
    }
    NRF_LOG_INFO("Found %d valid records, %d dirty records.\r\n", stat.valid_records, stat.dirty_records);

    if (stat.dirty_records >= MAX_DIRTY_RECORDS_BEFORE_GARBAGE_COLLECTION) {
        run_garbage_collector();
    }

#else
    ouid_ = ouid;
    if(!ouid_is_stored()){
        ouid_is_stored_ = true; 
    }
#endif
    return err_code;    
}

bool ouid_is_stored(void) {
#if STORE_IN_NVS
    return _is_ouid_data_valid;
#else
    return ouid_is_stored_;
#endif
}