/**
 * @file    shot_storage.c
 * @brief   Shot Storage Module
 * @author  Kenny Tran
 * @date    Sep 02 2025
 * @version 0.0.1
 *
 * @section LICENSES
 * Copyright (C) Morgan Advanced Programmable Systems, Inc.
 */


/*--------------------------------------------------------------------------------------------------
 * INCLUDE FILES
 *------------------------------------------------------------------------------------------------*/
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>

#include <zephyr/kernel.h>
#include <zephyr/sys/printk.h>
#include <zephyr/sys/util.h>  /* Zephyr MIN() macro */

#include "../app/app.h"
#include "../app/shot_capture.h"
#include "../app/ble.h"
#include "../app/motion.h"

#include "../lib/errno.h"
#include "../lib/fs.h"
#include "../lib/print.h"

#include "../bsp/imu.h"
#include "../bsp/nvm_ext.h"
#include "../bsp/wdt.h"

#include "../wrappers/spi.h"

#include "shot_storage.h"


/*--------------------------------------------------------------------------------------------------
 * TYPE DEFINITIONS & ENUMERATIONS
 *------------------------------------------------------------------------------------------------*/


/*--------------------------------------------------------------------------------------------------
 * MACROS & CONSTANTS
 *------------------------------------------------------------------------------------------------*/
/**
 * @brief For use with reading NVM external status bit.
 * */
#define _SHOT_STORAGE_NVM_STATUS_BUSY_MASK      (0x01)
#define _SHOT_STORAGE_NVM_EXT_STATUS_BUSY       (0x01)

#define _SHOT_STORAGE_NVM_PAGE_PROG_LEN         (4)


K_THREAD_STACK_DEFINE(_shot_storage_write_stack, SHOT_STORAGE_WRITE_THREAD_STACK_SIZE);
K_THREAD_STACK_DEFINE(_shot_storage_read_stack, SHOT_STORAGE_READ_THREAD_STACK_SIZE);

/*--------------------------------------------------------------------------------------------------
 * PUBLIC VARIABLE DEFINITIOPNS
 *------------------------------------------------------------------------------------------------*/


/*--------------------------------------------------------------------------------------------------
 * PRIVATE CONSTANTS & VARIABLES
 *------------------------------------------------------------------------------------------------*/
static SHOT_STORAGE_file_s  _shot_storage_files[SHOT_STORAGE_MAX_FILES];
static SHOT_STORAGE_file_s  _shot_storage_current_file = {
    .header_ver                = 0,
    .shot_data_format_ver      = 0,
    .reserved                  = 0,
    .shot_data_offset          = sizeof(SHOT_STORAGE_file_s),
    .gyro_bits_per_axis        = 16,
    .accel_bits_per_axis       = 16,
    .IMU_orientation           = 0x0FED, /* away from user: -x, to the right of user: -y, towards earth: -z */
};
K_SEM_DEFINE(write_shotdata_sem, 0, 1);
static char     _temp_name[SHOT_STORAGE_FILENAME_LEN];
static uint8_t  _shot_storage_buf[_SHOT_STORAGE_FILE_SIZE];
static uint32_t _shot_storage_addr = _SHOT_STORAGE_START_ADDR;
static uint32_t _shot_storage_write_pages = 0;
static uint32_t _shot_storage_bytes_remaining = 0;
static uint8_t  _shot_storage_staging_buf[_SHOT_STORAGE_NVM_PAGE_PROG_LEN + NVM_EXT_PAGE_SIZE];
static uint32_t _shot_storage_logic_stream_start = 0;
static FS_file_t _shot_storage_tracking_file; 

/* Circular file storage tracking (persisted to internal flash) */
static SHOT_STORAGE_tracking_s _shot_storage_tracking = {
    .write_idx     = 0,
    .file_cnt      = 0,
    .total_written = 0,
};

/* Read task synchronization */
K_SEM_DEFINE(read_shotdata_request, 0, 1);   /* Trigger read task */
K_SEM_DEFINE(read_shotdata_buf_ready, 0, 1); /* Buffer ready for consumer */
K_SEM_DEFINE(read_shotdata_consumed, 0, 1);  /* Consumer finished reading */

/* Idle fence for write task */
static K_SEM_DEFINE(shot_storage_write_idle_sem, 0, 1);
static atomic_t _shot_storage_write_idle_requested = ATOMIC_INIT(0);

/* Track who requested the read */
static atomic_t _read_request_source = ATOMIC_INIT(READ_REQ_NONE);

/*--------------------------------------------------------------------------------------------------
 * PRIVATE FUNCTION PROTOTYPES
 *------------------------------------------------------------------------------------------------*/
