#include "action_context.h"

#include <algorithm>
#include <array>
#include <functional>
#include <iostream>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "google/protobuf/timestamp.pb.h"
#include "action_validation.h"
#include "condition.h"
#include "convert_proto.h"
#include "daemon_context.h"
#include "flight_record.h"
#include "action.pb.h"
#include "safepower_agent.pb.h"
#include "state_persistence.pb.h"
#include "state_updater.h"
#include "utils.pb.h"
#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/functional/bind_front.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "bmc/status_macros.h"
#include "google/protobuf/repeated_ptr_field.h"

namespace safepower_agent {

using borg_mgmt::node_proxy::safepower::utils::FlightRecordRequest;
using ::safepower_agent_persistence_proto::SavedActions;
using ::safepower_agent_proto::ActionState;
using ::safepower_agent_proto::ActionStateChange;
using ::safepower_agent_proto::ActionStateLog;

static constexpr absl::Duration kMinRetryInterval = absl::Milliseconds(100);

ActionStateLog ActionContext::NewInitialState() {
  ActionStateLog initial_state;
  initial_state.set_epoch_ms(DaemonContext::Get().epoch_ms());
  initial_state.set_current_state(safepower_agent_proto::ACTION_STATE_INIT);
  SetTimestampToNow(*initial_state.add_history()->mutable_changed_at());
  return initial_state;
}

static absl::Time GetStartTime(
    const safepower_agent_proto::ActionStateLog& state) {
  if (state.history_size() == 0) [[unlikely]] {
    LOG(DFATAL) << "No start time in initial state";
    return DaemonContext::Get().now();
  }
  return ConvertTime(state.history(0).changed_at());
}

ActionContext::ActionContext(CreationToken /*token*/,
                             ActionContextManager& manager,
                             safepower_agent_proto::StartActionRequest request,
                             Action action_impl, ActionStateLog initial_state)

