// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef CHROME_BROWSER_ASH_APP_MODE_RETRY_RUNNER_H_
#define CHROME_BROWSER_ASH_APP_MODE_RETRY_RUNNER_H_

#include <memory>
#include <optional>
#include <variant>

#include "base/functional/bind.h"
#include "base/functional/callback.h"
#include "base/memory/ptr_util.h"
#include "base/memory/weak_ptr.h"
#include "base/time/time.h"
#include "chrome/browser/ash/app_mode/cancellable_job.h"

namespace ash {

template <typename Result>
using RetryResultCallback = base::OnceCallback<void(Result result)>;
template <typename Result>
using RetryJob = base::RepeatingCallback<std::unique_ptr<CancellableJob>(
    RetryResultCallback<Result> on_result)>;
template <typename Result>
using VoidRetryJob =
    base::RepeatingCallback<void(RetryResultCallback<Result> on_result)>;
template <typename Result>
using RetryPredicate = base::RepeatingCallback<bool(const Result& result)>;

// Runs the given `job` up to `n` times.
//
// `job` will be run `n` times or until `should_retry` returns false, whichever
// happens first.
//
// `should_retry` will be invoked with the result whenever `job` finishes and
// there are attempts remaining. The result of `should_retry` determines if we
// must retry the job, or if we're done and we can pass the result to `on_done`.
//
// `on_done` will be called with the result of the last `job`.
//
// Attempts at `job` are delayed with exponential backoff after failure.
//
// Destroying the returned `std::unique_ptr` cancels this task. In that case
// `on_done` will not be called.
template <typename Result>
[[nodiscard]] std::unique_ptr<CancellableJob> RunUpToNTimes(
    int n,
    VoidRetryJob<Result> job,
    RetryPredicate<Result> should_retry,
    RetryResultCallback<Result> on_done);

// Same as above but accepts jobs that return instances of `CancellableJob`.
//
// Each `CancellableJob` instance generated by attempts at `job` are eventually
// destroyed after they return. This guarantees resources allocated by the
// execution of `job` are not kept around indefinitely after the `job` ends.
template <typename Result>
[[nodiscard]] std::unique_ptr<CancellableJob> RunUpToNTimes(
    int n,
    RetryJob<Result> job,
    RetryPredicate<Result> should_retry,
    RetryResultCallback<Result> on_done);

// Returns a predicate suitable for `RunUpToNTimes()` jobs that return an
// `std::optional` and should retry when the result is `std::nullopt`.
template <typename T>
base::RepeatingCallback<bool(const std::optional<T>&)> RetryIfNullopt();

namespace internal {

// Computes the delay for `attempt_count` with exponential backoff.
base::TimeDelta DelayForAttempt(int attempt_count);

// Runs `task` after the given `delay` in the current sequence.
void PostDelayedTask(base::OnceClosure task, base::TimeDelta delay);

// Helper to retry tasks that can fail.
template <typename Result, typename HandleType = std::monostate>
class RetryRunner : public CancellableJob {
 public:
  using Job = base::RepeatingCallback<std::unique_ptr<HandleType>(
      RetryResultCallback<Result> on_result)>;

  [[nodiscard]] static std::unique_ptr<CancellableJob> Run(
      int max_attempts,
      Job job,
      RetryPredicate<Result> should_retry,
      RetryResultCallback<Result> on_done) {
    auto handle = base::WrapUnique(new RetryRunner<Result, HandleType>(
        max_attempts, job, should_retry, std::move(on_done)));
    handle->RunAndRetryOnFailure();
    return handle;
  }

  RetryRunner(const RetryRunner&) = delete;
  RetryRunner& operator=(const RetryRunner&) = delete;
  ~RetryRunner() override = default;

 private:
  RetryRunner(int max_attempts,
              Job job,
              RetryPredicate<Result> should_retry,
              RetryResultCallback<Result> on_done)
      : max_attempts_(max_attempts),
        job_(job),
        should_retry_(should_retry),
        on_done_(std::move(on_done)) {}

  void RunAndRetryOnFailure(int attempt_count = 1) {
    current_run_ = job_.Run(base::BindOnce(
        [](base::WeakPtr<RetryRunner> self, int attempt_count, Result result) {
          if (!self) {
            return;
          }

          self->current_run_.reset();
          if (attempt_count >= self->max_attempts_ ||
              !self->should_retry_.Run(result)) {
            return std::move(self->on_done_).Run(std::move(result));
          }

          internal::PostDelayedTask(
              base::BindOnce(&RetryRunner::RunAndRetryOnFailure, self,
                             attempt_count + 1),
              internal::DelayForAttempt(attempt_count));
        },
        weak_ptr_factory_.GetWeakPtr(), attempt_count));
    if (on_done_.is_null()) {
      // `job_` already called `on_done_` synchronously, reset `current_run_`.
      current_run_.reset();
    }
  }

  int max_attempts_;
  Job job_;
  RetryPredicate<Result> should_retry_;
  RetryResultCallback<Result> on_done_;
  std::unique_ptr<HandleType> current_run_;
  base::WeakPtrFactory<RetryRunner> weak_ptr_factory_{this};
};

}  // namespace internal

template <typename Result>
[[nodiscard]] std::unique_ptr<CancellableJob> RunUpToNTimes(
    int n,
    VoidRetryJob<Result> job,
    RetryPredicate<Result> should_retry,
    RetryResultCallback<Result> on_done) {
  return internal::RetryRunner<Result>::Run(
      /*max_attempts=*/n, /*job=*/
      base::BindRepeating(
          [](VoidRetryJob<Result> job,
             RetryResultCallback<Result> result_callback) {
            job.Run(std::move(result_callback));
            return std::make_unique<std::monostate>();
          },
          job),
      should_retry, std::move(on_done));
}

template <typename Result>
[[nodiscard]] std::unique_ptr<CancellableJob> RunUpToNTimes(
    int n,
    RetryJob<Result> job,
    RetryPredicate<Result> should_retry,
    RetryResultCallback<Result> on_done) {
  return internal::RetryRunner<Result, CancellableJob>::Run(
      /*max_attempts=*/n, job, should_retry, std::move(on_done));
}

template <typename T>
base::RepeatingCallback<bool(const std::optional<T>&)> RetryIfNullopt() {
  return base::BindRepeating(
      [](const std::optional<T>& optional) { return !optional.has_value(); });
}

}  // namespace ash

#endif  // CHROME_BROWSER_ASH_APP_MODE_RETRY_RUNNER_H_