// Copyright 2022 Google LLC
//
// 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.

#include "google/cloud/pubsub/internal/extend_leases_with_retry.h"
#include "google/cloud/pubsub/internal/exactly_once_policies.h"
#include "google/cloud/completion_queue.h"
#include "google/cloud/log.h"
#include "absl/strings/match.h"
#include <algorithm>
#include <memory>

namespace google {
namespace cloud {
namespace pubsub_internal {
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
namespace {

// This value is hard-coded in the spec.
auto constexpr kMaxAttempts = 3;

/// Updates a ModifyAckDeadlineRequest request to drop ack_ids with permanent
/// failures, and successful ack_ids.
google::pubsub::v1::ModifyAckDeadlineRequest UpdateRequest(
    google::pubsub::v1::ModifyAckDeadlineRequest request,
    Status const& status) {
  if (ExactlyOnceRetryable(status.code())) return request;

  auto permanent_error = [](std::string const& ack_id, Status const& status) {
    auto const& m = status.error_info().metadata();
    auto f = m.find(ack_id);
    auto const permanent =
        f == m.end() || !absl::StartsWith(f->second, "TRANSIENT_FAILURE_");
    if (permanent) {
      GCP_LOG(WARNING)
          << "permanent failure trying to extend the lease for ack_id="
          << ack_id << ", status=" << status;
    }
    return permanent;
  };
  auto& ids = *request.mutable_ack_ids();
  ids.erase(
      std::remove_if(ids.begin(), ids.end(),
                     [&](auto const& x) { return permanent_error(x, status); }),
      ids.end());
  return request;
}

/**
 * Implements `ExtendLeasesWithRetry()`.
 *
 * If we were using C++20, then `ExtendLeasesWithRetry()` would be a coroutine,
 * and this class would be automatically generated by the compiler and called
 * "the coroutine handle".  We need to support C++ >= 14, but we can use the
 * lingo from C++20 coroutines to describe how this works.
 *
 * Recall that with coroutines, all local variables become member variables of
 * the handle class, and each `co_await` becomes a callback.
 *
 * The coroutine would be:
 *
 * @code
 * future<Status> ExtendLeasesWithRetry(
 *     std::shared_ptr<SubscriberStub> stub, CompletionQueue cq,
 *     google::pubsub::v1::ModifyAckDeadlineRequest request) {
 *   auto backoff = google::cloud::pubsub::ExponentialBackoff(
 *       std::chrono::seconds(1), std::chrono::seconds(10), 2.0);
 *   Status last_status;
 *   for (int attempts = 0; attempts != 3; ++i) {
 *     last_status = co_await stub->AsyncModifyAckDeadline(
 *         cq, std::make_shared<grpc::ClientContext>(), request);
 *     if (status.ok()) co_return last_status;
 *     request = UpdateRequest(std::move(request), status);
 *     if (request.ack_ids().empty()) co_return last_status;
 *     co_await cq.MakeRelativeTimer(backoff.OnCompletion());
 *   }
 *   co_return last_status;
 * }
 * @enccode
 */
class ExtendLeasesWithRetryHandle
    : public std::enable_shared_from_this<ExtendLeasesWithRetryHandle> {
 public:
  ExtendLeasesWithRetryHandle(
      std::shared_ptr<SubscriberStub> stub, CompletionQueue cq,
      google::pubsub::v1::ModifyAckDeadlineRequest request)
      : stub_(std::move(stub)),
        cq_(std::move(cq)),
        request_(std::move(request)) {}

  future<Status> Start() {
    auto result = promise_.get_future();
    MakeAttempt();
    return result;
  }

 private:
  void MakeAttempt() {
    ++attempts_;
    auto context = std::make_shared<grpc::ClientContext>();
    stub_->AsyncModifyAckDeadline(cq_, std::move(context), request_)
        .then(
            [self = shared_from_this()](auto f) { self->OnAttempt(f.get()); });
  }

  void OnAttempt(Status status) {
    last_status_ = std::move(status);
    if (last_status_.ok()) {
      promise_.set_value(std::move(last_status_));
      return;
    }
    request_ = UpdateRequest(std::move(request_), last_status_);
    if (request_.ack_ids().empty()) {
      promise_.set_value(std::move(last_status_));
      return;
    }
    if (attempts_ >= kMaxAttempts) {
      for (auto const& a : request_.ack_ids()) {
        GCP_LOG(WARNING)
            << "retry policy exhausted while trying to extend lease for ack_id="
            << a << ", last_status=" << last_status_;
      }
      promise_.set_value(std::move(last_status_));
      return;
    }
    cq_.MakeRelativeTimer(backoff_->OnCompletion())
        .then([self = shared_from_this()](auto f) {
          self->OnBackoff(std::move(f.get()).status());
        });
  }

  void OnBackoff(Status status) {
    if (!status.ok()) {
      promise_.set_value(std::move(status));
      return;
    }
    return MakeAttempt();
  }

  promise<Status> promise_;
  std::shared_ptr<SubscriberStub> stub_;
  CompletionQueue cq_;
  google::pubsub::v1::ModifyAckDeadlineRequest request_;
  std::unique_ptr<pubsub::BackoffPolicy> backoff_ = ExactlyOnceBackoffPolicy();
  Status last_status_ = Status(StatusCode::kUnknown, "Exhausted before start");
  int attempts_ = 0;
};

}  // namespace

future<Status> ExtendLeasesWithRetry(
    std::shared_ptr<SubscriberStub> stub, CompletionQueue cq,
    google::pubsub::v1::ModifyAckDeadlineRequest request) {
  auto handle = std::make_shared<ExtendLeasesWithRetryHandle>(
      std::move(stub), std::move(cq), std::move(request));
  return handle->Start();
}

GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
}  // namespace pubsub_internal
}  // namespace cloud
}  // namespace google