    : manager_(manager),
      request_(std::move(request)),
      flight_record_(request_.flight_record()),
      action_impl_(std::move(action_impl)),
      precondition_(
          request_.has_precondition()
              ? std::make_optional<Condition>(request_.precondition(),
                                              GetStartTime(initial_state))
              : std::nullopt),
      validation_(request_.has_validation()
                      ? std::make_optional<Condition>(
                            request_.validation(), GetStartTime(initial_state))
                      : std::nullopt),
      action_state_updater_(
          std::make_shared<StateUpdater<ActionStateLog>>(initial_state)) {}

absl::Status ActionContext::StartCheckingCondition(
    Condition& condition, absl::string_view condition_id,
    absl::AnyInvocable<void(absl::Status, Condition::MatchList)> callback) {
  absl::Status status = condition.WaitForMatch(
      absl::StrCat(execution_task_name(), ".", condition_id),
      manager_.global_system_state_updater(), std::move(callback));
  if (!status.ok()) {
    LOG(ERROR) << "Failed to wait for condition for " << flight_record() << ": "
               << status;
    ActionStateChange change;
    SetStatus(*change.mutable_status(), status);
    SetState(ActionState::ACTION_STATE_ERROR, std::move(change));
    Finish(std::move(action_impl_));
    return status;
  }
  return absl::OkStatus();
}

absl::Status ActionContext::Activate() {
  absl::MutexLock lock(mutex_);
  const safepower_agent_proto::ActionState current_state =
      action_state_updater_->WithState(
          [](const ActionStateLog& state) { return state.current_state(); });
  LOG(INFO) << "Activate " << flight_record() << " in state "
            << safepower_agent_proto::ActionState_Name(current_state);
  if (!IsFinalState(current_state)) {
    keep_alive_listener_ = std::make_unique<KeepAliveListener>();
    keep_alive_listener_->Listen(manager_.system_state_updater());
  }
  switch (current_state) {
    case safepower_agent_proto::ACTION_STATE_UNDEFINED:
      return absl::InvalidArgumentError(
          absl::StrCat("Action state is unknown:",
                       safepower_agent_proto::ActionState_Name(current_state)));
    case safepower_agent_proto::ACTION_STATE_INIT: {
      RETURN_IF_ERROR(EnterStateInit());
      break;
    }
    case safepower_agent_proto::ACTION_STATE_CHECKING_PRECONDITION:
      if (precondition_.has_value()) {
        return EnterStateCheckingPrecondition();
      }

      LOG(WARNING)
          << flight_record()
          << " has no precondition but is in checking precondition state";
      SetState(ActionState::ACTION_STATE_RUNNING_ACTION);

      break;
    case safepower_agent_proto::ACTION_STATE_RUNNING_ACTION:
      LOG(INFO) << flight_record()
                << " resumed in running action state. Assuming the action was "
                   "already run successfully.";
      if (validation_.has_value()) {
        SetState(ActionState::ACTION_STATE_VALIDATING_FINAL_STATE);
      } else {
        SetState(ActionState::ACTION_STATE_SUCCESS);
      }
      break;
    case safepower_agent_proto::ACTION_STATE_VALIDATING_FINAL_STATE: {
      if (validation_.has_value()) {
        return EnterStateValidatingFinalState();
      }

      ActionStateChange change;
      SetStatus(*change.mutable_status(),
                absl::InvalidArgumentError(
                    "No validation specified in validating state"));
      LOG(WARNING) << flight_record()
                   << " has no validation but is in validating state";
      SetState(ActionState::ACTION_STATE_ERROR, std::move(change));
      break;
    }
    case safepower_agent_proto::ACTION_STATE_ABORTED:
    case safepower_agent_proto::ACTION_STATE_ERROR:
    case safepower_agent_proto::ACTION_STATE_SUCCESS:
      break;
    default: {
      LOG(DFATAL) << "Action state is unhandled: "
                  << safepower_agent_proto::ActionState_Name(current_state);
      return absl::UnimplementedError(
          absl::StrCat("Action state is unhandled:",
                       safepower_agent_proto::ActionState_Name(current_state)));
    }
  }

  return absl::OkStatus();
}

ActionContext::~ActionContext() {
  // Always try to cancel. This will fail if the action is already finished or
  // has not started yet, which is fine.
  const std::array<std::string, 4> tasks = {
      init_task_name(), execution_task_name(), response_task_name(),
      final_state_task_name()};
  for (absl::string_view task_name : tasks) {
    absl::Status cancel_status =
        DaemonContext::Get().scheduler().CancelCall(task_name);
    if (!cancel_status.ok() && !absl::IsNotFound(cancel_status)) {
      LOG(DFATAL) << "Failed to cancel action task " << task_name << " for "
                  << flight_record() << ": " << cancel_status;
    }
  }
  if (!mutex_.try_lock()) {
    LOG(DFATAL) << "Action " << flight_record() << " is still running";
    return;
  }
  if (keep_alive_listener_ != nullptr) {
    keep_alive_listener_->StopListening();
    keep_alive_listener_.reset();
  }
  Action action_impl = std::move(action_impl_);
  mutex_.unlock();
  Finish(std::move(action_impl));
}

void ActionContext::DumpDebugInfo(std::ostream& os) const {
  absl::MutexLock lock(mutex_);
  os << "  ActionContext for Flight Record: " << flight_record() << std::endl;
  os << "    Request: " << request_.DebugString() << std::endl;
  if (precondition_.has_value()) {
    os << "    Precondition:" << std::endl;
    precondition_->DumpDebugInfo(os);
  } else {
    os << "    Precondition: None" << std::endl;
  }
  if (validation_.has_value()) {
    os << "    Validation:" << std::endl;
    validation_->DumpDebugInfo(os);
  } else {
    os << "    Validation: None" << std::endl;
  }
  os << "    Current State: " << action_state_updater_->DebugString()
     << std::endl;
}

void ActionContext::Finish(Action action_impl) {
  if (action_impl != nullptr) {
    manager_.FinishAction(request_.action(), std::move(action_impl));
  }
}

static void PersistState(
    const ::borg_mgmt::node_proxy::safepower::utils::FlightRecordRequest&
        flight_record,
    safepower_agent_proto::ActionStateLog state_change) {
  SavedActions saved_actions;
  auto ptr_saved_action_record = saved_actions.add_saved_action_records();
  *(ptr_saved_action_record->mutable_actions()
        ->mutable_original_request()
        ->mutable_flight_record()) = flight_record;
  auto ptr_saved_action = ptr_saved_action_record->mutable_actions();
  *(ptr_saved_action->mutable_action_state_log()) = std::move(state_change);

  absl::Status write_status =
      DaemonContext::Get().persistent_storage_manager().WriteSavedActionsChange(
          saved_actions);
  if (!write_status.ok()) {
    LOG(ERROR) << "Failed to persist action state: " << write_status;
  }
}

static void AssignConditions(ActionStateChange& change,
                             const Condition::MatchList& matches) {
  for (const safepower_agent_proto::Condition*  condition :
       matches) {
    *change.add_matching_conditions() = *condition;
  }
}

void ActionContext::SetState(safepower_agent_proto::ActionState new_state,
                             ActionStateChange change_info) {
  using safepower_agent_proto::ActionStateChange;
  if (manager_.read_only_mode()) {
    LOG(DFATAL) << "SetState called in read-only mode: "
                << safepower_agent_proto::ActionState_Name(new_state) << " "
                << flight_record();
    return;
  }
  LOG(INFO) << "SetState " << flight_record() << " "
            << safepower_agent_proto::ActionState_Name(new_state) << ": "
            << change_info;

  bool is_final = false;
  absl::Status status;
  switch (new_state) {
    default:
      LOG(DFATAL) << "Unexpected state: "
                  << safepower_agent_proto::ActionState_Name(new_state);
      break;
    case safepower_agent_proto::ACTION_STATE_CHECKING_PRECONDITION:
      status = EnterStateCheckingPrecondition();
      break;
    case safepower_agent_proto::ACTION_STATE_RUNNING_ACTION:
      status = EnterStateRunningAction();
      break;
    case safepower_agent_proto::ACTION_STATE_VALIDATING_FINAL_STATE:
      status = EnterStateValidatingFinalState();
      break;
    case safepower_agent_proto::ACTION_STATE_SUCCESS:
    case safepower_agent_proto::ACTION_STATE_ERROR:
    case safepower_agent_proto::ACTION_STATE_ABORTED:
      is_final = true;
      break;
  }
  if (!status.ok()) {
    LOG(ERROR) << "Failed to enter state "
               << safepower_agent_proto::ActionState_Name(new_state) << ": "
               << status;
    ActionStateChange error_change;
    SetStatus(*error_change.mutable_status(), status);
    const safepower_agent_proto::ActionState current_state =
        action_state_updater_->WithState(
            [](const ActionStateLog& state) { return state.current_state(); });
    if (current_state == ActionState::ACTION_STATE_ERROR &&
        new_state == ActionState::ACTION_STATE_ERROR) {
      LOG(DFATAL) << "State change failed too many times";
      return;
    }
    SetState(ActionState::ACTION_STATE_ERROR, std::move(error_change));
    return;
  }

  const safepower_agent_proto::ActionState previous_state =
      action_state_updater_->WithState(
          [](const ActionStateLog& state) { return state.current_state(); });
  ActionStateLog state_change;
  state_change.set_epoch_ms(DaemonContext::Get().epoch_ms());
  state_change.set_current_state(new_state);
  ActionStateChange* history_item = state_change.add_history();
  *history_item = std::move(change_info);
  SetTimestampToNow(*history_item->mutable_changed_at());
  history_item->set_previous_state(previous_state);

  if (is_final) {
    if (keep_alive_listener_ != nullptr) {
      keep_alive_listener_->StopListening();
      keep_alive_listener_.reset();
    }
    absl::Status delay_status = DaemonContext::Get().scheduler().DelayCall(
        [this, action_impl = std::exchange(action_impl_, nullptr),
         state_change = std::move(state_change), is_final]() mutable {
          PersistState(flight_record().request(), state_change);
          Finish(std::move(action_impl));
          action_state_updater_->UpdateState(state_change, is_final);
        },
        absl::ZeroDuration(), final_state_task_name());
    if (!delay_status.ok()) {
      LOG(ERROR) << "Failed to schedule final state update: " << delay_status;
    }
  } else {
    PersistState(flight_record().request(), state_change);
    action_state_updater_->UpdateState(state_change, is_final);
  }
}

absl::Status ActionContext::EnterStateInit() {
  SavedActions saved_actions;
  safepower_agent_persistence_proto::SavedActionRecord*
      ptr_saved_action_record = saved_actions.add_saved_action_records();
  safepower_agent_persistence_proto::SavedAction* ptr_saved_action =
      ptr_saved_action_record->mutable_actions();
  *(ptr_saved_action->mutable_action_state_log()) =
      action_state_updater_->state();
  *(ptr_saved_action->mutable_original_request()) = request_;
  RETURN_IF_ERROR(
      DaemonContext::Get().persistent_storage_manager().WriteSavedActionsChange(
          saved_actions));
  return DaemonContext::Get().scheduler().DelayCall(
      std::bind(&ActionContext::NextStateInit, this), absl::ZeroDuration(),
      init_task_name());
}

void ActionContext::NextStateInit() {
  absl::MutexLock lock(mutex_);
  if (precondition_.has_value()) {
    SetState(ActionState::ACTION_STATE_CHECKING_PRECONDITION);
  } else {
    SetState(ActionState::ACTION_STATE_RUNNING_ACTION);
  }
}

absl::Status ActionContext::EnterStateCheckingPrecondition() {
  return StartCheckingCondition(
      *precondition_, "precondition",
      absl::bind_front(&ActionContext::NextStatePreconditionMatched, this));
}

void ActionContext::NextStatePreconditionMatched(absl::Status status,
                                                 Condition::MatchList matches) {
  absl::MutexLock lock(mutex_);
  ActionStateChange change;
  AssignConditions(change, matches);
  if (!status.ok()) {
    LOG(ERROR) << "Precondition failed: " << status;
    SetStatus(*change.mutable_status(), status);
    if (absl::IsAborted(status)) {
      SetState(safepower_agent_proto::ACTION_STATE_ABORTED, std::move(change));
    } else {
      SetState(safepower_agent_proto::ACTION_STATE_ERROR, std::move(change));
    }
    return;
  }
  SetState(ActionState::ACTION_STATE_RUNNING_ACTION, std::move(change));
}

std::optional<absl::Status> ActionContext::CheckValidation(
    const safepower_agent_proto::SystemState& state) {
  auto [match_status, match_list] = validation_->CheckState(state);
  if ((!match_status.ok() && !absl::IsUnavailable(match_status)) ||
      !match_list.empty()) {
    return DaemonContext::Get().scheduler().DelayCall(
        std::bind(&ActionContext::NextStateValidationCompleted, this,
                  match_status, std::move(match_list)),
        absl::ZeroDuration(), execution_task_name());
  }
  return std::nullopt;
}

absl::Status ActionContext::EnterStateRunningAction() {
  return DaemonContext::Get().scheduler().DelayCall(
      std::bind(&ActionContext::RunAction, this), absl::ZeroDuration(),
      execution_task_name());
}

void ActionContext::RunAction() {
  absl::MutexLock lock(mutex_);
  action_impl_(request_.action(), [this](absl::Status status) {
    // NextStateActionRan must be called asynchronously to avoid deadlocks. It
    // needs to be run as a different task name because the task running
    // RunAction may not have completed when the action callback is called.
    absl::Status delay_status = DaemonContext::Get().scheduler().DelayCall(
        std::bind(&ActionContext::NextStateActionRan, this, status),
        absl::ZeroDuration(), response_task_name());
    if (!delay_status.ok()) {
      LOG(DFATAL) << "Failed to delay next state: " << delay_status;
      NextStateActionRan(status);
    }
  });
}

void ActionContext::NextStateActionRan(absl::Status status) {
  absl::MutexLock lock(mutex_);
  if (!status.ok()) {
    LOG(ERROR) << "Action failed: " << status;
    ActionStateChange change;
    SetStatus(*change.mutable_status(), status);
    SetState(ActionState::ACTION_STATE_ERROR, std::move(change));
    return;
  }
  if (validation_.has_value()) {
    SetState(ActionState::ACTION_STATE_VALIDATING_FINAL_STATE);
  } else {
    SetState(ActionState::ACTION_STATE_SUCCESS);
  }
}

absl::Status ActionContext::EnterStateValidatingFinalState() {
  return StartCheckingCondition(
      *validation_, "validation",
      absl::bind_front(&ActionContext::NextStateValidationCompleted, this));
}

void ActionContext::NextStateValidationCompleted(absl::Status status,
                                                 Condition::MatchList matches) {
  absl::MutexLock lock(mutex_);
  ActionStateChange change;
  AssignConditions(change, matches);
  if (!status.ok()) {
    LOG(ERROR) << "Validation failed: " << status;
    SetStatus(*change.mutable_status(), status);
    if (absl::IsAborted(status)) {
      SetState(safepower_agent_proto::ACTION_STATE_ABORTED, std::move(change));
    } else {
      SetState(ActionState::ACTION_STATE_ERROR, std::move(change));
    }
    return;
  }
  SetState(ActionState::ACTION_STATE_SUCCESS, std::move(change));
}

absl::Status ActionContextManager::RegisterAction(
    const safepower_agent_proto::Action& action,
    ActionContext::Action action_impl) {
  absl::MutexLock lock(actions_mutex_);
  auto [_, inserted] = actions_.try_emplace(action, std::move(action_impl));
  if (!inserted) {
    return absl::AlreadyExistsError(
        absl::StrCat("Action already registered: ", action));
  }
  return absl::OkStatus();
}

absl::StatusOr< std::unique_ptr<ActionContext>>
ActionContextManager::ReinitActionContext(
    safepower_agent_persistence_proto::SavedAction saved_action) {
  Action action_impl = nullptr;
  if (!ActionContext::IsFinalState(
          saved_action.action_state_log().current_state())) {
    ASSIGN_OR_RETURN(action_impl,
                     ReserveAction(saved_action.original_request().action()));
  }

  auto action_context = std::make_unique<ActionContext>(
      ActionContext::CreationToken(), *this,
      std::move(*saved_action.mutable_original_request()),
      std::move(action_impl),
      std::move(*saved_action.mutable_action_state_log()));
  return action_context;
}

absl::Status ActionContextManager::ReloadAndInsertAction(
    safepower_agent_persistence_proto::SavedActionRecord saved_action_record,
    std::unique_ptr<ActionContext>& context_to_cleanup) {
  absl::Status validation_result = absl::OkStatus();
  if (!ActionContext::IsFinalState(
          saved_action_record.actions().action_state_log().current_state())) {
    validation_result = ValidateRequest(
        saved_action_record.actions().original_request(),
        GetStartTime(saved_action_record.actions().action_state_log()));
  }

  FlightRecordWrapper flight_record(
      saved_action_record.actions().original_request().flight_record());

  if (absl::IsUnavailable(validation_result)) {
    reloading_actions_.insert(flight_record);
    // In case validation failed because there was insufficient data to make a
    // decision, retry the validation later.
    LOG(ERROR) << "Failed to reload action; will retry: " << validation_result;
    auto& daemon_context = DaemonContext::Get();
    auto retry_interval =
        std::max(absl::Milliseconds(daemon_context.config()
                                        .condition_validation_options()
                                        .validation_retry_interval_ms()),
                 kMinRetryInterval);
    std::string task_name = absl::StrCat(
        "ReloadAndInsertAction",
        FlightRecordWrapper(
            saved_action_record.actions().original_request().flight_record()));
    absl::Status delay_status = daemon_context.scheduler().DelayCall(
        [this, saved_action_record = std::move(saved_action_record)]() mutable {
          std::unique_ptr<ActionContext> context_to_cleanup;
          {
            absl::MutexLock lock(actions_mutex_);
            absl::Status status = ReloadAndInsertAction(
                std::move(saved_action_record), context_to_cleanup);
            if (!status.ok()) {
              LOG(ERROR) << "Failed to reload action: " << status;
            }
          }
        },
        retry_interval, task_name);
    if (!delay_status.ok()) {
      LOG(DFATAL) << "Failed to delay reload action: " << delay_status;
      return delay_status;
    }
    return absl::OkStatus();
  }
  reloading_actions_.erase(flight_record);

  if (!validation_result.ok()) {
    LOG(ERROR) << "Request validation failed: " << validation_result;
    return absl::OkStatus();
  }

  absl::StatusOr<std::unique_ptr<ActionContext>> action_context_or =
      ReinitActionContext(saved_action_record.actions());
  if (!action_context_or.ok()) {
    LOG(ERROR) << "Failed to reinit action: " << action_context_or.status();
    return absl::OkStatus();
  }

  std::unique_ptr<ActionContext> action_context = *std::move(action_context_or);
  auto [it, inserted] =
      running_actions_.try_emplace(flight_record, std::move(action_context));
  if (!inserted) {
    // action_context was not moved by try_emplace, so add it to
    // context_to_cleanup.
    context_to_cleanup = std::move(action_context);
    return absl::AlreadyExistsError(
        absl::StrCat("Duplicate action ID: ", flight_record));
  }
  if (!read_only_mode_) {
    absl::Status status = it->second->Activate();
    if (!status.ok()) {
      context_to_cleanup = std::move(running_actions_.extract(it).mapped());
      return status;
    }
  }
  return absl::OkStatus();
}

absl::Status ActionContextManager::LoadSavedActions() {
  ASSIGN_OR_RETURN(
      SavedActions saved_actions,
      DaemonContext::Get().persistent_storage_manager().ReadSavedActions());
  // This will hold the failed contexts so that they can be destroyed
  // after actions_mutex_ is released.
  std::vector<std::unique_ptr<ActionContext>> contexts_to_cleanup;
  {
    absl::MutexLock lock(actions_mutex_);
    for (auto& saved_action_record : saved_actions.saved_action_records()) {
      if (!saved_action_record.has_actions() ||
          !saved_action_record.actions().has_original_request() ||
          !saved_action_record.actions()
               .original_request()
               .has_flight_record()) {
        LOG(ERROR) << "Saved action record is missing flight record: "
                   << saved_action_record.DebugString();
        continue;
      }
      std::unique_ptr<ActionContext> context_to_cleanup;
      if (absl::Status status =
              ReloadAndInsertAction(saved_action_record, context_to_cleanup);
          !status.ok()) {
        LOG(ERROR) << "Failed to reload saved action: " << status
                   << "; Action Record: "
                   << saved_action_record.ShortDebugString();
      }
      if (context_to_cleanup != nullptr) {
        contexts_to_cleanup.push_back(std::move(context_to_cleanup));
      }
    }
  }
  return absl::OkStatus();
}

absl::StatusOr<ActionContext::Action> ActionContextManager::ReserveAction(
    const safepower_agent_proto::Action& action) {
  auto found_action_impl = actions_.find(action);
  if (found_action_impl == actions_.end()) {
    return absl::NotFoundError(absl::StrCat("Action not found: ", action));
  }
  if (found_action_impl->second == nullptr) {
    return absl::FailedPreconditionError(
        absl::StrCat("Action already started: ", action));
  }
  return std::move(found_action_impl->second);
}

absl::Status ActionContextManager::validateFlightRecord(
    const FlightRecordRequest& record) {
  if (!record.has_flight_name()) {
    return absl::InvalidArgumentError("Flight ID is required");
  }
  if (!record.has_step_id()) {
    return absl::InvalidArgumentError("Step ID is required");
  }
  if (record.flight_name().empty()) {
    return absl::InvalidArgumentError("Flight ID must not be empty");
  }
  if (record.step_id().empty()) {
    return absl::InvalidArgumentError("Step ID must not be empty");
  }
  return absl::OkStatus();
}

absl::Status ActionContextManager::ValidateRequest(
    const safepower_agent_proto::StartActionRequest& request,
    absl::Time start_time) {
  RETURN_IF_ERROR(validateFlightRecord(request.flight_record()));
  auto& daemon_context = DaemonContext::Get();
  auto& offline_node_entities = daemon_context.offline_node_entities();
  if (offline_node_entities.entity_tag().empty()) {
    return absl::FailedPreconditionError("No offline node entities found");
  }
  return global_state_updater_->WithState(
      [&](const safepower_agent_proto::SystemState& state) {
        return ::safepower_agent::ValidateRequest(
            request, offline_node_entities.entity_tag(), state, start_time,
            daemon_context.config().condition_validation_options());
      });
}

absl::StatusOr<ActionContext* > ActionContextManager::StartAction(
    safepower_agent_proto::StartActionRequest request) {
  // This will hold a context that fails activation so that it can be destroyed
  // after actions_mutex_ is released.
  std::unique_ptr<ActionContext> context_to_destroy;
  RETURN_IF_ERROR(ValidateRequest(request, DaemonContext::Get().now()));
  FlightRecordWrapper flight_record(request.flight_record());
  {
    absl::MutexLock lock(actions_mutex_);
    ASSIGN_OR_RETURN(ActionContext::Action action_impl,
                     ReserveAction(request.action()));
    auto action_context = std::make_unique<ActionContext>(
        ActionContext::CreationToken{}, *this, std::move(request),
        std::move(action_impl));

    auto [it, inserted] =
        running_actions_.try_emplace(flight_record, std::move(action_context));
    if (!inserted) {
      // If emplace fails, action_context is not moved.
      // Move to context_to_destroy to delay destruction until after lock is
      // released.
      context_to_destroy = std::move(action_context);
      return absl::AlreadyExistsError(
          absl::StrCat("Duplicate flight record: ", flight_record));
    }
    absl::Status status = it->second->Activate();
    if (!status.ok()) {
      // If activation fails, we must remove the action context from
      // running_actions_ but destroy it *after* releasing actions_mutex_ to
      // avoid deadlock.
      // We are expecting C++ to clean up local variables LIFO.
      context_to_destroy = std::move(running_actions_.extract(it).mapped());
      return status;
    }
    return it->second.get();
  }
}

void ActionContextManager::FinishAction(
    const safepower_agent_proto::Action& action,
    ActionContext::Action action_impl) {
  absl::MutexLock lock(actions_mutex_);
  auto found_action_impl = actions_.find(action);
  if (found_action_impl == actions_.end()) {
    LOG(DFATAL) << "Action not found: " << action;
    return;
  }
  if (found_action_impl->second != nullptr) {
    LOG(DFATAL) << "Action not started: " << action;
    return;
  }
  found_action_impl->second = std::move(action_impl);
}

std::vector<FlightRecordWrapper> ActionContextManager::ListRunningActions()
    const {
  absl::MutexLock lock(actions_mutex_);
  std::vector<FlightRecordWrapper> running_actions;
  for (const auto& [flight_record, _] : running_actions_) {
    running_actions.push_back(flight_record);
  }
  return running_actions;
}
void ActionContextManager::DumpDebugInfo(std::ostream& os) const {
  auto running_actions = ListRunningActions();
  absl::MutexLock lock(actions_mutex_);
  os << "Running actions: " << running_actions.size() << "\n";
  for (const auto& flight_record : running_actions) {
    auto found_action_context = running_actions_.find(flight_record);
    if (found_action_context == running_actions_.end()) {
      LOG(ERROR) << "Flight record not found: " << flight_record;
      continue;
    }
    found_action_context->second->DumpDebugInfo(os);
  }
}

ActionContext*  ActionContextManager::GetActionContext(
    const FlightRecordRequest& flight_record_request) {
  absl::Status status = validateFlightRecord(flight_record_request);
  if (!status.ok()) {
    LOG(ERROR) << "Invalid flight record: " << status;
    return nullptr;
  }

  absl::MutexLock lock(actions_mutex_);
  FlightRecordWrapper flight_record(flight_record_request);
  auto found_action_context = running_actions_.find(flight_record);
  if (found_action_context == running_actions_.end()) {
    LOG(ERROR) << "Flight record not found: " << flight_record;
    return nullptr;
  }
  return found_action_context->second.get();
}

bool ActionContextManager::IsActionReloading(
    const borg_mgmt::node_proxy::safepower::utils::FlightRecordRequest&
        flight_record_request) const {
  absl::MutexLock lock(actions_mutex_);
  FlightRecordWrapper flight_record(flight_record_request);
  return reloading_actions_.contains(flight_record);
}

void ActionContextManager::GetSupportedActions(
    google::protobuf::RepeatedPtrField<safepower_agent_proto::Action>& actions) const {
  absl::MutexLock lock(actions_mutex_);
  for (const auto& [action, _] : actions_) {
    *actions.Add() = action;
  }
}

}  // namespace safepower_agent