#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_AUTHZ_SUBSCRIPTION_TRACKER_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_AUTHZ_SUBSCRIPTION_TRACKER_H_

#include <cstdint>
#include <map>

#include "absl/base/thread_annotations.h"
#include "absl/status/status.h"
#include "absl/synchronization/mutex.h"
#include "config_parser.h"

namespace milotic::authz {

class SubscriptionTracker {
 public:
  SubscriptionTracker();
  explicit SubscriptionTracker(int universal_subscription_limit)
      : universal_subscription_limit_(universal_subscription_limit) {}

  // Records a new subscription for the given `peer`.
  // Returns error if the client has reached to the subscription limit.
  absl::Status RecordNewSubscription(const PeerSpiffeIdentity& peer);

  // Records the an unsubscription for the given `peer`.
  absl::Status RecordNewUnsubscription(const PeerSpiffeIdentity& peer);

  // Returns the subscription limit for the given `peer`. Currently returns the
  // universal subscription limit.
  int64_t GetSubscriptionLimit(const PeerSpiffeIdentity&) {
    absl::MutexLock lock(&mutex_);
    return universal_subscription_limit_;
  }

  // Sets the limit for the universal subscription limit.
  // Will be called when authorization policy is reloaded.
  void SetUniversalSubscriptionLimit(int64_t limit) {
    absl::MutexLock lock(&mutex_);
    universal_subscription_limit_ = limit;
  }

 private:
  absl::Mutex mutex_;
  int64_t universal_subscription_limit_ ABSL_GUARDED_BY(mutex_) = 0;
  std::map<AuthzConfiguration::SpiffeIdentityMatcher, int64_t>
      peer_to_subscription_count_ ABSL_GUARDED_BY(mutex_);
};
}  // namespace milotic::authz

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_AUTHZ_SUBSCRIPTION_TRACKER_H_