// Copyright 2020 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

#define DUMP_KVS_STATE_TO_FILE 0
#define USE_MEMORY_BUFFER 1

#include "pw_kvs/key_value_store.h"

#include <array>
#include <cstdio>
#include <cstring>

#include "pw_span/span.h"

#if DUMP_KVS_STATE_TO_FILE
#include <vector>
#endif  // DUMP_KVS_STATE_TO_FILE

#include "pw_assert/check.h"
#include "pw_bytes/array.h"
#include "pw_checksum/crc16_ccitt.h"
#include "pw_kvs/crc16_checksum.h"
#include "pw_kvs/fake_flash_memory.h"
#include "pw_kvs/flash_memory.h"
#include "pw_kvs/internal/entry.h"
#include "pw_kvs_private/config.h"
#include "pw_log/log.h"
#include "pw_status/status.h"
#include "pw_string/string_builder.h"
#include "pw_unit_test/framework.h"

namespace pw::kvs {
namespace {

using internal::EntryHeader;

constexpr size_t kMaxEntries = 256;
constexpr size_t kMaxUsableSectors = 256;

// This is a self contained flash unit with both memory and a single partition.
template <uint32_t kSectorSizeBytes, uint16_t kSectorCount>
struct FlashWithPartitionFake {
  // Default to 16 byte alignment, which is common in practice.
  FlashWithPartitionFake() : FlashWithPartitionFake(16) {}
  FlashWithPartitionFake(size_t alignment_bytes)
      : memory(alignment_bytes),
        partition(&memory, 0, static_cast<uint32_t>(memory.sector_count())) {}

  FakeFlashMemoryBuffer<kSectorSizeBytes, kSectorCount> memory;
  FlashPartition partition;

