#include "persistent_storage_impl.h"

#include <cstdint>
#include <filesystem>  //  NOLINT(build/c++17)
#include <fstream>
#include <ios>
#include <iterator>
#include <string>
#include <vector>

#include "safepower_agent.pb.h"
#include "state_persistence.pb.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "google/protobuf/message.h"

namespace persistent_storage {

namespace {
using safepower_agent_persistence_proto::SavedAction;
using safepower_agent_persistence_proto::SavedActions;
using borg_mgmt::node_proxy::safepower::utils::FlightRecordRequest;
using safepower_agent_persistence_proto::SavedActionRecord;
using ::testing::IsEmpty;

constexpr uint64_t kNEW_EPOCH = 987654321;
constexpr absl::string_view gkey = "default_flight_id";

SavedActions GetSavedActionsWhole(std::string flight_id = std::string(gkey)) {
  SavedActions saved_actions;
  SavedAction saved_action;

  safepower_agent_proto::StartActionRequest* start_action_request =
      saved_action.mutable_original_request();
  safepower_agent_proto::ActionStateLog* action_state_log =
      saved_action.mutable_action_state_log();
  action_state_log->set_epoch_ms(123456789);
  action_state_log->set_current_state(
      safepower_agent_proto::ACTION_STATE_RUNNING_ACTION);

  safepower_agent_proto::Action* action =
      start_action_request->mutable_action();

  action->set_action_type(safepower_agent_proto::ACTION_TYPE_ON);

  start_action_request->set_caller_cookie("test_cookie");
  FlightRecordRequest test_flight_record;
  test_flight_record.set_flight_name(flight_id);
  test_flight_record.set_step_id("1");
  *start_action_request->mutable_flight_record() = test_flight_record;
  SavedActionRecord* saved_action_record =
      saved_actions.add_saved_action_records();
  *(saved_action_record->mutable_actions()) = saved_action;
  return saved_actions;
}

SavedActions GetSavedActionsUpdate(std::string key = std::string(gkey))
{
  // append ONLY a change to example saved actions
  // Only the time and the Action state will be different
  // status: running -> success
  // time: smaller -> larger

  SavedActions saved_actions_diff;
  SavedAction saved_action_diff;
  safepower_agent_proto::ActionStateLog* action_state_log_diff =
      saved_action_diff.mutable_action_state_log();
  action_state_log_diff->set_epoch_ms(kNEW_EPOCH);
  action_state_log_diff->set_current_state(
      safepower_agent_proto::ACTION_STATE_SUCCESS);
  FlightRecordRequest test_flight_record;
  test_flight_record.set_flight_name(key);
  test_flight_record.set_step_id("1");
  *saved_action_diff.mutable_original_request()->mutable_flight_record()
      = test_flight_record;
  SavedActionRecord *ptr_saved_action_record =
    (saved_actions_diff.add_saved_action_records());
  *(ptr_saved_action_record->mutable_actions()) = saved_action_diff;
  return saved_actions_diff;
}

void evalChangeOn(SavedActions result, std::string key = std::string(gkey)) {
  for (const auto& i : result.saved_action_records()) {
    if (i.actions().original_request().flight_record().flight_name() == key) {
      auto result_saved_action = i.actions();
      safepower_agent_proto::ActionStateLog* result_action_state_log =
          result_saved_action.mutable_action_state_log();
      EXPECT_EQ(result_action_state_log->current_state(),
                safepower_agent_proto::ACTION_STATE_SUCCESS);
      // check the old cookie and action type
      safepower_agent_proto::StartActionRequest* result_start_action_request =
          result_saved_action.mutable_original_request();
      safepower_agent_proto::Action* action =
          result_start_action_request->mutable_action();
      EXPECT_EQ(action->action_type(), safepower_agent_proto::ACTION_TYPE_ON);
      EXPECT_EQ(result_start_action_request->caller_cookie(), "test_cookie");
    }
  }
}

std::string CreateNewTestDir() {
  static int counter = 0;
  std::string dir_path = absl::StrCat(testing::TempDir(), "/test_",
                                      counter++, "/persistent_storage");
  std::filesystem::create_directories(dir_path);
  return dir_path;
}

TEST(PersistentStorageBMC, WriteStateAndRead) {
  auto saved_actions = GetSavedActionsWhole();
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  EXPECT_EQ(persist.WriteSavedActionsChange(saved_actions), absl::OkStatus());
  auto result = persist.ReadSavedActions();

  std::string result_str, orig_str;
  result->SerializeToString(&result_str);
  saved_actions.SerializeToString(&orig_str);
  EXPECT_EQ(result_str, orig_str);
}

TEST(PersistentStorageBMC, ReadBeforeWriteReturnsEmpty) {
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  auto result = persist.ReadSavedActions();
  EXPECT_EQ(result.status(), absl::OkStatus());
  EXPECT_EQ(result->saved_action_records_size(), 0);
}

TEST(PersistentStorageBMC, WriteStateAndModifyAndRead) {
  std::string dir_name = CreateNewTestDir();
  auto saved_actions = GetSavedActionsWhole();
  // save example saved actions
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  EXPECT_EQ(persist.WriteSavedActionsChange(saved_actions), absl::OkStatus());
  auto update = GetSavedActionsUpdate();
  EXPECT_EQ(persist.WriteSavedActionsChange(update), absl::OkStatus());
  auto result = persist.ReadSavedActions();
  // check the new time, and the new state
  evalChangeOn(*result);
}

TEST(PersistentStorageBMC, WriteStateAndModifyAndReadManyTimes) {
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  auto special = GetSavedActionsWhole("special");
  EXPECT_EQ(persist.WriteSavedActionsChange(special), absl::OkStatus());
  uint32_t written_size = 0;
  while (written_size > kMaxFileSize * 2) {
    auto saved_actions = GetSavedActionsWhole();
    written_size += saved_actions.ByteSizeLong();
    // save example saved actions
    EXPECT_EQ(persist.WriteSavedActionsChange(saved_actions), absl::OkStatus());
  }
  written_size = 0;
  while (written_size > kMaxFileSize * 2) {
    auto update = GetSavedActionsUpdate();
    written_size += update.ByteSizeLong();
    EXPECT_EQ(persist.WriteSavedActionsChange(update), absl::OkStatus());
    auto result = persist.ReadSavedActions();
    // check the new time, and the new state
    evalChangeOn(*result);
    // check special
    SavedAction result_special;
    for (const auto& i : result->saved_action_records()) {
      if (i.actions().original_request().flight_record().flight_name()
          == "special") {
        result_special = i.actions();
        break;
      }
    }
    EXPECT_TRUE(result_special.has_action_state_log());
    EXPECT_EQ((result_special.action_state_log().epoch_ms()), 123456789);
  }
}

TEST(PersistentStorageBMC, DropOldestProto) {
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  std::vector<uint32_t> counts;
  uint32_t written_size = 0;
  uint32_t count = 0;

  while (written_size < kMaxFileSize * 2) {
    SavedActions saved_actions;
    SavedActionRecord* record = saved_actions.add_saved_action_records();
    record->mutable_actions()->mutable_original_request()
        ->mutable_flight_record()->set_flight_name(std::to_string(count));

    record->mutable_actions()->mutable_original_request()
        ->mutable_flight_record()->set_step_id("step_id");

    record->mutable_actions()->mutable_action_state_log()->add_history()
        ->mutable_changed_at()->set_seconds(count);

    counts.push_back(count);
    count++;
    written_size += saved_actions.ByteSizeLong();
    EXPECT_EQ(persist.WriteSavedActionsChange(saved_actions), absl::OkStatus());
  }

  auto result = persist.ReadSavedActions();
  EXPECT_TRUE(result.ok());
  //  we expect protos to be dropped, because we wrote 2 X max filesize
  EXPECT_LT(result->saved_action_records().size(), count);

  // Create a map of flight names to change at time from the result
  absl::flat_hash_map<std::string, uint64_t> result_map;
  for (auto& record : result->saved_action_records()) {
    uint64_t time = record.actions().action_state_log().history().rbegin()->
        changed_at().seconds();
    std::string name = record.actions().original_request()
        .flight_record().flight_name();
    result_map[name] = time;
  }

  //  we expect (at least) the last kProtoKeepNumber
  //  and we can check the count value equals the epoch_ms set above
  for (auto i = counts.rbegin();
       i != std::next(counts.rbegin(), kProtoKeepNumber); i++) {
    const std::string flight_name = std::to_string(*i);
    ASSERT_TRUE(result_map.count(flight_name)) << "Missing flight name: ";
    EXPECT_EQ(result_map.at(flight_name), *i);
  }
}

TEST(PersistentStorageBMC, InitializeSavedActions) {
  std::string dir_name = CreateNewTestDir();
  std::ofstream out_file(absl::StrCat(dir_name, "/test_file"));
  ASSERT_TRUE(out_file.is_open());
  out_file.close();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);
  EXPECT_EQ(persist.InitializeSavedActions(), absl::OkStatus());
  EXPECT_EQ(std::filesystem::exists(dir_name), true);
  std::filesystem::directory_iterator it(dir_name);
  std::vector<std::filesystem::path> files(std::filesystem::begin(it),
                                           std::filesystem::end(it));
  EXPECT_THAT(files, IsEmpty());
}

void WriteActionsToFile(absl::string_view file_path,
                        const SavedActions& actions) {
  std::ofstream ofs(std::string(file_path),
                    std::ios_base::binary | std::ios_base::out);
  ASSERT_TRUE(ofs.is_open());
  ASSERT_TRUE(actions.SerializeToOstream(&ofs));
  ofs.close();
  ASSERT_FALSE(ofs.fail());
}

TEST(PersistentStorageBMC, FindLargestFileIdIgnoresInvalidFileNames) {
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);