static void _shot_storage_files_init (void);
static bool _shot_storage_is_empty (uint8_t *p_data, uint32_t len);
static ERRNO_e _shot_storage_nvm_page_write (uint32_t addr, uint8_t *p_src, uint32_t rem_len);
static void _shot_storage_nvm_wait_til_ready (void);
static void _shot_storage_write_task (void *p1, void *p2, void *p3);
static void _shot_storage_read_task (void *p1, void *p2, void *p3);
static void _shot_storage_tracking_read (void);
static ERRNO_e _shot_storage_tracking_write (void);


/*--------------------------------------------------------------------------------------------------
 * PUBLIC FUNCTION DEFINITIONS
 *------------------------------------------------------------------------------------------------*/
void SHOT_STORAGE_init (void)
{
    PRINT (INFO, "Initializing SHOT STORAGE...\r\n");

    if (NVM_EXT_init () < ERRNO_SUCCESS)
    {
        PRINT (CRITICAL, "Failure to initialize NVM_EXT.\r\n");
    }

    _shot_storage_files_init ();
}

void SHOT_STORAGE_task_install (void)
{
    static        k_tid_t           write_id;
    static struct k_thread          write_data;
    static        k_tid_t           read_id;
    static struct k_thread          read_data;

    write_id = k_thread_create (&write_data,
                          _shot_storage_write_stack,
                          SHOT_STORAGE_WRITE_THREAD_STACK_SIZE,
                          _shot_storage_write_task,
                          NULL,
                          NULL,
                          NULL,
                          SHOT_STORAGE_WRITE_THREAD_PRIORITY,
                          0,
                          K_NO_WAIT);

    k_thread_name_set(&write_data, "SHOT_STORAGE_write_task");

    read_id = k_thread_create (&read_data,
                          _shot_storage_read_stack,
                          SHOT_STORAGE_READ_THREAD_STACK_SIZE,
                          _shot_storage_read_task,
                          NULL,
                          NULL,
                          NULL,
                          SHOT_STORAGE_READ_THREAD_PRIORITY,
                          0,
                          K_NO_WAIT);

    k_thread_name_set(&read_data, "SHOT_STORAGE_read_task");
}

ERRNO_e SHOT_STORAGE_get_buf(uint32_t size, uint8_t **pp_data)
{
    if (size > _SHOT_STORAGE_FILE_SIZE) {
        *pp_data = NULL;
        PRINT (CRITICAL, "Failure to get buffer.  Requested size too large!\r\n");
        return ERRNO_EXCEED_FILE_SIZE;
    }

    /* return pointer to the statically allocated shot data buffer */
    *pp_data =  &_shot_storage_buf[sizeof(SHOT_STORAGE_file_s)];

    return ERRNO_SUCCESS;
}

ERRNO_e SHOT_STORAGE_file_create(void)
{
    if (_shot_storage_current_file.shot_data_size > _SHOT_STORAGE_FILE_SIZE) {
        PRINT (CRITICAL, "Failure to create new file.  File size too large!\r\n");
        return ERRNO_EXCEED_FILE_SIZE;
    }
    
    /* Assign file name. */
    // int len = snprintf(_temp_name, SHOT_STORAGE_FILENAME_LEN, "%s", p_filename);
    // if (len >= SHOT_STORAGE_FILENAME_LEN) {
    //     PRINT(WARNING, "Filename truncated: %s\n", p_filename);
    //     return ERRNO_EXCEED_NAME_SIZE;
    // }

    /* fill file header info */
    k_mutex_lock(&shot_data_mutex, K_FOREVER);
    _shot_storage_logic_stream_start = shot_data.start;
    _shot_storage_current_file.time_of_shot = shot_data.time;
    _shot_storage_current_file.ambient_light_intensity = shot_data.ambient_light_intensity;
    _shot_storage_current_file.shot_data_size = shot_data.size;
    k_mutex_unlock(&shot_data_mutex);

    _shot_storage_current_file.power_cnt = APP_info_get ().power_cnt;
    _shot_storage_current_file.gyro_sample_rate_hz  = LSM6DS3TR_C_ODR_G_TO_HZ(SHOT_CAPTURE_config_get()->odr_g);
    _shot_storage_current_file.accel_sample_rate_hz = LSM6DS3TR_C_ODR_XL_TO_HZ(SHOT_CAPTURE_config_get()->odr_xl);
    _shot_storage_current_file.gyro_full_scale_dps  = LSM6DS3TR_C_FS_G_TO_DPS(SHOT_CAPTURE_config_get()->fs_g);
    _shot_storage_current_file.accel_full_scale_g   = LSM6DS3TR_C_FS_XL_TO_G(SHOT_CAPTURE_config_get()->fs_xl);
    _shot_storage_current_file.accel_cal            = MOTION_config_get ()->accel_cal;
    _shot_storage_current_file.gyro_cal             = MOTION_config_get ()->gyro_cal;

    return ERRNO_SUCCESS;
}

