#include <string.h>
#include <zephyr/kernel.h>
#include <zephyr/drivers/flash.h>
#include <zephyr/storage/flash_map.h>
#include <zephyr/fs/nvs.h>

#include "device_data_model.h"

#ifdef CONFIG_LOG
#include <zephyr/logging/log.h>
LOG_MODULE_REGISTER(data_storage, CONFIG_LOG_DEFAULT_LEVEL);
#endif

#define NVS_PARTITION            app_storage_partition
#define NVS_PARTITION_DEVICE     FIXED_PARTITION_DEVICE(NVS_PARTITION)
#define NVS_PARTITION_OFFSET     FIXED_PARTITION_OFFSET(NVS_PARTITION)
#define NVS_PARTITION_SIZE       FIXED_PARTITION_SIZE(NVS_PARTITION)

#define DEFAULT_CAL_RAW_A   1637
#define DEFAULT_CAL_REAL_C  129
#define DEFAULT_CAL_RAW_B   2870
#define DEFAULT_CAL_REAL_D  83
#define DEFAULT_CAL_SLOPE   -0.037307f
#define DEFAULT_CAL_OFFSET  190.072f

/* Single mutex guarding every NVS access in this file. K_FOREVER everywhere
 * (a prior incident was caused by silent timeout failures on this lock). */
static K_MUTEX_DEFINE(nvs_mtx);

static struct nvs_fs fs = {
	.flash_device = NVS_PARTITION_DEVICE,
	.offset       = NVS_PARTITION_OFFSET,
};

static bool                  nvs_mounted;
static uint16_t              max_storage_records;
static struct active_session session;

static int      stream_idx;
static uint16_t stream_remaining;

static int nvs_mount_safe(void)
{
	if (nvs_mounted) {
		return 0;
	}

	if (!device_is_ready(fs.flash_device)) {
#ifdef CONFIG_LOG
		LOG_ERR("Flash device %s not ready", fs.flash_device->name);
#endif
		return -ENODEV;
	}

	struct flash_pages_info info;
	int rc = flash_get_page_info_by_offs(fs.flash_device,
					     fs.offset, &info);
	if (rc) {
		return rc;
	}

	fs.sector_size  = info.size;
	fs.sector_count = NVS_PARTITION_SIZE / fs.sector_size;

	rc = nvs_mount(&fs);
	if (rc) {
#ifdef CONFIG_LOG
		LOG_ERR("NVS mount failed (%d)", rc);
#endif
		return rc;
	}

	max_storage_records = ((fs.sector_count - 1) * fs.sector_size) /
			      (sizeof(struct measurement_payload) + 8);

	nvs_mounted = true;

#ifdef CONFIG_LOG
	LOG_INF("NVS mounted. Capacity: %u records", max_storage_records);
#endif
	return 0;
}

int data_storage_init(void)
{
	k_mutex_lock(&nvs_mtx, K_FOREVER);

	int rc = nvs_mount_safe();
	if (rc) {
		k_mutex_unlock(&nvs_mtx);
		return rc;
	}

	int rd = nvs_read(&fs, NVS_SESSION_META_ID,
			  &session, sizeof(session));
	if (rd <= 0) {
		memset(&session, 0, sizeof(session));
	}

	k_mutex_unlock(&nvs_mtx);
	return 0;
}

int calibration_load(struct calibration_params *out)
{
	if (!out) {
		return -EINVAL;
	}

	k_mutex_lock(&nvs_mtx, K_FOREVER);

	int rc = nvs_mount_safe();
	if (rc) {
		k_mutex_unlock(&nvs_mtx);
		return rc;
	}

	rc = nvs_read(&fs, NVS_CALIBRATION_ID, out,
		      sizeof(struct calibration_params));

	if (rc <= 0) {
		out->raw_a    = DEFAULT_CAL_RAW_A;
		out->real_c   = DEFAULT_CAL_REAL_C;
		out->raw_b    = DEFAULT_CAL_RAW_B;
		out->real_d   = DEFAULT_CAL_REAL_D;
		out->slope    = DEFAULT_CAL_SLOPE;
		out->offset   = DEFAULT_CAL_OFFSET;
		out->is_valid = true;
#ifdef CONFIG_LOG
		LOG_WRN("No calibration found, using defaults");
#endif
	}

	k_mutex_unlock(&nvs_mtx);
	return 0;
}

int calibration_save(const struct calibration_params *p)
{
	if (!p) {
		return -EINVAL;
	}

	k_mutex_lock(&nvs_mtx, K_FOREVER);

	int rc = nvs_mount_safe();
	if (rc) {
		k_mutex_unlock(&nvs_mtx);
		return rc;
	}

	rc = nvs_write(&fs, NVS_CALIBRATION_ID, p,
		       sizeof(struct calibration_params));
	k_mutex_unlock(&nvs_mtx);

	if (rc < 0) {
#ifdef CONFIG_LOG
		LOG_ERR("Calibration save failed (%d)", rc);
#endif
		return rc;
	}
	return 0;
}