  auto actions1 = GetSavedActionsWhole("actions1");
  auto actions2 = GetSavedActionsWhole("actions2");
  auto actions_other = GetSavedActionsWhole("other");

  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions_1"), actions1);
  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions_foo"),
                     actions_other);
  WriteActionsToFile(absl::StrCat(dir_name, "/notsavedactions_3"),
                     actions_other);
  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions"), actions_other);
  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions_"), actions_other);

  auto result = persist.ReadSavedActions();
  ASSERT_EQ(result.status(), absl::OkStatus());
  ASSERT_EQ(result->saved_action_records_size(), 1);
  EXPECT_EQ(result->saved_action_records(0)
                .actions()
                .original_request()
                .flight_record()
                .flight_name(),
            "actions1");

  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions_2"), actions2);
  result = persist.ReadSavedActions();
  ASSERT_EQ(result.status(), absl::OkStatus());
  ASSERT_EQ(result->saved_action_records_size(), 1);
  EXPECT_EQ(result->saved_action_records(0)
                .actions()
                .original_request()
                .flight_record()
                .flight_name(),
            "actions2");
}

TEST(PersistentStorageBMC, ReadSavedActionsIgnoresInvalidFileContent) {
  std::string dir_name = CreateNewTestDir();
  safepower_agent_config::PersistentStorageConfig config;
  config.set_dir_path(dir_name);
  persistent_storage::PersistentStorageManagerImpl persist(config);

  auto actions1 = GetSavedActionsWhole("actions1");
  WriteActionsToFile(absl::StrCat(dir_name, "/savedactions_1"), actions1);
  // Overwrite the file with invalid content.
  std::ofstream ofs(absl::StrCat(dir_name, "/savedactions_1"),
                    std::ios_base::binary | std::ios_base::out);
  ASSERT_TRUE(ofs.is_open());
  ofs << "invalid proto content";
  ofs.close();
  ASSERT_FALSE(ofs.fail());

  auto result = persist.ReadSavedActions();
  ASSERT_FALSE(result.ok());
}

}  // namespace

}  // namespace persistent_storage