void SHOT_STORAGE_begin_write(void) {
    k_sem_give(&write_shotdata_sem);
}

void SHOT_STORAGE_write (void) {
    ERRNO_e err = ERRNO_SUCCESS;
    struct stream_cursor cur;
    uint32_t copied = 0;
    uint32_t slot = _shot_storage_tracking.write_idx;

    /* Generate unique name using total_written counter */
    //int len = snprintf(_shot_storage_current_file.name, SHOT_STORAGE_FILENAME_LEN, "%s_%u", _temp_name, (_shot_storage_tracking.total_written + 1));
    //if (len >= SHOT_STORAGE_FILENAME_LEN) {
    //    PRINT(WARNING, "Generated filename truncated\n");
    //}

    /* If all slots are full, erase the oldest slot before writing */
    if (_shot_storage_tracking.file_cnt >= SHOT_STORAGE_MAX_FILES) {
        uint32_t oldest_addr = _SHOT_STORAGE_START_ADDR + (slot * _SHOT_STORAGE_TOTAL_FILE_SIZE);
        if (NVM_EXT_erase_block(oldest_addr, NVM_EXT_BLOCK_SIZE_64KB)) {
            PRINT(CRITICAL, "Failure to erase oldest slot %u!\r\n", slot);
            return;
        }
        PRINT(INFO, "Overwrote oldest file: power_cnt: %u, time: %.3f s, (slot %u)\r\n", _shot_storage_files[slot].power_cnt, (float) _shot_storage_files[slot].time_of_shot / 1000.0, slot);
    }

    SHOT_STORAGE_stream_cursor_init(&cur, (const uint8_t *)&_shot_storage_current_file, sizeof(SHOT_STORAGE_file_s),
                                    &_shot_storage_buf[sizeof(SHOT_STORAGE_file_s)], _shot_storage_current_file.shot_data_size, _shot_storage_logic_stream_start);

    _shot_storage_addr = _SHOT_STORAGE_START_ADDR + (slot * _SHOT_STORAGE_TOTAL_FILE_SIZE);

    while (stream_cursor_remaining(&cur)) {
        copied = stream_cursor_read(&cur, &_shot_storage_staging_buf[4], NVM_EXT_PAGE_SIZE);
        err = _shot_storage_nvm_page_write(_shot_storage_addr, _shot_storage_staging_buf, copied);
        if (err != ERRNO_SUCCESS) break;
        _shot_storage_addr += NVM_EXT_PAGE_SIZE;
    }

    if (err != ERRNO_SUCCESS) {
        /* Erase the partially written slot */
        uint32_t fail_addr = _SHOT_STORAGE_START_ADDR + (slot * _SHOT_STORAGE_TOTAL_FILE_SIZE);
        NVM_EXT_erase_block(fail_addr, NVM_EXT_BLOCK_SIZE_64KB);
        PRINT(CRITICAL, "Write failed, erased partial data at slot %u\r\n", slot);
    } else {
        SHOT_CAPTURE_config_get()->shot_cnt++;
        SHOT_CAPTURE_config_get()->user_shot_cnt++;
        BLE_notify(BLE_CHAR_USER_SHOT_CNT, &(SHOT_CAPTURE_config_get()->user_shot_cnt), sizeof(SHOT_CAPTURE_config_get()->user_shot_cnt));
        /* Update tracking */
        _shot_storage_files[slot] = _shot_storage_current_file;
        if (_shot_storage_tracking.file_cnt < SHOT_STORAGE_MAX_FILES) {
            _shot_storage_tracking.file_cnt++;
        }
        _shot_storage_tracking.write_idx = (slot + 1) % SHOT_STORAGE_MAX_FILES;
        _shot_storage_tracking.total_written++;
        _shot_storage_tracking_write();
    }
}

bool SHOT_STORAGE_write_done (void)
{
    return (_shot_storage_write_pages == 0 && _shot_storage_bytes_remaining == 0);
}