 public:
#if DUMP_KVS_STATE_TO_FILE
  Status Dump(const char* filename) {
    std::FILE* out_file = std::fopen(filename, "w+");
    if (out_file == nullptr) {
      PW_LOG_ERROR("Failed to dump to %s", filename);
      return Status::DataLoss();
    }
    std::vector<std::byte> out_vec(memory.size_bytes());
    Status status =
        memory.Read(0, span<std::byte>(out_vec.data(), out_vec.size()));
    if (status != OkStatus()) {
      fclose(out_file);
      return status;
    }

    size_t written =
        std::fwrite(out_vec.data(), 1, memory.size_bytes(), out_file);
    if (written != memory.size_bytes()) {
      PW_LOG_ERROR("Failed to dump to %s, written=%u",
                   filename,
                   static_cast<unsigned>(written));
      status = Status::DataLoss();
    } else {
      PW_LOG_INFO("Dumped to %s", filename);
      status = OkStatus();
    }

    fclose(out_file);
    return status;
  }
#else
  Status Dump(const char*) { return OkStatus(); }
#endif  // DUMP_KVS_STATE_TO_FILE
};

typedef FlashWithPartitionFake<4 * 128 /*sector size*/, 6 /*sectors*/> Flash;

FakeFlashMemoryBuffer<1024, 60> large_test_flash(8);
FlashPartition large_test_partition(
    &large_test_flash,
    0,
    static_cast<uint32_t>(large_test_flash.sector_count()));

constexpr std::array<const char*, 3> keys{"TestKey1", "Key2", "TestKey3"};

ChecksumCrc16 checksum;
// For KVS magic value always use a random 32 bit integer rather than a
// human readable 4 bytes. See pw_kvs/format.h for more information.
constexpr EntryFormat default_format{.magic = 0xa6cb3c16,
                                     .checksum = &checksum};

}  // namespace

TEST(InitCheck, TooFewSectors) {
  // Use test flash with 1 x 4k sectors, 16 byte alignment
  FakeFlashMemoryBuffer<4 * 1024, 1> test_flash(16);
  FlashPartition test_partition(
      &test_flash, 0, static_cast<uint32_t>(test_flash.sector_count()));

  // For KVS magic value always use a random 32 bit integer rather than a
  // human readable 4 bytes. See pw_kvs/format.h for more information.
  constexpr EntryFormat format{.magic = 0x89bb14d2, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&test_partition,
                                                          format);

  EXPECT_EQ(kvs.Init(), Status::FailedPrecondition());
}

TEST(InitCheck, ZeroSectors) {
  // Use test flash with 1 x 4k sectors, 16 byte alignment
  FakeFlashMemoryBuffer<4 * 1024, 1> test_flash(16);

  // Set FlashPartition to have 0 sectors.
  FlashPartition test_partition(&test_flash, 0, 0);

  // For KVS magic value always use a random 32 bit integer rather than a
  // human readable 4 bytes. See pw_kvs/format.h for more information.
  constexpr EntryFormat format{.magic = 0xd1da57c1, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&test_partition,
                                                          format);

  EXPECT_EQ(kvs.Init(), Status::FailedPrecondition());
}

TEST(InitCheck, TooManySectors) {
  // Use test flash with 1 x 4k sectors, 16 byte alignment
  FakeFlashMemoryBuffer<4 * 1024, 5> test_flash(16);

  // Set FlashPartition to have 0 sectors.
  FlashPartition test_partition(
      &test_flash, 0, static_cast<uint32_t>(test_flash.sector_count()));

  // For KVS magic value always use a random 32 bit integer rather than a
  // human readable 4 bytes. See pw_kvs/format.h for more information.
  constexpr EntryFormat format{.magic = 0x610f6d17, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, 2> kvs(&test_partition, format);

  EXPECT_EQ(kvs.Init(), Status::FailedPrecondition());
}

TEST(InMemoryKvs, WriteOneKeyMultipleTimes) {
  // Create and erase the fake flash. It will persist across reloads.
  Flash flash;
  PW_TEST_ASSERT_OK(flash.partition.Erase());

  int num_reloads = 2;
  for (int reload = 0; reload < num_reloads; ++reload) {
    PW_LOG_DEBUG("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
    PW_LOG_DEBUG("xxx                                      xxxx");
    PW_LOG_DEBUG("xxx               Reload %2d              xxxx", reload);
    PW_LOG_DEBUG("xxx                                      xxxx");
    PW_LOG_DEBUG("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");

    // Create and initialize the KVS. For KVS magic value always use a random 32
    // bit integer rather than a human readable 4 bytes. See pw_kvs/format.h for
    // more information.
    constexpr EntryFormat format{.magic = 0x83a9257, .checksum = nullptr};
    KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                            format);
    PW_TEST_ASSERT_OK(kvs.Init());

    // Write the same entry many times.
    const char* key = "abcd";
    const size_t num_writes = 99;
    uint32_t written_value;
    EXPECT_EQ(kvs.size(), (reload == 0) ? 0 : 1u);
    for (uint32_t i = 0; i < num_writes; ++i) {
      PW_LOG_DEBUG(
          "PUT #%zu for key %s with value %zu", size_t(i), key, size_t(i));

      written_value = i + 0xfc;  // Prevent accidental pass with zero.
      PW_TEST_EXPECT_OK(kvs.Put(key, written_value));
      EXPECT_EQ(kvs.size(), 1u);
    }

    // Verify that we can read the value back.
    PW_LOG_DEBUG("GET final value for key: %s", key);
    uint32_t actual_value;
    PW_TEST_EXPECT_OK(kvs.Get(key, &actual_value));
    EXPECT_EQ(actual_value, written_value);

    char fname_buf[64] = {'\0'};
    snprintf(&fname_buf[0],
             sizeof(fname_buf),
             "WriteOneKeyMultipleTimes_%d.bin",
             reload);
    ASSERT_EQ(OkStatus(), flash.Dump(fname_buf));
  }
}

TEST(InMemoryKvs, WritingMultipleKeysIncreasesSize) {
  // Create and erase the fake flash.
  Flash flash;
  PW_TEST_ASSERT_OK(flash.partition.Erase());

  // Create and initialize the KVS. For KVS magic value always use a random 32
  // bit integer rather than a human readable 4 bytes. See pw_kvs/format.h for
  // more information.
  constexpr EntryFormat format{.magic = 0x2ed3a058, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          format);
  PW_TEST_ASSERT_OK(kvs.Init());

  // Write the same entry many times.
  const size_t num_writes = 10;
  EXPECT_EQ(kvs.size(), 0u);
  for (size_t i = 0; i < num_writes; ++i) {
    StringBuffer<150> key;
    key << "key_" << i;
    PW_LOG_DEBUG("PUT #%zu for key %s with value %zu", i, key.c_str(), i);

    size_t value = i + 77;  // Prevent accidental pass with zero.
    PW_TEST_EXPECT_OK(kvs.Put(key.view(), value));
    EXPECT_EQ(kvs.size(), i + 1);
  }
  ASSERT_EQ(OkStatus(), flash.Dump("WritingMultipleKeysIncreasesSize.bin"));
}

TEST(InMemoryKvs, WriteAndReadOneKey) {
  // Create and erase the fake flash.
  Flash flash;
  PW_TEST_ASSERT_OK(flash.partition.Erase());

  // Create and initialize the KVS.
  // For KVS magic value always use a random 32 bit integer rather than a
  // human readable 4 bytes. See pw_kvs/format.h for more information.
  constexpr EntryFormat format{.magic = 0x5d70896, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          format);
  PW_TEST_ASSERT_OK(kvs.Init());

  // Add one entry.
  const char* key = "Key1";
  PW_LOG_DEBUG("PUT value for key: %s", key);
  uint8_t written_value = 0xDA;
  PW_TEST_ASSERT_OK(kvs.Put(key, written_value));
  EXPECT_EQ(kvs.size(), 1u);

  PW_LOG_DEBUG("GET value for key: %s", key);
  uint8_t actual_value;
  PW_TEST_ASSERT_OK(kvs.Get(key, &actual_value));
  EXPECT_EQ(actual_value, written_value);

  EXPECT_EQ(kvs.size(), 1u);
}

TEST(InMemoryKvs, WriteOneKeyValueMultipleTimes) {
  // Create and erase the fake flash.
  Flash flash;
  PW_TEST_ASSERT_OK(flash.partition.Erase());

  // Create and initialize the KVS.
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          default_format);
  PW_TEST_ASSERT_OK(kvs.Init());

  // Add one entry, with the same key and value, multiple times.
  const char* key = "Key1";
  uint8_t written_value = 0xDA;
  for (int i = 0; i < 50; i++) {
    PW_LOG_DEBUG("PUT [%d] value for key: %s", i, key);
    PW_TEST_ASSERT_OK(kvs.Put(key, written_value));
    EXPECT_EQ(kvs.size(), 1u);
  }

  PW_LOG_DEBUG("GET value for key: %s", key);
  uint8_t actual_value;
  PW_TEST_ASSERT_OK(kvs.Get(key, &actual_value));
  EXPECT_EQ(actual_value, written_value);

  // Verify that only one entry was written to the KVS.
  EXPECT_EQ(kvs.size(), 1u);
  EXPECT_EQ(kvs.transaction_count(), 1u);
  KeyValueStore::StorageStats stats = kvs.GetStorageStats();
  EXPECT_EQ(stats.reclaimable_bytes, 0u);
}

TEST(InMemoryKvs, Basic) {
  const char* key1 = "Key1";
  const char* key2 = "Key2";

  // Create and erase the fake flash.
  Flash flash;
  ASSERT_EQ(OkStatus(), flash.partition.Erase());

  // Create and initialize the KVS.
  // For KVS magic value always use a random 32 bit integer rather than a
  // human readable 4 bytes. See pw_kvs/format.h for more information.
  constexpr EntryFormat format{.magic = 0x7bf19895, .checksum = nullptr};
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          format);
  PW_TEST_ASSERT_OK(kvs.Init());

