// 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 <memory>
#include <type_traits>
#include <utility>

#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.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_; }

  // Attaches an extension of type `T` to this context, replacing any existing
  // extension of the same type. Passing nullptr removes the extension.
  //
  // Because `AsyncContext` is threaded through asynchronous calls and captured
  // into continuations, attached extensions survive executor hops.
  //
  // Note: Lookup in `GetExtension<T>()` is keyed by the exact type `T`. When
  // attaching an implementation of an interface, specify the interface type
  // explicitly (e.g. `context.SetExtension<MyInterface>(impl)`) so consumers
  // can retrieve it via `context.GetExtension<MyInterface>()`.
  template <typename T>
  void SetExtension(absl_nullable std::shared_ptr<T> extension) {
    using KeyType = std::remove_cv_t<T>;
    if (extension == nullptr) {
      extensions_.erase(ExtensionKey<KeyType>());
      return;
    }
    extensions_[ExtensionKey<KeyType>()] = std::move(extension);
  }

  // Returns the extension of type `T` attached to this context, or nullptr if
  // none is present.
  template <typename T>
  absl_nullable std::shared_ptr<T> GetExtension() const {
    using KeyType = std::remove_cv_t<T>;
    auto it = extensions_.find(ExtensionKey<KeyType>());
    if (it == extensions_.end()) {
      return nullptr;
    }
    return std::static_pointer_cast<T>(it->second);
  }

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

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

 private:
  template <typename T>
  static const void* ExtensionKey() {
    static const char kKey = 0;
    return &kKey;
  }

  Executor* absl_nullable executor_ = nullptr;
  CancellationSource cancellation_source_;
  absl::flat_hash_map<const void*, std::shared_ptr<void>> extensions_;
};

// 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_