uint8_t *SHOT_STORAGE_last_read (SHOT_STORAGE_file_s *p_file)
{
    uint32_t addr;
    uint16_t len;
    uint32_t newest_idx;

    if (!_shot_storage_tracking.file_cnt)
    {
        return NULL;
    }

    /* Newest file is the one most recently written */
    newest_idx = (_shot_storage_tracking.write_idx - 1 + SHOT_STORAGE_MAX_FILES) % SHOT_STORAGE_MAX_FILES;
    addr = _SHOT_STORAGE_START_ADDR + (newest_idx * _SHOT_STORAGE_TOTAL_FILE_SIZE);
    len  = _shot_storage_files[newest_idx].shot_data_size + sizeof(SHOT_STORAGE_file_s);

    if (NVM_EXT_read(addr, _shot_storage_buf, (uint32_t) len))
    {
        PRINT(CRITICAL, "Failure to read shot data from external NVM...\r\n");
    }

    *p_file = _shot_storage_files[newest_idx];

    return _shot_storage_buf;
}

void SHOT_STORAGE_last_erase(void)
{
    uint32_t addr;
    uint32_t newest_idx;

    if (!_shot_storage_tracking.file_cnt) {
        PRINT(INFO, "No files exist!\r\n");
        return;
    }

    newest_idx = (_shot_storage_tracking.write_idx - 1 + SHOT_STORAGE_MAX_FILES) % SHOT_STORAGE_MAX_FILES;
    addr = _SHOT_STORAGE_START_ADDR + (newest_idx * _SHOT_STORAGE_TOTAL_FILE_SIZE);

    if (NVM_EXT_erase_block (addr, NVM_EXT_BLOCK_SIZE_64KB)) {
        PRINT(CRITICAL, "Failure to erase shot data from external NVM...\r\n");
        return;
    }

    memset(&_shot_storage_files[newest_idx], 0xFF, sizeof(SHOT_STORAGE_file_s));
    _shot_storage_tracking.file_cnt--;
    _shot_storage_tracking.write_idx = newest_idx;
    _shot_storage_tracking_write();
}

ERRNO_e SHOT_STORAGE_eraseall(void)
{
    uint32_t addr;

    for (addr = _SHOT_STORAGE_START_ADDR; addr < _SHOT_STORAGE_END_ADDR; addr += NVM_EXT_BLOCK_SIZE_64KB) {
        if (NVM_EXT_erase_block(addr, NVM_EXT_BLOCK_SIZE_64KB)) {
            PRINT(CRITICAL, "Failure to erase shot data from external NVM...\r\n");
            return ERRNO_NVM_FAILURE;
        }
        k_msleep(5); /* Sleep to allow other tasks to run and prevent WDT reset during long erase */
    }

    /* Reset tracking and file headers */
    memset(_shot_storage_files, 0xFF, sizeof(_shot_storage_files));
    _shot_storage_tracking.write_idx    = 0;
    _shot_storage_tracking.file_cnt     = 0;
    _shot_storage_tracking.total_written = 0;
    _shot_storage_tracking_write();

    return ERRNO_SUCCESS;
}

uint32_t SHOT_STORAGE_get_file_cnt(void)
{
    return _shot_storage_tracking.file_cnt;
}