  // Add two entries with different keys and values.
  uint8_t value1 = 0xDA;
  PW_TEST_ASSERT_OK(kvs.Put(key1, as_bytes(span(&value1, sizeof(value1)))));
  EXPECT_EQ(kvs.size(), 1u);

  uint32_t value2 = 0xBAD0301f;
  PW_TEST_ASSERT_OK(kvs.Put(key2, value2));
  EXPECT_EQ(kvs.size(), 2u);

  // Verify data
  uint32_t test2;
  PW_TEST_EXPECT_OK(kvs.Get(key2, &test2));

  uint8_t test1;
  PW_TEST_ASSERT_OK(kvs.Get(key1, &test1));

  EXPECT_EQ(test1, value1);
  EXPECT_EQ(test2, value2);

  EXPECT_EQ(kvs.size(), 2u);
}

TEST(InMemoryKvs, CallingEraseTwice_NothingWrittenToFlash) {
  // Create and erase the fake flash.
  Flash flash;
  ASSERT_EQ(OkStatus(), flash.partition.Erase());

  // Create and initialize the KVS.
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          default_format);
  PW_TEST_ASSERT_OK(kvs.Init());

  const uint8_t kValue = 0xDA;
  ASSERT_EQ(OkStatus(), kvs.Put(keys[0], kValue));
  ASSERT_EQ(OkStatus(), kvs.Delete(keys[0]));

  // Compare before / after checksums to verify that nothing was written.
  const uint16_t crc = checksum::Crc16Ccitt::Calculate(flash.memory.buffer());

  EXPECT_EQ(kvs.Delete(keys[0]), Status::NotFound());

  EXPECT_EQ(crc, checksum::Crc16Ccitt::Calculate(flash.memory.buffer()));
}

