#include "persistent_storage_impl.h"

#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <filesystem>  // NOLINT(build/c++17)
#include <iterator>
#include <optional>
#include <ostream>
#include <string>
#include <system_error>  // NOLINT(build/c++11)
#include <utility>
#include <vector>

#include "google/protobuf/timestamp.pb.h"
#include "action_context.h"
#include "flight_record.h"
#include "action.pb.h"
#include "safepower_agent.pb.h"
#include "safepower_agent_config.pb.h"
#include "state_persistence.pb.h"
#include "proto_reader.h"
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "bmc/status_macros.h"
#include "google/protobuf/io/zero_copy_stream_impl.h"

namespace persistent_storage {

using safepower_agent::ActionContext;
using safepower_agent::FlightRecordWrapper;
using safepower_agent_persistence_proto::SavedAction;
using safepower_agent_persistence_proto::SavedActionRecord;
using safepower_agent_persistence_proto::SavedActions;

// TODO: b/493660834 -
// Convert this to use posix file APIs, and not std::filesystem
inline absl::Status WriteSavedActions(const std::string& proto_path,
                                      const SavedActions& in_proto,
                                      bool append = true) {
  int flags = O_WRONLY | O_CREAT;
  if (append) {
    flags |= O_APPEND;
  } else {
    flags |= O_TRUNC;
  }
  int fd = open(proto_path.c_str(), flags, 0644);
  if (fd < 0) {
    int ec = errno;
    LOG(ERROR) << "Failed to open file " << proto_path << " " << strerror(ec);
    return absl::UnavailableError(
        absl::StrCat("Failed to open file: ", proto_path, ":", strerror(ec)));
  }

  google::protobuf::io::FileOutputStream stream(fd);
  stream.SetCloseOnDelete(true);
  if (!in_proto.SerializeToZeroCopyStream(&stream)) {
    int ec = stream.GetErrno();
    LOG(ERROR) << "Failed to serialize the proto file" << proto_path
               << " error:" << strerror(ec);
    return absl::InternalError(
        absl::StrCat("Failed to serialize the proto file: ", proto_path,
                     " error:", strerror(ec)));
  }

  if (!stream.Flush()) {
    int ec = stream.GetErrno();
    LOG(ERROR) << "Failed to flush the proto file" << proto_path
               << " error:" << strerror(ec);
    return absl::InternalError(
        absl::StrCat("Failed to flush the proto file: ", proto_path,
                     " error:", strerror(ec)));
  }

  int ret = fsync(fd);
  if (ret < 0) {
    int ec = errno;
    LOG(ERROR) << "Failed to sync the file " << proto_path << " "
               << strerror(ec);
    return absl::DataLossError(absl::StrCat(
        "Failed to sync the file: ", proto_path, ":", strerror(ec)));
  }
  return absl::OkStatus();
}

PersistentStorageManagerImpl::PersistentStorageManagerImpl(
    const safepower_agent_config::PersistentStorageConfig& config) {
  if (!config.dir_path().empty()) {
    dir_path_ = config.dir_path();
  }
  if (config.max_file_size() > 0) {
    max_file_size_ = config.max_file_size();
  }
  if (config.proto_keep_number() > 0) {
    proto_keep_number_ = config.proto_keep_number();
  }
}

// Writes a delta of actions. If the active file exceeds the size limit,
// it rotates (compresses and writes to the other file).
// TODO:  b/469169011 - Rename these functions
absl::Status PersistentStorageManagerImpl::WriteSavedActionsChange(
    const SavedActions& actions) {
  absl::MutexLock lock(mu_);
  auto info = GetActiveFileInfo();
  if (!info.ok()) {
    return info.status();
  }

  uint64_t write_id = 0;
  int64_t next_sequence = 1;

  if (info->exists) {
    if (info->id > 1) {
      // Migrating from old unbounded files
      write_id = 0;
      next_sequence = 1;

      SavedActions merged_actions = info->proto;
      for (const auto& record : actions.saved_action_records()) {
        *merged_actions.add_saved_action_records() = record;
      }
      if (!actions.boot_id().empty()) {
        merged_actions.set_boot_id(actions.boot_id());
      }
      if (actions.boot_counter() != 0) {
        merged_actions.set_boot_counter(actions.boot_counter());
      }

      auto compressed = CompressLogs(merged_actions, /*prune=*/true);
      if (!compressed.ok()) {
        return compressed.status();
      }
      compressed->set_sequence_number(next_sequence);

      RETURN_IF_ERROR(WriteSavedActions(MakeFilePath(dir_path_, write_id),
                                        *compressed, /*append=*/false));
      RETURN_IF_ERROR(DeleteOldFiles());
      return absl::OkStatus();
    } else {
      write_id = info->id;
      next_sequence = info->sequence_number;
    }
  } else {
    write_id = 0;
    next_sequence = 1;
  }

  std::string file_path = MakeFilePath(dir_path_, write_id);
  SavedActions actions_to_write = actions;
  actions_to_write.set_sequence_number(next_sequence);

  RETURN_IF_ERROR(
      WriteSavedActions(file_path, actions_to_write, /*append=*/true));

  // if the file is too large, merge the files
  std::error_code ec;
  std::uintmax_t file_size = std::filesystem::file_size(file_path, ec);
  if (ec) {
    return absl::UnavailableError(
        absl::StrCat("Failed to get file size", ec.message()));
  }

  if (file_size > max_file_size_) {
    LOG(INFO) << "file size too large merging" << std::endl;
    auto full_actions = ReadProtoForId(write_id);
    if (!full_actions.ok()) {
      return full_actions.status();
    }

    uint64_t next_write_id = 1 - write_id;
    uint64_t next_seq = next_sequence + 1;

    auto compressed = CompressLogs(*full_actions, /*prune=*/true);
    if (!compressed.ok()) {
      return compressed.status();
    }
    compressed->set_sequence_number(next_seq);

    RETURN_IF_ERROR(WriteSavedActions(MakeFilePath(dir_path_, next_write_id),
                                      *compressed, /*append=*/false));
  }
  return absl::OkStatus();
}

absl::Status PersistentStorageManagerImpl::ConvertAppendToMap(
    absl::flat_hash_map<FlightRecordWrapper, SavedAction>& map,
    const SavedActions& saved_actions) {
  for (const auto& i : saved_actions.saved_action_records()) {
    FlightRecordWrapper flight_record(
        i.actions().original_request().flight_record());
    auto [it, inserted] = map.try_emplace(flight_record, i.actions());
    if (!inserted) {
      it->second.MergeFrom(i.actions());
    }
  }
  return absl::OkStatus();
}

absl::StatusOr<SavedActions> PersistentStorageManagerImpl::CompressLogs(
    const SavedActions& saved_actions, bool prune) {
  // convert to map, and back to append log
  absl::flat_hash_map<FlightRecordWrapper, SavedAction> new_saved_actions_map;
  absl::Status status =
      this->ConvertAppendToMap(new_saved_actions_map, saved_actions);
  if (!status.ok()) {
    return status;
  }
  if (prune) {
    absl::Status keep_status = keepMostRecentProtos(new_saved_actions_map);
    if (!keep_status.ok()) {
      return keep_status;
    }
  }
  SavedActions new_log;
  for (auto& i : new_saved_actions_map) {
    SavedActionRecord* ptr_record = new_log.add_saved_action_records();
    *(ptr_record->mutable_actions()) = i.second;
  }
  if (!saved_actions.boot_id().empty()) {
    new_log.set_boot_id(saved_actions.boot_id());
  }
  if (saved_actions.boot_counter() != 0) {
    new_log.set_boot_counter(saved_actions.boot_counter());
  }
  new_log.set_sequence_number(saved_actions.sequence_number());
  return new_log;
}

absl::StatusOr<SavedActions> PersistentStorageManagerImpl::ReadSavedActions() {
  absl::MutexLock lock(mu_);
  auto info = GetActiveFileInfo();
  if (!info.ok()) {
    return info.status();
  }
  if (!info->exists) {
    LOG(INFO) << "No saved actions files found in " << dir_path_;
    return SavedActions();
  }
  auto compressed = CompressLogs(info->proto);
  if (!compressed.ok()) {
    LOG(ERROR) << "Unable to compress SavedActions: " << compressed.status();
    return compressed.status();
  }
  return *compressed;
}

namespace {
google::protobuf::Timestamp GetChangeTime(const SavedAction& action,
                                          absl::Time now) {
  if (!action.has_action_state_log() ||
      (action.action_state_log().history().empty())) {
    // if there is no history, return the time of the last 24 hours
    absl::Time one_day_ago = now - absl::Hours(24);
    google::protobuf::Timestamp timestamp;
    timestamp.set_seconds(absl::ToUnixSeconds(one_day_ago));
    return timestamp;
  }
  return action.action_state_log().history().rbegin()->changed_at();
}

bool AtEndState(const SavedAction& action) {
  if ((!action.has_action_state_log()) ||
      !action.action_state_log().has_current_state()) {
    return false;  // if there is no state
    // we should depend on the timing logic to clear it
  }
  const safepower_agent_proto::ActionState& state =
      action.action_state_log().current_state();
  return ActionContext::IsFinalState(state);
}
}  // namespace

// when the persistence proto is cleaned up, we need to remove the protos that
// we need to keep the most recent N protos
absl::Status PersistentStorageManagerImpl::keepMostRecentProtos(
    absl::flat_hash_map<FlightRecordWrapper, SavedAction>& map) {
  if (map.size() < proto_keep_number_) {
    LOG(WARNING) << " Assumptions about the size of the"
                 << " saved protos are wrong" << std::endl
                 << "The number of saved protos is " << map.size()
                 << "less than the number we want to keep "
                 << proto_keep_number_;
    return absl::OkStatus();
  }

  // sort the protos by time
  std::vector<SavedActionRecord> records;
  for (const auto& i : map) {
    SavedActionRecord record;
    *(record.mutable_actions()) = i.second;
    records.push_back(record);
  }

  absl::Time now = absl::Now();
  // Sort the records so that all records in a final state come before those
  // that do not. Within each grouping, sort from oldest to youngest.
  auto compare_records = [now](const SavedActionRecord& a,
                               const SavedActionRecord& b) {
    bool final_a = AtEndState(a.actions());
    bool final_b = AtEndState(b.actions());
    if (final_a != final_b) {
      // If a is in a final state, it is "<" b, which is not.
      // If a is _not_ in a final state, b is, and a ">" b.
      return final_a;
    }
    // In ascending time, so oldest to youngest
    auto time_a = GetChangeTime(a.actions(), now);
    auto time_b = GetChangeTime(b.actions(), now);
    return (time_a.seconds() < time_b.seconds()) ||
           (time_a.seconds() == time_b.seconds() &&
            time_a.nanos() < time_b.nanos());
  };
  absl::c_stable_sort(records, compare_records);

  // erase the oldest protos, keeping proto_keep_number_ protos
  LOG(INFO) << "cleaning up the saved protos";
  LOG(INFO) << "current:" << records.size();
  LOG(INFO) << "keeping:" << proto_keep_number_;
  const size_t num_to_delete = records.size() - proto_keep_number_;
  LOG(INFO) << "deleting:" << num_to_delete;
  if (proto_keep_number_ > records.size()) {
    LOG(INFO) << "keeping all protos";
    return absl::OkStatus();
  }

  auto end_delete = std::next(records.begin(), num_to_delete);
  for (auto i = records.begin(); i != end_delete; i++) {
    FlightRecordWrapper action_to_delete(
        i->actions().original_request().flight_record());
    // In ideal circumstances, only actions that are completed are deleted.
    if (AtEndState(i->actions())) {
      LOG(INFO) << "deleting end-state action record:" << action_to_delete;
    } else {
      LOG(ERROR) << "deleting an unfinished action record:" << action_to_delete;
    }
    map.erase(action_to_delete);
  }
  return absl::OkStatus();
}

// find the largest file id, and return it

// the file id is a number, and the file name is in the format of
// "savedactions_1"
// "savedactions_2"
// where the largest id is the most recent file
std::optional<uint64_t> PersistentStorageManagerImpl::findLargestFileId(
    absl::Span<const std::string> files) {
  std::optional<uint64_t> largest_id;
  for (const auto& entry : files) {
    std::vector<absl::string_view> file_name_parts = absl::StrSplit(entry, '_');
    if (file_name_parts.size() != 2) {
      LOG(WARNING) << "unable to parse the file name in the directory " << entry
                   << " size:" << file_name_parts.size();
      continue;
    }
    uint64_t file_id;
    if (!absl::SimpleAtoi(file_name_parts[1], &file_id)) {
      LOG(WARNING) << "unable to parse the file ID as unint64_t" << entry;
      continue;
    }
    if (file_name_parts[0] != kfileNamePrefix) {
      LOG(WARNING) << "unknown file name prefix" << entry;
      continue;
    }
    if (!largest_id.has_value() || file_id > *largest_id) {
      largest_id = file_id;
    }
  }
  return largest_id;
}

absl::StatusOr<std::vector<std::string> >
PersistentStorageManagerImpl::listFiles(absl::string_view file_path) {
  // check if the file path exists and is a directory
  std::error_code ec_is_directory;
  bool is_directory = std::filesystem::is_directory(file_path, ec_is_directory);
  if (ec_is_directory) {
    const std::string error_message =
        absl::StrCat("Failed to check if file is a directory ", file_path,
                     " error:", ec_is_directory.message());
    LOG(ERROR) << error_message;
    return absl::InternalError(error_message);
  }
  if (!is_directory) {
    const std::string error_message =
        absl::StrCat("File path is not a directory:", file_path,
                     "is_directory:  ", is_directory);
    LOG(ERROR) << error_message;
    return absl::InternalError(error_message);
  }
  // list the files in the directory
  std::vector<std::string> files;
  std::error_code ec;
  std::filesystem::directory_iterator it(file_path, ec), end;
  if (ec) {
    const std::string error_message =
        absl::StrCat("Failed to create directory iterator ", file_path,
                     " error:", ec.message());
    LOG(ERROR) << error_message;
    return absl::InternalError(error_message);
  }
  while (it != end) {
    files.push_back(it->path().filename().string());
    std::error_code ec_increment;
    it.increment(ec_increment);
    if (ec_increment) {
      const std::string error_message =
          absl::StrCat("Failed to iterate the directory ", file_path,
                       " error:", ec_increment.message());
      LOG(ERROR) << error_message;
      return absl::InternalError(error_message);
    }
  }
  return files;
}

absl::StatusOr<PersistentStorageManagerImpl::ActiveFileInfo>
PersistentStorageManagerImpl::GetFileInfo(uint64_t id) {
  ActiveFileInfo info;
  info.id = id;

  std::error_code ec;
  info.exists = std::filesystem::exists(MakeFilePath(dir_path_, id), ec);
  if (ec && ec != std::errc::no_such_file_or_directory) {
    return absl::InternalError(absl::StrCat("Failed to check if file ", id,
                                            " exists: ", ec.message()));
  }

  if (!info.exists) {
    return info;
  }

  auto proto = ReadProtoForId(id);
  if (!proto.ok()) {
    return absl::DataLossError(absl::StrCat(
        "Savedactions file ", id, " is corrupt: ", proto.status().message()));
  }

  info.sequence_number = proto->sequence_number();
  info.proto = *proto;
  return info;
}

absl::StatusOr<PersistentStorageManagerImpl::ActiveFileInfo>
PersistentStorageManagerImpl::GetActiveFileInfo() {
  auto res_0 = GetFileInfo(0);
  if (!res_0.status().ok() && !absl::IsDataLoss(res_0.status())) {
    return res_0.status();
  }
  auto res_1 = GetFileInfo(1);
  if (!res_1.status().ok() && !absl::IsDataLoss(res_1.status())) {
    return res_1.status();
  }

  bool corrupt_0 = absl::IsDataLoss(res_0.status());
  bool corrupt_1 = absl::IsDataLoss(res_1.status());

  if (corrupt_0 && corrupt_1) {
    return absl::DataLossError(absl::StrCat(
        "Both savedactions files are corrupt: ", res_0.status().message(), "; ",
        res_1.status().message()));
  }

  bool ok_0 = res_0.ok() && res_0->exists;
  bool ok_1 = res_1.ok() && res_1->exists;

  if (ok_0 && ok_1) {
    // If both files have the same sequence number (e.g., they are both
    // unpopulated old files from before the transition), we favor file 1
    // because it would have been the newer file in the old scheme.
    if (res_0->sequence_number > res_1->sequence_number) {
      return *res_0;
    } else {
      return *res_1;
    }
  }

  if (ok_0) {
    return *res_0;
  }

  if (ok_1) {
    return *res_1;
  }

  if (corrupt_0) {
    return res_0.status();
  }
  if (corrupt_1) {
    return res_1.status();
  }

  DCHECK(res_0.ok() && res_1.ok());
  DCHECK(!res_0->exists && !res_1->exists);

  // Migration path: If neither savedactions_0 nor savedactions_1 exists, but
  // older files (with IDs > 1) exist, we find the largest one and treat it as
  // the base state. This allows seamless migration to the new ping-pong storage
  // scheme.
  auto files = listFiles(dir_path_);
  if (!files.ok()) {
    return files.status();
  }
  std::optional<uint64_t> old_id = findLargestFileId(*files);
  if (old_id.has_value()) {
    auto res_old = GetFileInfo(old_id.value());
    if (res_old.ok() && res_old->exists) {
      return *res_old;
    }
    if (old_id.value() > 0) {
      auto res_fallback = GetFileInfo(old_id.value() - 1);
      if (res_fallback.ok() && res_fallback->exists) {
        return *res_fallback;
      }
    }
    return absl::DataLossError("Old files exist but are corrupt");
  }

  ActiveFileInfo info;
  info.exists = false;
  return info;
}

absl::StatusOr<SavedActions> PersistentStorageManagerImpl::ReadProtoForId(
    uint64_t file_id) {
  std::string file_path = MakeFilePath(dir_path_, file_id);
  return proto_reader::ReadProto<SavedActions>(file_path);
}

absl::Status PersistentStorageManagerImpl::DeleteOldFiles() {
  auto files = listFiles(dir_path_);
  if (!files.ok()) {
    return files.status();
  }
  for (const auto& entry : *files) {
    if (entry == MakeFileName(0) || entry == MakeFileName(1)) {
      continue;
    }
    const std::string del_path = absl::StrCat(dir_path_, "/", entry);
    std::error_code ec;
    bool removed = std::filesystem::remove(del_path, ec);
    if (!removed || ec) {
      LOG(WARNING) << "Failed to delete old file " << del_path
                   << " remove result:" << removed << " error:" << ec.message();
    }
  }
  return absl::OkStatus();
}

absl::Status PersistentStorageManagerImpl::InitializeSavedActions() {
  absl::MutexLock lock(mu_);
  std::error_code ec;
  if (!std::filesystem::is_directory(dir_path_, ec)) {
    return absl::FailedPreconditionError(
        absl::StrCat("Directory does not exist: ", dir_path_));
  }
  if (ec) {
    const std::string error_message =
        absl::StrCat("Failed to check if directory exists ", dir_path_,
                     " error:", ec.message());
    LOG(ERROR) << error_message;
    return absl::InternalError(error_message);
  }
  std::filesystem::directory_iterator it(dir_path_, ec), end;
  if (ec) {
    const std::string error_message =
        absl::StrCat("Failed to create directory iterator ", dir_path_,
                     " error:", ec.message());
    LOG(ERROR) << error_message;
    return absl::InternalError(error_message);
  }
  while (it != end) {
    const std::filesystem::path& file = it->path();
    std::error_code remove_ec;
    bool removed = std::filesystem::remove(file, remove_ec);
    if (!removed || remove_ec) {
      const std::string error_message =
          absl::StrCat("Failed to delete the file ", file.string(),
                       ": error:", remove_ec.message());
      LOG(ERROR) << error_message;
      return absl::UnavailableError(error_message);
    }
    LOG(WARNING) << "Deleted " << file.string();
    std::error_code ec_increment;
    it.increment(ec_increment);
    if (ec_increment) {
      LOG(ERROR) << "Failed to iterate the directory " << dir_path_
                 << " error:" << ec_increment.message();
      return absl::InternalError(absl::StrCat(
          "Failed to iterate the directory, file_path: ", dir_path_,
          " error:", ec_increment.message()));
    }
  }
  return absl::OkStatus();
}

std::string PersistentStorageManagerImpl::MakeFileName(const uint64_t file_id) {
  return absl::StrCat(kfileNamePrefix, "_", std::to_string(file_id));
}

std::string PersistentStorageManagerImpl::MakeFilePath(
    absl::string_view dir_path, const uint64_t file_id) {
  return absl::StrCat(dir_path, "/", MakeFileName(file_id));
}
}  // namespace persistent_storage