void SHOT_STORAGE_cli_cmd (int32_t argc, char **argv)
{
    if (argc == 0) {
        PRINT (INFO, "SHOT STORAGE\r\n");
        PRINT (INFO, "-------------------------------------------\r\n");
        PRINT (INFO, "Files       : %u / %u\r\n", _shot_storage_tracking.file_cnt, SHOT_STORAGE_MAX_FILES);
        PRINT (INFO, "Write index : %u\r\n", _shot_storage_tracking.write_idx);
        PRINT (INFO, "Total saved : %u\r\n", _shot_storage_tracking.total_written);
    }
    else if (argc == 1) {
        if (!strcmp (argv[0], "--list") || !strcmp (argv[0], "-l")) {
            uint32_t i;

            if (!_shot_storage_tracking.file_cnt) {
                PRINT (INFO, "No files stored.\r\n");
            }
            /* List from oldest to newest */
            for (i = 0; i < _shot_storage_tracking.file_cnt; i++) {
                uint32_t idx = (_shot_storage_tracking.write_idx - _shot_storage_tracking.file_cnt + i + SHOT_STORAGE_MAX_FILES) % SHOT_STORAGE_MAX_FILES;
                PRINT (INFO, "Slot %2u: power_cnt: %u, time: %.3f s\r\n", idx, _shot_storage_files[idx].power_cnt, (float) _shot_storage_files[idx].time_of_shot / 1000.0);
            }
        }
        else if (!strcmp (argv[0], "--test")) {
            #define FNAME  ("shot_storage")
            uint8_t *p_data;
            SHOT_STORAGE_get_buf(50000, &p_data);
            SHOT_STORAGE_file_create();

            memset (p_data, 9, 50000);

            PRINT (INFO, "writing...\r\n");

            SHOT_STORAGE_write();

            PRINT (INFO, "writing done. Current file count in flash: %u \r\n", _shot_storage_tracking.file_cnt);
        }
        else if (!strcmp (argv[0], "--done")) {
            PRINT (INFO, "%s\r\n", SHOT_STORAGE_write_done () ? "done" : "not done");
        }
        else if (!strcmp (argv[0], "--read") || !strcmp (argv[0], "-r")) {
            SHOT_STORAGE_file_s file;

            /* Request the read task to load data from flash */
            SHOT_STORAGE_request_read(READ_REQ_CLI);

            /* Wait for buffer to be ready */
            if (SHOT_STORAGE_wait_buf_ready() != 0) {
                PRINT(ERROR, "Failed to get read buffer\r\n");
                return;
            }

            uint8_t *p_data = _shot_storage_buf;
 
             /* Copy the entire file header struct (includes name) */
            memcpy(&file, p_data, sizeof(SHOT_STORAGE_file_s));

            /* Print file name with guaranteed null termination */
            PRINT(INFO, "Power Count: %u\r\n", file.power_cnt);
            PRINT(INFO, "Time of shot: %u ms\r\n", file.time_of_shot);
            PRINT(INFO, "Header version: %u\r\n", file.header_ver);
            PRINT(INFO, "Shot data format version: %u\r\n", file.shot_data_format_ver);
            PRINT(INFO, "Ambient light intensity: %u\r\n", file.ambient_light_intensity);
            PRINT(INFO, "Shot data offset: %u bytes\r\n", file.shot_data_offset);
            PRINT(INFO, "Shot data size: %u bytes\r\n", file.shot_data_size);
            PRINT(INFO, "Gyro bits per axis: %u\r\n", file.gyro_bits_per_axis);
            PRINT(INFO, "Accel bits per axis: %u\r\n", file.accel_bits_per_axis);
            PRINT(INFO, "Gyro sample rate: %u Hz\r\n", file.gyro_sample_rate_hz);
            PRINT(INFO, "Accel sample rate: %u Hz\r\n", file.accel_sample_rate_hz);
            PRINT(INFO, "Gyro full scale: %u dps\r\n", file.gyro_full_scale_dps);
            PRINT(INFO, "Accel full scale: %u g\r\n", file.accel_full_scale_g);
            PRINT(INFO, "IMU orientation: 0x%04X\r\n", file.IMU_orientation);      
            
            /*The shot data in the buffer is already in order thanks to the stream cursor function */
            PRINT(INFO, "GyroX, GyroY, GyroZ, AccX, AccY, AccZ\r\n");
            for (int i = sizeof(SHOT_STORAGE_file_s); i < sizeof(SHOT_STORAGE_file_s) + file.shot_data_size; i += 12) {
            PRINT(INFO, "%d,%d,%d,%d,%d,%d\r\n", *((int16_t *) &p_data[i]), \
                                                  *((int16_t *) &p_data[i + 2]), \
                                                  *((int16_t *) &p_data[i + 4]), \
                                                  *((int16_t *) &p_data[i + 6]), \
                                                  *((int16_t *) &p_data[i + 8]), \
                                                  *((int16_t *) &p_data[i + 10]));
                WDT_feed(); /* prevent watchdog timeout due to slow print on serial port over uart */
            }

            SHOT_STORAGE_last_erase();
            PRINT (INFO, "Power cnt: %u, time: %.3f s erased...\r\n", file.power_cnt, (float) file.time_of_shot/1000.0);
            
            /* Signal read task that we're done */
            SHOT_STORAGE_read_done();
        }
        else if (!strcmp (argv[0], "--eraseall") || !strcmp (argv[0], "-ra")) {
            if (SHOT_STORAGE_eraseall() == ERRNO_SUCCESS) {
                PRINT(INFO, "Erased all 64kB sectors and reset tracking!\r\n");
            }
        }   
    }
}

void SHOT_STORAGE_stream_cursor_init(struct stream_cursor *cur,
                        const uint8_t *header, uint32_t header_size,
                        const uint8_t *buffer, uint32_t data_size,
                        uint32_t write_ptr)
{
    cur->header      = header;
    cur->header_size = header_size;
    cur->buffer      = buffer;
    cur->data_size   = data_size;
    cur->write_ptr   = write_ptr;
    cur->total       = header_size + data_size;
    cur->offset      = 0;
}