class LargeEmptyInitializedKvs : public ::testing::Test {
 protected:
  LargeEmptyInitializedKvs() : kvs_(&large_test_partition, default_format) {
    PW_CHECK_OK(large_test_partition.Erase());
    PW_CHECK_OK(kvs_.Init());
  }

  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs_;
};

TEST_F(LargeEmptyInitializedKvs, Basic) {
  const uint8_t kValue1 = 0xDA;
  const uint8_t kValue2 = 0x12;
  uint8_t value;
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[0], kValue1));
  EXPECT_EQ(kvs_.size(), 1u);
  ASSERT_EQ(OkStatus(), kvs_.Delete(keys[0]));
  EXPECT_EQ(kvs_.Get(keys[0], &value), Status::NotFound());
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[1], kValue1));
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[2], kValue2));
  ASSERT_EQ(OkStatus(), kvs_.Delete(keys[1]));
  EXPECT_EQ(OkStatus(), kvs_.Get(keys[2], &value));
  EXPECT_EQ(kValue2, value);
  ASSERT_EQ(kvs_.Get(keys[1], &value), Status::NotFound());
  EXPECT_EQ(kvs_.size(), 1u);
}

TEST_F(LargeEmptyInitializedKvs, FullMaintenance) {
  const uint8_t kValue1 = 0xDA;
  const uint8_t kValue2 = 0x12;

  // Write a key and write again with a different value, resulting in a stale
  // entry from the first write.
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[0], kValue1));
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[0], kValue2));
  EXPECT_EQ(kvs_.size(), 1u);

  KeyValueStore::StorageStats stats = kvs_.GetStorageStats();
  EXPECT_EQ(stats.sector_erase_count, 0u);
  EXPECT_GT(stats.reclaimable_bytes, 0u);

  // Do regular FullMaintenance, which should not touch the sector with valid
  // data.
  EXPECT_EQ(OkStatus(), kvs_.FullMaintenance());
  stats = kvs_.GetStorageStats();
  EXPECT_EQ(stats.sector_erase_count, 0u);
  EXPECT_GT(stats.reclaimable_bytes, 0u);

  // Do aggressive FullMaintenance, which should GC the sector with valid data,
  // resulting in no reclaimable bytes and an erased sector.
  EXPECT_EQ(OkStatus(), kvs_.HeavyMaintenance());
  stats = kvs_.GetStorageStats();
  EXPECT_EQ(stats.sector_erase_count, 1u);
  EXPECT_EQ(stats.reclaimable_bytes, 0u);
}

