// Copyright 2026 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.
//

#ifndef THIRD_PARTY_CREDENTIO_ASYNC_ASYNC_CONTEXT_H_
#define THIRD_PARTY_CREDENTIO_ASYNC_ASYNC_CONTEXT_H_
#include <utility>

#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "async/cancellation_token.h"
#include "async/executor.h"

namespace credentio {

// Provides an execution context for asynchronous operations.
class AsyncContext {
 public:
  explicit AsyncContext() = default;

  // The caller retains ownership of `executor`, which must outlive the
  // `AsyncContext`.
  explicit AsyncContext(Executor* absl_nullable executor)
      : executor_(executor) {}

  Executor* absl_nullable executor() const { return executor_; }

  void Cancel() { cancellation_source_.Cancel(); }

  CancellationToken GetCancellationToken() const {
    return cancellation_source_.GetToken();
  }

 private:
  Executor* absl_nullable executor_ = nullptr;
  CancellationSource cancellation_source_;
};

// Helper class to resolve a callback on the given AsyncContext using a
// StatusOr<T>. If the context is cancelled, the callback is executed
// inline with a CancelledError. If the context (or its executor) is missing, it
// resolves inline.
template <typename T>
class AsyncResolver {
 public:
  AsyncResolver(AsyncContext context,
                absl::AnyInvocable<void(absl::StatusOr<T>) &&> cb)
      : context_(std::move(context)), cb_(std::move(cb)) {}

  AsyncResolver(AsyncResolver&&) = default;
  AsyncResolver& operator=(AsyncResolver&&) = default;
  AsyncResolver(const AsyncResolver&) = delete;
  AsyncResolver& operator=(const AsyncResolver&) = delete;

  void Return(absl::StatusOr<T> result) {
    if (context_.GetCancellationToken().IsCancelled()) {
      std::move(cb_)(absl::CancelledError("AsyncContext cancelled"));
      return;
    }
    if (context_.executor()) {
      context_.executor()->Schedule(
          [cb = std::move(cb_), r = std::move(result)]() mutable {
            std::move(cb)(std::move(r));
          });
      return;
    }
    std::move(cb_)(std::move(result));
  }

 private:
  AsyncContext context_;
  absl::AnyInvocable<void(absl::StatusOr<T>) &&> cb_;
};

// Specialization for void callbacks (absl::Status).
template <>
class AsyncResolver<void> {
 public:
  AsyncResolver(AsyncContext context,
                absl::AnyInvocable<void(absl::Status) &&> cb)
      : context_(std::move(context)), cb_(std::move(cb)) {}

  AsyncResolver(AsyncResolver&&) = default;
  AsyncResolver& operator=(AsyncResolver&&) = default;
  AsyncResolver(const AsyncResolver&) = delete;
  AsyncResolver& operator=(const AsyncResolver&) = delete;

  void Return(absl::Status status) {
    if (context_.GetCancellationToken().IsCancelled()) {
      std::move(cb_)(absl::CancelledError("AsyncContext cancelled"));
      return;
    }
    if (context_.executor()) {
      context_.executor()->Schedule(
          [cb = std::move(cb_), s = std::move(status)]() mutable {
            std::move(cb)(std::move(s));
          });
      return;
    }
    std::move(cb_)(std::move(status));
  }

 private:
  AsyncContext context_;
  absl::AnyInvocable<void(absl::Status) &&> cb_;
};

}  // namespace credentio

#endif  // THIRD_PARTY_CREDENTIO_ASYNC_ASYNC_CONTEXT_H_