uint32_t stream_cursor_read(struct stream_cursor *cur, uint8_t *dst, uint32_t len)
{
    uint32_t copied = 0;

    while (copied < len && cur->offset < cur->total) {
        uint32_t want = len - copied;

        if (cur->offset < cur->header_size) {
            /*
             * Region 1: Header
             * Straightforward linear copy from header array.
             */
            uint32_t hdr_remaining = cur->header_size - cur->offset;
            uint32_t chunk = MIN(want, hdr_remaining);

            memcpy(dst + copied, cur->header + cur->offset, chunk);
            copied      += chunk;
            cur->offset += chunk;

        } else {
            /*
             * Region 2: Circular buffer data
             *
             * Map the logical data offset to a physical buffer position,
             * then copy the largest contiguous chunk available without
             * crossing the buffer boundary.
             */
            uint32_t data_offset = cur->offset - cur->header_size;
            uint32_t buf_pos     = (cur->write_ptr + data_offset) % cur->data_size;

            /* How much data is left in the logical stream */
            uint32_t data_remaining = cur->data_size - data_offset;

            /* How many contiguous bytes from buf_pos to end of effective region */
            uint32_t contiguous = cur->data_size - buf_pos;

            uint32_t chunk = MIN(want, MIN(data_remaining, contiguous));

            memcpy(dst + copied, cur->buffer + buf_pos, chunk);
            copied      += chunk;
            cur->offset += chunk;
        }
    }

    return copied;
}

uint32_t stream_cursor_remaining(struct stream_cursor *cur)
{
    return cur->total - cur->offset;
}

void SHOT_STORAGE_wait_write_task_idle(void)
{
    /* 1. Request idle confirmation from _shot_storage_write_task */
    k_sem_reset(&shot_storage_write_idle_sem);
    atomic_set(&_shot_storage_write_idle_requested, 1);

    /* 
     * 2. Wake the write task so it can see the idle request.
     *    If it's already waiting on shot_capture_sem, this
     *    lets it loop once and post the idle confirmation.
     */
    k_sem_give(&write_shotdata_sem);  /* nudge write task so it can see the request */

    /* 3. Wait for confirmation that write task is idle */
    k_sem_take(&shot_storage_write_idle_sem, K_FOREVER);
    atomic_set(&_shot_storage_write_idle_requested, 0);
}

void SHOT_STORAGE_request_read(read_request_source_e source)
{
    atomic_set(&_read_request_source, (atomic_val_t)source);
    k_sem_give(&read_shotdata_request);
}

int SHOT_STORAGE_wait_buf_ready(void)
{
    return k_sem_take(&read_shotdata_buf_ready, K_FOREVER);
}

void SHOT_STORAGE_read_done(void)
{
    k_sem_give(&read_shotdata_consumed);
}


/*--------------------------------------------------------------------------------------------------
 * PRIVATE FUNCTION DEFINITIONS
 *------------------------------------------------------------------------------------------------*/
static void _shot_storage_files_init (void)
{
    /* Read tracking info from internal flash */
    _shot_storage_tracking_read();

    /* Read file headers for all valid slots based on tracking */
    if (_shot_storage_tracking.file_cnt > 0) {
        for (uint32_t i = 0; i < _shot_storage_tracking.file_cnt; i++) {
            uint32_t idx = (_shot_storage_tracking.write_idx - _shot_storage_tracking.file_cnt + i + SHOT_STORAGE_MAX_FILES) % SHOT_STORAGE_MAX_FILES;
            uint32_t addr = _SHOT_STORAGE_START_ADDR + (idx * _SHOT_STORAGE_TOTAL_FILE_SIZE);

            if (NVM_EXT_read(addr, (uint8_t *) &_shot_storage_files[idx], sizeof(SHOT_STORAGE_file_s))) {
                PRINT(CRITICAL, "Failure to read file header from slot %u\r\n", idx);
            }
        }
    }

    PRINT(INFO, "Shot storage: %u files, write_idx=%u, total_written=%u\r\n",
          _shot_storage_tracking.file_cnt, _shot_storage_tracking.write_idx,
          _shot_storage_tracking.total_written);
}

static bool _shot_storage_is_empty (uint8_t *p_data, uint32_t len)
{
    uint32_t i;

    /* Loop through the read data. */
    for (i = 0; i < len; i++) {
        /* Scan for empty space. */
        if (p_data[i] != 0xFF) {
            return false;
        }
    }

    return true;
}

static ERRNO_e _shot_storage_nvm_page_write (uint32_t addr, uint8_t *p_src, uint32_t rem_len)
{
    /* Write enable command must set before performing a write. */
    p_src[0] = 0x06;

    if (SPI_xfer (p_src, 1, NULL, 0) != ERRNO_SUCCESS)
    {
        PRINT (CRITICAL, "SPI failure transferring write enable.\r\n");
        return ERRNO_SPI_FAILURE;
    }

    /* Page program opcode is 0x02. */
    p_src[0] = 0x02;
    p_src[1] = _SHOT_STORAGE_SPLIT_BYTE (addr, 2);      /* A23 - A16 */
    p_src[2] = _SHOT_STORAGE_SPLIT_BYTE (addr, 1);      /* A15 -  A8 */
    p_src[3] = _SHOT_STORAGE_SPLIT_BYTE (addr, 0);      /*  A7 -  A0 */

    /* 256 bytes per page + 4 bytes for op code and address. */
    rem_len = rem_len + 4;

    /* Perform write. */
    if (SPI_xfer (p_src, rem_len, NULL, 0) != ERRNO_SUCCESS)
    {
        return ERRNO_SPI_FAILURE;
    }

    /* Wait for write to complete. */
    _shot_storage_nvm_wait_til_ready ();

    return ERRNO_SUCCESS;
}
 