TEST_F(LargeEmptyInitializedKvs, KeyDeletionMaintenance) {
  const uint8_t kValue1 = 0xDA;
  const uint8_t kValue2 = 0x12;
  uint8_t val = 0;

  // Write and delete a key. The key should be gone, but the size should be 1.
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[0], kValue1));
  ASSERT_EQ(kvs_.size(), 1u);
  ASSERT_EQ(OkStatus(), kvs_.Delete(keys[0]));

  // Ensure the key is indeed gone and the size is 1 before continuing.
  ASSERT_EQ(kvs_.Get(keys[0], &val), Status::NotFound());
  ASSERT_EQ(kvs_.size(), 0u);
  ASSERT_EQ(kvs_.total_entries_with_deleted(), 1u);

  KeyValueStore::StorageStats stats = kvs_.GetStorageStats();
  EXPECT_EQ(stats.sector_erase_count, 0u);
  EXPECT_GT(stats.reclaimable_bytes, 0u);

  // Do aggressive FullMaintenance, which should GC the sector with valid data,
  // resulting in no reclaimable bytes and an erased sector.
  EXPECT_EQ(OkStatus(), kvs_.HeavyMaintenance());
  stats = kvs_.GetStorageStats();
  EXPECT_EQ(stats.reclaimable_bytes, 0u);
  ASSERT_EQ(kvs_.size(), 0u);

  if (PW_KVS_REMOVE_DELETED_KEYS_IN_HEAVY_MAINTENANCE) {
    EXPECT_GT(stats.sector_erase_count, 1u);
    ASSERT_EQ(kvs_.total_entries_with_deleted(), 0u);
  } else {  // The deleted entries are only removed if that feature is enabled.
    EXPECT_EQ(stats.sector_erase_count, 1u);
    ASSERT_EQ(kvs_.total_entries_with_deleted(), 1u);
  }

  // Do it again but with 2 keys and keep one.
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[0], kValue1));
  ASSERT_EQ(OkStatus(), kvs_.Put(keys[1], kValue2));
  ASSERT_EQ(kvs_.size(), 2u);
  ASSERT_EQ(OkStatus(), kvs_.Delete(keys[0]));

  // Ensure the key is indeed gone and the size is 1 before continuing.
  ASSERT_EQ(kvs_.Get(keys[0], &val), Status::NotFound());
  ASSERT_EQ(kvs_.size(), 1u);
  ASSERT_EQ(kvs_.total_entries_with_deleted(), 2u);

  // Do aggressive FullMaintenance, which should GC the sector with valid data,
  // resulting in no reclaimable bytes and an erased sector.
  EXPECT_EQ(OkStatus(), kvs_.HeavyMaintenance());
  stats = kvs_.GetStorageStats();
  ASSERT_EQ(kvs_.size(), 1u);

  if (PW_KVS_REMOVE_DELETED_KEYS_IN_HEAVY_MAINTENANCE) {
    ASSERT_EQ(kvs_.total_entries_with_deleted(), 1u);
  } else {  // The deleted entries are only removed if that feature is enabled.
    ASSERT_EQ(kvs_.total_entries_with_deleted(), 2u);
  }

  // Read back the second key to make sure it is still valid.
  ASSERT_EQ(kvs_.Get(keys[1], &val), OkStatus());
  ASSERT_EQ(val, kValue2);
}

