#include "common/constants.h"
#include "io/nvs_handler.h"
#include "nvs_flash.h"
#include "string.h"
#include "esp_log.h"

#define TAG "NVSStorageHandler"

NVSStorageHandler::NVSStorageHandler(
  const char* name_space
) : name_space(name_space) { }

NVSStorageHandler::~NVSStorageHandler() {
  if (this->nvsHandle != 0) {
    nvs_close(this->nvsHandle);
    this->nvsHandle = 0;
  }
}

void NVSStorageHandler::init(const EventGroupHandle_t& system_event_group) {
  esp_err_t err = nvs_flash_init();
  if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
    ESP_LOGW(TAG, "NVS Flash init failed with %s, erasing and retrying...", esp_err_to_name(err));
    nvs_flash_erase();
    err = nvs_flash_init();
  }
  ESP_ERROR_CHECK(err);

  err = nvs_open(this->name_space, NVS_READWRITE, &this->nvsHandle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) opening NVS handle!", esp_err_to_name(err));
  } else {
    if (system_event_group != nullptr) {
      xEventGroupSetBits(system_event_group, STORAGE_READY_BIT);
    }
    ESP_LOGI(TAG, "NVS Storage initialized.");
  }
}

void NVSStorageHandler::put(const std::string& key, const std::string& value) {
  if (this->nvsHandle == 0) {
    ESP_LOGE(TAG, "NVS handle is not initialized.");
    return;
  }

  esp_err_t err = nvs_set_str(this->nvsHandle, key.c_str(), value.c_str());
  if (err == ESP_ERR_NVS_NOT_ENOUGH_SPACE) {
    ESP_LOGE(TAG, "NVS storage full! Cannot store key '%s'. Consider clearing old data.", key.c_str());
    ESP_LOGI(TAG, "Attempting to erase and retry...");
    // Try to commit pending changes first
    nvs_commit(this->nvsHandle);
    // Retry once
    err = nvs_set_str(this->nvsHandle, key.c_str(), value.c_str());
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Retry failed: %s", esp_err_to_name(err));
      return;
    }
  } else if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) setting key-value pair in NVS!", esp_err_to_name(err));
    return;
  }

  // Commit successful write
  err = nvs_commit(this->nvsHandle);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) committing to NVS!", esp_err_to_name(err));
  }
}

std::string NVSStorageHandler::get(const std::string& key) const {
  if (this->nvsHandle == 0) {
    ESP_LOGE(TAG, "NVS handle is not initialized.");
    return "";
  }

  size_t required_size = 0;
  esp_err_t err = nvs_get_str(this->nvsHandle, key.c_str(), nullptr, &required_size);
  if (err == ESP_ERR_NVS_NOT_FOUND) {
    ESP_LOGW(TAG, "Key %s not found in NVS.", key.c_str());
    return "";
  } else if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) getting size for key %s from NVS!", esp_err_to_name(err), key.c_str());
    return "";
  }

  // Allocate string buffer with correct size (includes null terminator)
  std::string value(required_size - 1, '\0');
  err = nvs_get_str(this->nvsHandle, key.c_str(), &value[0], &required_size);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) getting value for key %s from NVS!", esp_err_to_name(err), key.c_str());
    return "";
  }

  return value;
}

NVSIteratorGuard NVSStorageHandler::create_iterator() const {
  nvs_iterator_t it = nullptr;
  esp_err_t err = nvs_entry_find(NVS_DEFAULT_PART_NAME, this->name_space, NVS_TYPE_ANY, &it);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) creating NVS iterator!", esp_err_to_name(err));
    return NVSIteratorGuard(nullptr, err);
  }

  return NVSIteratorGuard(it, ESP_OK);
}

esp_err_t NVSStorageHandler::process_all(KeyValueProcessor processor, void* arg) const {
  NVSIteratorGuard iterator_guard = this->create_iterator();
  if (!iterator_guard.is_valid()) {
    return iterator_guard.get_error();
  }
  const nvs_iterator_t& it = iterator_guard.get_iterator();

  for (; it != NULL; iterator_guard.advance_iter()) {
    nvs_entry_info_t info;
    esp_err_t err = nvs_entry_info(it, &info);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Error (%s) getting NVS entry info!", esp_err_to_name(err));
      return err;
    }
    nvs_handle_t temp_handle;
    err = nvs_open(this->name_space, NVS_READONLY, &temp_handle);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Error (%s) opening NVS handle for reading!", esp_err_to_name(err));
      return err;
    }

    // call the processor with the key and value
    std::string key_str = info.key;
    processor(arg, key_str, this->get(key_str));
  }
  return ESP_OK;
}
esp_err_t NVSStorageHandler::process_filtered(const std::string& key_prefix, KeyValueProcessor processor, void* arg) const {
  NVSIteratorGuard iterator_guard = this->create_iterator();
  if (!iterator_guard.is_valid()) {
    return iterator_guard.get_error();
  }
  const nvs_iterator_t& it = iterator_guard.get_iterator();

  for (; it != NULL; iterator_guard.advance_iter()) {
    nvs_entry_info_t info;
    esp_err_t err = nvs_entry_info(it, &info);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Error (%s) getting NVS entry info!", esp_err_to_name(err));
      return err;
    }
    // check if the key matches the prefix
    if (strncmp(info.key, key_prefix.c_str(), key_prefix.length()) == 0) {
      nvs_handle_t temp_handle;
      err = nvs_open(this->name_space, NVS_READONLY, &temp_handle);
      if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error (%s) opening NVS handle for reading!", esp_err_to_name(err));
        return err;
      }
      // call the processor with the key and value
      processor(arg, std::string(info.key), this->get(std::string(info.key)));
    }
  }
  return ESP_OK;
}

esp_err_t NVSStorageHandler::process_filtered(FilterFunc filter_func, KeyValueProcessor processor, void* arg) const {
  NVSIteratorGuard iterator_guard = this->create_iterator();
  if (!iterator_guard.is_valid()) {
    return iterator_guard.get_error();
  }
  const nvs_iterator_t& it = iterator_guard.get_iterator();

  for (; it != NULL; iterator_guard.advance_iter()) {
    nvs_entry_info_t info;
    esp_err_t err = nvs_entry_info(it, &info);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Error (%s) getting NVS entry info!", esp_err_to_name(err));
      return err;
    }
    // check if the key matches the filter function
    std::string key_str(info.key);
    if (filter_func(key_str)) {
      nvs_handle_t temp_handle;
      err = nvs_open(this->name_space, NVS_READONLY, &temp_handle);
      if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error (%s) opening NVS handle for reading!", esp_err_to_name(err));
        return err;
      }
      // call the processor with the key and value
      processor(arg, key_str, this->get(key_str));
    }
  }
  return ESP_OK;
}

void NVSStorageHandler::remove(const std::string& key) {
  if (this->nvsHandle == 0) {
    ESP_LOGE(TAG, "NVS handle is not initialized.");
    return;
  }

  esp_err_t err = nvs_erase_key(this->nvsHandle, key.c_str());
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Error (%s) deleting key %s from NVS!", esp_err_to_name(err), key.c_str());
  } else {
    nvs_commit(this->nvsHandle);
    ESP_LOGI(TAG, "Key %s deleted from NVS.", key.c_str());
  }
}