static void _shot_storage_nvm_wait_til_ready (void)
{ 
    uint8_t op_code = 0x05;
    uint8_t rx;

    /* Wait for erase to complete by polling READYn/BUSY bit. */
    do
    {
        if (SPI_xfer (&op_code, 1, &rx, 1) != ERRNO_SUCCESS)
        {
            PRINT (CRITICAL, "SPI failure polling ready bit.\r\n");
        }
    } while ((rx & _SHOT_STORAGE_NVM_STATUS_BUSY_MASK) == _SHOT_STORAGE_NVM_STATUS_BUSY_MASK);
}

static void _shot_storage_write_task(void *p1, void *p2, void *p3)
{
    ERRNO_e err = ERRNO_SUCCESS;
    uint8_t *p_data;

    while (1) {
        int ret = k_sem_take(&write_shotdata_sem, K_FOREVER);
        if (ret != 0) {
            /* if k_sem_reset is called ret is -EAGAIN, in that case just continue and wait for the next give */
            continue;
        }

        /* Check if someone requested us to confirm idle */
        if (atomic_get(&_shot_storage_write_idle_requested)) {
            k_sem_give(&shot_storage_write_idle_sem);
            continue;
        }

        /* Create a file for the current shot data */
        err = SHOT_STORAGE_file_create();
        if (err == ERRNO_SUCCESS) {
            SHOT_STORAGE_write();
            err = SHOT_STORAGE_get_buf(shot_data.size, &p_data);
            if (err != ERRNO_SUCCESS) {
                SHOT_CAPTURE_enable(false);
                PRINT(CRITICAL, "Cannot enable shot capture, shot storage buffer allocation failed!\r\n");
            } else {
                k_mutex_lock(&FIFO_ring_buf_mutex, K_FOREVER);
                FIFO_ring_buf.buf = p_data;
                k_mutex_unlock(&FIFO_ring_buf_mutex);
            }
        } else {
            SHOT_CAPTURE_enable(false);
            if (err == ERRNO_EXCEED_FILE_SIZE) {
                PRINT(CRITICAL, "Cannot enable shot capture, shot data size exceed max file size!\r\n");
            } else {
                PRINT(CRITICAL, "Cannot enable shot capture, shot storage filename size exceed limit!\r\n");
            }
        }

        k_mutex_lock(&shot_data_mutex, K_FOREVER);
        shot_data.marked = false;
        shot_data.ready  = false;
        shot_data.saved  = true;
        k_mutex_unlock(&shot_data_mutex);
    }
}