TEST(InMemoryKvs, Put_MaxValueSize) {
  // Create and erase the fake flash.
  Flash flash;
  ASSERT_EQ(OkStatus(), flash.partition.Erase());

  // Create and initialize the KVS.
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          default_format);
  PW_TEST_ASSERT_OK(kvs.Init());

  size_t max_key_value_size = kvs.max_key_value_size_bytes();
  EXPECT_EQ(max_key_value_size,
            KeyValueStore::max_key_value_size_bytes(
                flash.partition.sector_size_bytes()));

  size_t max_value_size =
      flash.partition.sector_size_bytes() - sizeof(EntryHeader) - 1;
  EXPECT_EQ(max_key_value_size, (max_value_size + 1));

  // Use the large_test_flash as a big chunk of data for the Put statement.
  ASSERT_GT(sizeof(large_test_flash), max_value_size + 2 * sizeof(EntryHeader));
  auto big_data = as_bytes(span(&large_test_flash, 1));

  EXPECT_EQ(OkStatus(), kvs.Put("K", big_data.subspan(0, max_value_size)));

  // Larger than maximum is rejected.
  EXPECT_EQ(Status::InvalidArgument(),
            kvs.Put("K", big_data.subspan(0, max_value_size + 1)));
  EXPECT_EQ(Status::InvalidArgument(), kvs.Put("K", big_data));
}

TEST(InMemoryKvs, SectorWritableBytesUnderflow) {
  // 1. Create a fake flash with 4096-byte sectors, and 4 sectors.
  FakeFlashMemoryBuffer<4096, 4> test_flash(16);
  FlashPartition test_partition(
      &test_flash, 0, static_cast<uint32_t>(test_flash.sector_count()));

  // 2. Erase the partition.
  PW_TEST_ASSERT_OK(test_partition.Erase());

  // 3. Define format.
  constexpr EntryFormat format{.magic = 0x12345678, .checksum = nullptr};

  // 4. Manually construct and write a spanning entry header in Sector 1.
  // We will place the crafted entry at the beginning of Sector 1 (address
  // 0x1000). Value size is 8000 bytes, so it spans way beyond Sector 1 (size
  // 4096).
  EntryHeader header;
  header.magic = format.magic;
  header.checksum = 0;
  header.alignment_units = 0;   // 16-byte alignment
  header.key_length_bytes = 1;  // key: "k"
  header.value_size_bytes = 8000;
  header.transaction_id = 1;

  // Write the entry header at address 0x1000
  PW_TEST_ASSERT_OK(test_partition.Write(0x1000, as_bytes(span(&header, 1))));
  // Write the key "k" at address 0x1010, padded/aligned to 16 bytes
  char key[16] = {'k'};
  PW_TEST_ASSERT_OK(test_partition.Write(0x1010, as_bytes(span(key))));

  // 5. Initialize KVS with manual error recovery.
  Options options;
  options.recovery = ErrorRecovery::kManual;
  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(
      &test_partition, format, options);
  // Init() should now return Status::DataLoss() because corruption (spanning
  // entry) is detected.
  EXPECT_EQ(kvs.Init(), Status::DataLoss());
  EXPECT_TRUE(kvs.error_detected());

  // 6. Get storage stats and assert that NO underflow occurred.
  // Sector 1 is marked as corrupt and has 0 writable bytes.
  // Sector 0 is empty (4096), but excluded as the first empty sector.
  // Sector 2 and Sector 3 are empty and have 4096 writable bytes each.
  // Total writable bytes = 4096 + 4096 = 8192.
  KeyValueStore::StorageStats stats = kvs.GetStorageStats();
  EXPECT_EQ(stats.writable_bytes, 8192u);
}