int session_start(void)
{
	k_mutex_lock(&nvs_mtx, K_FOREVER);

	int rc = nvs_mount_safe();
	if (rc) {
		k_mutex_unlock(&nvs_mtx);
		return rc;
	}

	/* Auto-wipe prior measurements: with a single fixed session id, leaving
	 * old ATEs around would make nvs_read_hist return them on the next
	 * stream and mix two recordings. SRS §6, user-selected option. */
	(void)nvs_delete(&fs, NVS_MEASUREMENT_ID);

	session.record_count = 0;
	session.is_active    = 1;

	rc = nvs_write(&fs, NVS_SESSION_META_ID,
		       &session, sizeof(session));
	if (rc < 0) {
		k_mutex_unlock(&nvs_mtx);
#ifdef CONFIG_LOG
		LOG_ERR("Session meta write failed (%d)", rc);
#endif
		return rc;
	}

	uint16_t zero = 0;
	(void)nvs_write(&fs, NVS_ACTIVE_COUNT_ID, &zero, sizeof(zero));

	k_mutex_unlock(&nvs_mtx);

#ifdef CONFIG_LOG
	LOG_INF("Session started (auto-wiped prior data)");
#endif
	return 0;
}

void session_stop(void)
{
	k_mutex_lock(&nvs_mtx, K_FOREVER);

	if (!nvs_mounted) {
		k_mutex_unlock(&nvs_mtx);
		return;
	}

	session.is_active = 0;
	(void)nvs_write(&fs, NVS_SESSION_META_ID,
			&session, sizeof(session));

	k_mutex_unlock(&nvs_mtx);

#ifdef CONFIG_LOG
	LOG_INF("Session stopped. Records: %u", session.record_count);
#endif
}

void session_recover(void)
{
	k_mutex_lock(&nvs_mtx, K_FOREVER);

	if (!nvs_mounted || session.is_active == 0) {
		k_mutex_unlock(&nvs_mtx);
		return;
	}

	/* SRS §8: NVS_ACTIVE_COUNT_ID is the per-sample anchor — it always
	 * holds the live count of the last successful measurement_put. */
	uint16_t recovered = 0;
	int rc = nvs_read(&fs, NVS_ACTIVE_COUNT_ID,
			  &recovered, sizeof(recovered));
	if (rc > 0) {
		session.record_count = recovered;
	}

	session.is_active = 0;
	(void)nvs_write(&fs, NVS_SESSION_META_ID,
			&session, sizeof(session));

	k_mutex_unlock(&nvs_mtx);

#ifdef CONFIG_LOG
	LOG_INF("Recovered session: %u records", session.record_count);
#endif
}

uint16_t session_get_record_count(void)
{
	return session.record_count;
}

bool session_is_active(void)
{
	return session.is_active != 0;
}

int measurement_put(const struct measurement_payload *m)
{
	if (!m) {
		return -EINVAL;
	}

	if (session.record_count >= max_storage_records) {
#ifdef CONFIG_LOG
		LOG_ERR("Flash full (%u records)", session.record_count);
#endif
		return -ENOSPC;
	}

	k_mutex_lock(&nvs_mtx, K_FOREVER);

	int rc = nvs_write(&fs, NVS_MEASUREMENT_ID, m, sizeof(*m));
	if (rc < 0) {
		k_mutex_unlock(&nvs_mtx);
		return rc;
	}

	session.record_count++;

	/* SRS §8: rewrite the live counter in the same critical section as the
	 * sample. NVS GC keeps only the latest value at this ID, so this is
	 * the recovery anchor that survives an unclean shutdown. */
	uint16_t cnt = session.record_count;
	(void)nvs_write(&fs, NVS_ACTIVE_COUNT_ID, &cnt, sizeof(cnt));

	k_mutex_unlock(&nvs_mtx);
	return 0;
}

int stream_begin(uint8_t session_id)
{
	if (session_id != FIXED_SESSION_ID || session.record_count == 0) {
		return -ENOENT;
	}

	stream_remaining = session.record_count;
	stream_idx       = 0;

#ifdef CONFIG_LOG
	LOG_INF("Stream begin: up to %u records", stream_remaining);
#endif
	return 0;
}

int stream_next(struct measurement_payload *out)
{
	if (!out || stream_remaining == 0) {
		return -ENOENT;
	}

	k_mutex_lock(&nvs_mtx, K_FOREVER);
	int rc = nvs_read_hist(&fs, NVS_MEASUREMENT_ID, out,
			       sizeof(*out), stream_idx);
	k_mutex_unlock(&nvs_mtx);

	/* Tolerate short reads: NVS GC may have dropped older ATEs when the
	 * partition wrapped, so record_count is an upper bound (SRS §6.2). */
	if (rc <= 0) {
		stream_remaining = 0;
		return -ENOENT;
	}

	stream_idx++;
	stream_remaining--;
	return 0;
}

bool stream_has_more(void)
{
	return stream_remaining > 0;
}

void erase_all(void)
{
	k_mutex_lock(&nvs_mtx, K_FOREVER);

	if (nvs_mounted) {
		(void)nvs_delete(&fs, NVS_MEASUREMENT_ID);
		(void)nvs_delete(&fs, NVS_SESSION_META_ID);
		(void)nvs_delete(&fs, NVS_ACTIVE_COUNT_ID);
	}

	memset(&session, 0, sizeof(session));
	stream_idx       = 0;
	stream_remaining = 0;

	k_mutex_unlock(&nvs_mtx);

#ifdef CONFIG_LOG
	LOG_INF("All data erased");
#endif
}
