#include "subscription_tracker.h"

#include <optional>

#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.h"
#include "config_parser.h"

namespace milotic::authz {

namespace {
constexpr int kDefaultSubscriptionLimit = 1;
}  // namespace

SubscriptionTracker::SubscriptionTracker()
    : SubscriptionTracker(kDefaultSubscriptionLimit) {}

absl::Status SubscriptionTracker::RecordNewUnsubscription(
    const PeerSpiffeIdentity& peer) {
  std::optional<AuthzConfiguration::SpiffeIdentityMatcher> peer_matcher =
      AuthzConfiguration::CreateIdentityMatcherFromPeerSpiffeIdentity(peer);

  if (!peer_matcher) {
    return absl::InvalidArgumentError(
        absl::StrCat("Peer spiffe id '", peer.spiffe_id, "' is malformed"));
  }

  absl::MutexLock lock(mutex_);
  auto it = peer_to_subscription_count_.find(*peer_matcher);
  if (it == peer_to_subscription_count_.end()) {
    return absl::OkStatus();
  }

  it->second--;
  if (it->second == 0) {
    peer_to_subscription_count_.erase(it);
  }

  return absl::OkStatus();
}

absl::Status SubscriptionTracker::RecordNewSubscription(
    const PeerSpiffeIdentity& peer) {
  std::optional<AuthzConfiguration::SpiffeIdentityMatcher> peer_matcher =
      AuthzConfiguration::CreateIdentityMatcherFromPeerSpiffeIdentity(peer);

  if (!peer_matcher) {
    return absl::InvalidArgumentError(
        absl::StrCat("Peer spiffe id '", peer.spiffe_id, "' is malformed"));
  }

  absl::MutexLock lock(mutex_);
  auto it = peer_to_subscription_count_.find(*peer_matcher);
  if (it == peer_to_subscription_count_.end()) {
    peer_to_subscription_count_.insert({*peer_matcher, 1});
    return absl::OkStatus();
  }

  if (it->second >= universal_subscription_limit_) {
    return absl::PermissionDeniedError(absl::Substitute(
        "Peer '$0' has reached the limit; the limit is '$1'; currently used "
        "'$2'",
        peer.DebugString(), universal_subscription_limit_, it->second));
  }
  it->second++;
  return absl::OkStatus();
}

}  // namespace milotic::authz