#include "condition.h"

#include <concepts>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "google/protobuf/timestamp.pb.h"
#include "convert_proto.h"
#include "daemon_context.h"
#include "safepower_agent.pb.h"
#include "state_fields.h"
#include "state_merge.h"
#include "state_updater.h"
#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/time/time.h"
#include "google/protobuf/repeated_ptr_field.h"
#include "bmc/status_macros.h"

namespace safepower_agent {

using ::safepower_agent_proto::AllOfCondition;
using ::safepower_agent_proto::AnyOfCondition;
using ::safepower_agent_proto::ComparisonType;
using ::safepower_agent_proto::NodeState;
using ::safepower_agent_proto::StateMatchedCondition;
using ::safepower_agent_proto::SystemState;
using ::safepower_agent_proto::TimeoutCondition;

template <typename T>
absl::StatusOr<bool> Compare(ComparisonType comparison_type, T&& a,
                             T&& expected) {
  switch (comparison_type) {
    case ComparisonType::COMPARISON_TYPE_EQ:
      return a == expected;
    case ComparisonType::COMPARISON_TYPE_NE:
      return a != expected;
    case ComparisonType::COMPARISON_TYPE_GT:
      return a > expected;
    case ComparisonType::COMPARISON_TYPE_GE:
      return a >= expected;
    case ComparisonType::COMPARISON_TYPE_LT:
      return a < expected;
    case ComparisonType::COMPARISON_TYPE_LE:
      return a <= expected;
    default:
      return absl::InvalidArgumentError(
          absl::StrCat("Unknown comparison type: ", comparison_type));
      return false;
  }
}

template <typename Proto>
  requires(std::derived_from<Proto, google::protobuf::Message>)
absl::StatusOr<bool> Compare(ComparisonType comparison_type, const Proto& a,
                             const Proto& expected);

static bool CombineResults(absl::StatusOr<bool>& result,
                           const absl::StatusOr<bool>& new_result) {
  if (absl::IsNotFound(new_result.status())) {
    return true;
  }
  result = new_result;
  return result.ok() && *result;
}

template <typename Proto, typename FieldT>
absl::StatusOr<bool> CompareField(const MapFieldDef<Proto, FieldT>& field,
                                  ComparisonType comparison_type,
                                  const Proto& a, const Proto& expected) {
  absl::StatusOr<bool> result = absl::NotFoundError("No matching keys");
  for (const auto& [key, value] : field.get_from(expected)) {
    auto a_it = field.get_from(a).find(key);
    if (!CombineResults(
            result, (a_it == field.get_from(a).end()
                         ? absl::NotFoundError("Key not found")
                         : Compare(comparison_type, a_it->second, value)))) {
      break;
    }
  }
  return result;
}

template <typename Proto, typename FieldDefT>
absl::StatusOr<bool> CompareField(const FieldDefT& field,
                                  ComparisonType comparison_type,
                                  const Proto& a, const Proto& expected) {
  if (field.is_in(a) && field.is_in(expected)) {
    return Compare(comparison_type, field.get_from(a),
                   field.get_from(expected));
  } else {
    return absl::NotFoundError("Field not set");
  }
}

template <typename Proto>
  requires(std::derived_from<Proto, google::protobuf::Message>)
absl::StatusOr<bool> Compare(ComparisonType comparison_type, const Proto& a,
                             const Proto& expected) {
  return std::apply(
      [&](auto&&... field_defs) {
        absl::StatusOr<bool> result = absl::NotFoundError("No matching fields");
        (CombineResults(
             result, CompareField(field_defs, comparison_type, a, expected)) &&
         ...);
        if (absl::IsNotFound(result.status())) {
          VLOG(1) << "No matches found: " << result.status();
          result = false;
        }
        return result;
      },
      ProtoFields<Proto>::kFields);
}

using MatchResult = std::tuple<absl::Status, Condition::MatchList>;

static MatchResult ConditionMatches(
    const safepower_agent_proto::Condition& condition,
    const safepower_agent_proto::SystemState& system_state,
    absl::Time start_time);

static MatchResult ConditionMatches(const AnyOfCondition& any_of,
                                    const SystemState& system_state,
                                    absl::Time start_time) {
  Condition::MatchList result;
  absl::Status status = absl::OkStatus();
  for (const safepower_agent_proto::Condition& condition :
       any_of.conditions()) {
    auto [match_status, match_list] =
        ConditionMatches(condition, system_state, start_time);
    result.insert(result.end(), match_list.begin(), match_list.end());
    status.Update(match_status);
  }
  return {status, result};
}

static MatchResult ConditionMatches(const AllOfCondition& all_of,
                                    const SystemState& system_state,
                                    absl::Time start_time) {
  bool matches_all = true;
  absl::Status status = absl::OkStatus();
  Condition::MatchList result;
  for (const safepower_agent_proto::Condition& condition :
       all_of.conditions()) {
    auto [match_status, match_list] =
        ConditionMatches(condition, system_state, start_time);
    result.insert(result.end(), match_list.begin(), match_list.end());
    if (match_list.empty()) {
      // Make sure we don't return early if there are any other conditions. It
      // is possible that the other conditions will return an error, which will
      // override the return value.
      matches_all = false;
    }
    status.Update(match_status);
  }
  if (!matches_all && status.ok()) {
    result.clear();
  }
  return {status, result};
}

static absl::StatusOr<bool> ConditionMatches(
    const StateMatchedCondition& condition, const SystemState& system_state) {
  auto node_state_found =
      system_state.node_state().find(condition.node_entity_tag());
  if (node_state_found == system_state.node_state().end()) {
    return false;
  }

  ASSIGN_OR_RETURN(bool matches,
                   Compare(condition.comparison_type(),
                           node_state_found->second, condition.node_state()));
  if (matches && condition.abort()) {
    return absl::AbortedError(
        absl::StrCat("Abort condition matched: ", condition));
  }
  return matches;
}

static absl::StatusOr<bool> ConditionMatches(const TimeoutCondition& condition,
                                             absl::Time start_time) {
  bool matches = false;
  switch (condition.time_case()) {
    case TimeoutCondition::kTimeoutMs:
      matches = start_time + absl::Milliseconds(condition.timeout_ms()) <=
                DaemonContext::Get().now();
      break;
    case TimeoutCondition::kDeadline:
      matches = ConvertTime(condition.deadline()) <= DaemonContext::Get().now();
      break;
    default:
      return absl::InvalidArgumentError(
          absl::StrCat("Unknown timeout type: ", condition));
  }
  if (matches && condition.abort()) {
    return absl::AbortedError(
        absl::StrCat("Abort condition matched: ", condition));
  }
  return matches;
}

static MatchResult GetMatchResult(
    const safepower_agent_proto::Condition& condition,
    absl::StatusOr<bool> matches) {
  if (!matches.ok()) {
    return {matches.status(), {&condition}};
  }
  if (*matches) {
    return {absl::OkStatus(), {&condition}};
  }
  return {absl::OkStatus(), {}};
}

static MatchResult ConditionMatches(
    const safepower_agent_proto::Condition& condition,
    const safepower_agent_proto::SystemState& system_state,
    absl::Time start_time) {
  switch (condition.condition_type_case()) {
    case safepower_agent_proto::Condition::kAnyOf:
      return ConditionMatches(condition.any_of(), system_state, start_time);
    case safepower_agent_proto::Condition::kAllOf:
      return ConditionMatches(condition.all_of(), system_state, start_time);
    case safepower_agent_proto::Condition::kStateCondition:
      return GetMatchResult(
          condition,
          ConditionMatches(condition.state_condition(), system_state));
    case safepower_agent_proto::Condition::kTimeout:
      return GetMatchResult(condition,
                            ConditionMatches(condition.timeout(), start_time));
    default:
      return {absl::InvalidArgumentError(
                  absl::StrCat("Unknown condition type: ", condition)),
              {&condition}};
  }
}

Condition::Condition(const safepower_agent_proto::Condition& condition,
                     absl::Time original_start_time)
    : listener_(*this),
      condition_(condition),
      start_time_(original_start_time) {}

Condition::~Condition() {
  for (const std::string& job_name : pending_jobs_) {
    absl::Status status = DaemonContext::Get().scheduler().CancelCall(job_name);
    if (!status.ok() && !absl::IsNotFound(status)) {
      LOG(ERROR) << "Failed to cancel job: " << status;
    }
  }
}

absl::Status Condition::WaitForMatch(
    absl::string_view unique_id,
     std::shared_ptr<StateUpdater<SystemState>> state_updater,
    absl::AnyInvocable<void(absl::Status, Condition::MatchList) &&> callback) {
  complete_job_name_ = absl::StrCat(unique_id, ".complete");
  state_updater_ = std::move(state_updater);
  if (callback == nullptr) {
    return absl::InvalidArgumentError("Callback must not be null");
  }
  callback_ = std::move(callback);
  RETURN_IF_ERROR(RegisterTimeouts(condition_, unique_id));
  listener_.Listen(state_updater_);
  return absl::OkStatus();
}

std::tuple<absl::Status, Condition::MatchList> Condition::CheckState(
    const SystemState& state) {
  return ConditionMatches(condition_, state, start_time_);
}

absl::Status Condition::TriggerCallback(absl::Status match_status,
                                        MatchList match_list) {
  if (callback_ == nullptr) {
    return absl::FailedPreconditionError("Condition is not active");
  }
  return DaemonContext::Get().scheduler().DelayCall(
      [callback = std::move(callback_), match_status = std::move(match_status),
       match_list = std::move(match_list)]() mutable {
        std::move(callback)(std::move(match_status), std::move(match_list));
      },
      absl::ZeroDuration(), complete_job_name_);
}

void Condition::Listener::UpdateState(const SystemState& current_state,
                                      const SystemState& update) {
  auto [match_status, match_list] =
      condition_.CheckState(MergeState(current_state, update));
  if (!match_status.ok() || !match_list.empty()) {
    absl::Status status = condition_.TriggerCallback(std::move(match_status),
                                                     std::move(match_list));
    if (!status.ok()) {
      LOG(ERROR) << "Failed to trigger callback: " << status;
    }
  }
}

absl::Status Condition::RegisterTimeouts(
    const safepower_agent_proto::Condition& condition,
    absl::string_view job_name_prefix) {
  switch (condition.condition_type_case()) {
    case safepower_agent_proto::Condition::kAnyOf:
      RETURN_IF_ERROR(
          RegisterTimeouts(condition.any_of().conditions(), job_name_prefix));
      break;
    case safepower_agent_proto::Condition::kAllOf:
      RETURN_IF_ERROR(
          RegisterTimeouts(condition.all_of().conditions(), job_name_prefix));
      break;
    case safepower_agent_proto::Condition::kStateCondition:
      break;
    case safepower_agent_proto::Condition::kTimeout: {
      absl::Time deadline;
      switch (condition.timeout().time_case()) {
        case TimeoutCondition::kTimeoutMs:
          deadline = start_time_ +
                     absl::Milliseconds(condition.timeout().timeout_ms());
          break;
        case TimeoutCondition::kDeadline:
          deadline = ConvertTime(condition.timeout().deadline());
          break;
        default:
          return absl::InvalidArgumentError(
              absl::StrCat("Unknown timeout type: ", condition));
      }
      pending_jobs_.push_back(absl::StrCat(job_name_prefix, ".timer"));
      RETURN_IF_ERROR(DaemonContext::Get().scheduler().DelayCall(
          [this] {
            // Push empty update to trigger a check.
            state_updater_->UpdateState({});
          },
          deadline - DaemonContext::Get().now(), pending_jobs_.back()));
      break;
    }
    default:
      return absl::InvalidArgumentError(
          absl::StrCat("Unknown condition type: ", condition));
  }
  return absl::OkStatus();
}

absl::Status Condition::RegisterTimeouts(
    const google::protobuf::RepeatedPtrField<safepower_agent_proto::Condition>&
        conditions,
    absl::string_view job_name_prefix) {
  int index = 0;
  for (const safepower_agent_proto::Condition& condition : conditions) {
    RETURN_IF_ERROR(RegisterTimeouts(
        condition, absl::StrCat(job_name_prefix, ".", index++)));
  }
  return absl::OkStatus();
}

}  // namespace safepower_agent