// 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_CANCELLATION_TOKEN_H_
#define THIRD_PARTY_CREDENTIO_ASYNC_CANCELLATION_TOKEN_H_

#include <cstdint>
#include <memory>
#include <utility>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/synchronization/mutex.h"

namespace credentio {

class CancellationToken;

// Internal shared state for cancellation.
class CancellationState {
 public:
  CancellationState() = default;

  bool IsCancelled() const {
    absl::MutexLock lock(mutex_);
    return cancelled_;
  }

  void Cancel();

  // Returns >0 registration ID if successfully registered.
  // Returns 0 if already cancelled (and executes callback immediately inline).
  uint64_t RegisterCallback(absl::AnyInvocable<void()> cb);

  void DeregisterCallback(uint64_t id);

 private:
  mutable absl::Mutex mutex_;
  bool cancelled_ ABSL_GUARDED_BY(mutex_) = false;
  uint64_t next_id_ ABSL_GUARDED_BY(mutex_) = 1;
  absl::flat_hash_map<uint64_t, absl::AnyInvocable<void()>> callbacks_
      ABSL_GUARDED_BY(mutex_);
};

// Represents a read-only view of a cancellation state.
class CancellationToken {
 public:
  CancellationToken() = default;

  bool IsCancelled() const {
    if (!state_) return false;
    return state_->IsCancelled();
  }

  // Copyable and movable.
  CancellationToken(const CancellationToken&) = default;
  CancellationToken& operator=(const CancellationToken&) = default;
  CancellationToken(CancellationToken&&) = default;
  CancellationToken& operator=(CancellationToken&&) = default;

 private:
  friend class CancellationSource;
  friend class CancellationCallback;

  explicit CancellationToken(std::shared_ptr<CancellationState> state)
      : state_(std::move(state)) {}

  std::shared_ptr<CancellationState> state_;
};

// The source of cancellation.
class CancellationSource {
 public:
  CancellationSource() : state_(std::make_shared<CancellationState>()) {}

  void Cancel() {
    if (state_) {
      state_->Cancel();
    }
  }

  CancellationToken GetToken() const { return CancellationToken(state_); }

 private:
  std::shared_ptr<CancellationState> state_;
};

// Represents a registration of a callback against a CancellationToken.
// When this object falls out of scope, the callback is automatically
// deregistered.
class CancellationCallback {
 public:
  CancellationCallback(CancellationToken token, absl::AnyInvocable<void()> cb)
      : state_(std::move(token.state_)) {
    if (state_) {
      id_ = state_->RegisterCallback(std::move(cb));
    }
  }

  ~CancellationCallback() {
    if (state_ && id_ != 0) {
      state_->DeregisterCallback(id_);
    }
  }

  // Non-copyable, non-movable.
  CancellationCallback(const CancellationCallback&) = delete;
  CancellationCallback& operator=(const CancellationCallback&) = delete;

 private:
  std::shared_ptr<CancellationState> state_;
  uint64_t id_ = 0;
};

}  // namespace credentio

#endif  // THIRD_PARTY_CREDENTIO_ASYNC_CANCELLATION_TOKEN_H_