TEST(InMemoryKvs, Get_PartialAndOffsetChecksumFailure) {
  Flash flash;
  ASSERT_EQ(OkStatus(), flash.partition.Erase());

  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          default_format);
  ASSERT_EQ(OkStatus(), kvs.Init());

  const char* key = "MyKey";
  const char* value = "CorruptedData12345";  // 18 bytes
  span<const std::byte> value_bytes = as_bytes(span(value, std::strlen(value)));

  ASSERT_EQ(OkStatus(), kvs.Put(key, value_bytes));

  span<std::byte> flash_buf = flash.memory.buffer();
  bool found = false;
  for (size_t i = 0; i <= flash_buf.size() - value_bytes.size(); ++i) {
    if (std::memcmp(&flash_buf[i], value_bytes.data(), value_bytes.size()) ==
        0) {
      flash_buf[i] = std::byte{'X'};
      found = true;
      break;
    }
  }
  ASSERT_TRUE(found);

  // Test a partial read detects data loss
  std::array<char, 5> small_buf = {};
  StatusWithSize result = kvs.Get(key, as_writable_bytes(span(small_buf)));
  EXPECT_EQ(Status::DataLoss(), result.status());
  EXPECT_EQ(0u, result.size());
  for (char c : small_buf) {
    EXPECT_EQ(c, '\0');
  }

  // Test an offset read detects data loss
  std::array<char, 13> offset_buf = {};
  result = kvs.Get(key, as_writable_bytes(span(offset_buf)), 5);
  EXPECT_EQ(Status::DataLoss(), result.status());
  EXPECT_EQ(0u, result.size());
  for (char c : offset_buf) {
    EXPECT_EQ(c, '\0');
  }

  // Test offset read that is partial (ResourceExhausted)
  std::array<char, 5> small_offset_buf = {};
  result = kvs.Get(key, as_writable_bytes(span(small_offset_buf)), 5);
  EXPECT_EQ(Status::DataLoss(), result.status());
  EXPECT_EQ(0u, result.size());
  for (char c : offset_buf) {
    EXPECT_EQ(c, '\0');
  }
}

TEST(InMemoryKvs, Get_PartialAndOffsetReadSuccess) {
  Flash flash;
  ASSERT_EQ(OkStatus(), flash.partition.Erase());

  KeyValueStoreBuffer<kMaxEntries, kMaxUsableSectors> kvs(&flash.partition,
                                                          default_format);
  ASSERT_EQ(OkStatus(), kvs.Init());

  const char* key = "MyKey";
  const char* value = "UncorruptedData123";  // 18 bytes
  span<const std::byte> value_bytes = as_bytes(span(value, std::strlen(value)));

  ASSERT_EQ(OkStatus(), kvs.Put(key, value_bytes));

  // Test partial read (ResourceExhausted)
  std::array<char, 5> small_buf = {};
  StatusWithSize result = kvs.Get(key, as_writable_bytes(span(small_buf)));
  EXPECT_EQ(Status::ResourceExhausted(), result.status());
  EXPECT_EQ(5u, result.size());
  EXPECT_EQ(std::string_view(small_buf.data(), 5), "Uncor");

  // Test offset read that fits entirely in buffer
  std::array<char, 15> offset_buf = {};
  result = kvs.Get(key, as_writable_bytes(span(offset_buf)), 5);
  EXPECT_EQ(OkStatus(), result.status());
  EXPECT_EQ(13u, result.size());
  EXPECT_EQ(std::string_view(offset_buf.data(), 13), "ruptedData123");

  // Test offset read that is partial (ResourceExhausted)
  std::array<char, 5> small_offset_buf = {};
  result = kvs.Get(key, as_writable_bytes(span(small_offset_buf)), 5);
  EXPECT_EQ(Status::ResourceExhausted(), result.status());
  EXPECT_EQ(5u, result.size());
  EXPECT_EQ(std::string_view(small_offset_buf.data(), 5), "rupte");
}

}  // namespace pw::kvs