static void _shot_storage_read_task(void *p1, void *p2, void *p3)
{
    SHOT_STORAGE_file_s file;

    while (1) {
        /* 1. Wait for a read request from CLI or BLE */
        int ret = k_sem_take(&read_shotdata_request, K_FOREVER);
        if (ret != 0) {
            continue;
        }

        read_request_source_e source = (read_request_source_e)atomic_get(&_read_request_source);

        /* Disable shot capture and confrim idle (serialized with FSM via mutex)*/
        SHOT_CAPTURE_task_disable_and_wait();

        /* Wait for write task to finish and confirm idle */
        SHOT_STORAGE_wait_write_task_idle();

        /* read last shot data from flash to _shot_storage_buf */
        uint8_t *p_data = SHOT_STORAGE_last_read(&file);

        if (NULL == p_data) {
            PRINT(INFO, "No file found!\r\n");

            if (source == READ_REQ_BLE) {
                /* Notify phone that no file was found */
                BLE_shot_data_not_found();
            } else if (source == READ_REQ_CLI) {
                /* reset read_shotdata_buf_ready so --read command does not block */
                k_sem_reset(&read_shotdata_buf_ready);
            }
            /* Resume capture. FSM will call SHOT_CAPTURE_enable(true) if appropriate */
            SHOT_CAPTURE_task_resume();
            continue;
        }

        if (source == READ_REQ_BLE) {
            /* Notify phone that shot data is ready to be read */
            BLE_shot_data_ready(_shot_storage_buf,
                                sizeof(SHOT_STORAGE_file_s) + file.shot_data_size);
        } else {
            /* Notify the CLI read command that buffer is ready */
            k_sem_give(&read_shotdata_buf_ready);
        }

        /* Wait for the requester to finish consuming the data */
        if (source == READ_REQ_BLE) {
            /*
             * BLE path: wait for phone to confirm it received the data.
             * Use a timeout so we don't block forever if phone is unresponsive.
             * If BLE disconnects during the read, the disconnect callback will
             * call k_sem_reset, causing k_sem_take to return -EAGAIN.
             */
            ret = k_sem_take(&read_shotdata_consumed, K_MSEC(10000));
            if (ret == 0) {
                /* Phone confirmed receipt — safe to erase */
                SHOT_STORAGE_last_erase();
                PRINT(INFO, "Power cnt: %u, time: %.3f s erased...\r\n", file.power_cnt, (float) file.time_of_shot / 1000.0);
            } else {
                /* Timeout or disconnect — do NOT erase so phone can retry */
                BLE_shot_data_timeout();
                PRINT(WARNING, "BLE read not confirmed (ret=%d), keeping flash data for retry\r\n", ret);
            }
        } else if (source == READ_REQ_CLI) {
            /*
             * CLI path: CLI handler will take read_shotdata_buf_ready,
             * print the data, call SHOT_STORAGE_last_erase(), then
             * give read_shotdata_consumed.
             */
            k_sem_take(&read_shotdata_consumed, K_FOREVER);
        }

        atomic_set(&_read_request_source, READ_REQ_NONE);
        /* Resume shot capture. FSM will re-enable if appropriate */
        SHOT_CAPTURE_task_resume();
    }
}

static void _shot_storage_tracking_read (void)
{
    ERRNO_e err;
    SHOT_STORAGE_tracking_s trk;

    /* allocate two sectors to store this tracking file because it updates for every shot */
    err = FS_open_or_create(SHOT_STORAGE_TRACKING_FILENAME, &_shot_storage_tracking_file,
                            sizeof(SHOT_STORAGE_tracking_s), (FS_SECTOR_SIZE / (sizeof(SHOT_STORAGE_tracking_s) + 
                                                              FS_FILE_INDEX_SIZE +
                                                              FS_FILE_CHECKSUM_SIZE)) * 
                                                              (SHOT_STORAGE_TRACKING_FILE_SECTOR - 1)); 

    if (err == ERRNO_SUCCESS) {
        err = FS_read(_shot_storage_tracking_file, (uint8_t *)&trk, 0, sizeof(SHOT_STORAGE_tracking_s));
        if (err != ERRNO_SUCCESS) {
            /* First boot or corrupt — write defaults */
            _shot_storage_tracking_write();
            return;
        }

        /* Validate read data */
        if (trk.write_idx >= SHOT_STORAGE_MAX_FILES || trk.file_cnt > SHOT_STORAGE_MAX_FILES) {
            PRINT(WARNING, "Invalid tracking data, resetting to defaults\r\n");
            _shot_storage_tracking_write();
            return;
        }

        _shot_storage_tracking = trk;
    } else if (err == ERRNO_SIZE_MISMATCH) {
        /* Struct size changed — delete old file and recreate */
        err = FS_delete(SHOT_STORAGE_TRACKING_FILENAME);
        if (err != ERRNO_SUCCESS) {
            PRINT(ERROR, "Failure to delete old tracking file.\r\n");
            return;
        }

        err = FS_open_or_create(SHOT_STORAGE_TRACKING_FILENAME, &_shot_storage_tracking_file,
                        sizeof(SHOT_STORAGE_tracking_s), (FS_SECTOR_SIZE / (sizeof(SHOT_STORAGE_tracking_s) + 
                                                          FS_FILE_INDEX_SIZE +
                                                          FS_FILE_CHECKSUM_SIZE)) * 
                                                          (SHOT_STORAGE_TRACKING_FILE_SECTOR - 1)); 
        if (err != ERRNO_SUCCESS) {
            PRINT(ERROR, "Failure to create tracking file.\r\n");
        } else {
            PRINT(INFO, "New shot storage tracking file created successfully\r\n");
        }

        _shot_storage_tracking_write();
    } else {
        _shot_storage_tracking_write();
    }
}

static ERRNO_e _shot_storage_tracking_write (void)
{
    ERRNO_e err;

    err = FS_write(_shot_storage_tracking_file, (uint8_t *)&_shot_storage_tracking);
    if (err != ERRNO_SUCCESS) {
        PRINT(ERROR, "Failed to write shot storage tracking to file. Errro: %d \r\n", err);
    }

    return err;
}