#include "config_parser.h"

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <fstream>
#include <optional>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "gmi/machine_identity.pb.h"
#include "one/node_entities_api.pb.h"
#include "one/offline_node_entities.pb.h"
#include "one/resolved_entities.pb.h"
#include "one/public_offline_node_entities.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/ascii.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "nlohmann/json.hpp"
#include "nlohmann/json_fwd.hpp"
#include "json_utils.h"
#include "authorized_entity.pb.h"
#include "zatar/certificate_metadata.h"

namespace milotic::authz {

using ::milotic::redfish::CertificateMetadataParser;
using ::production_msv::node_entities::ReadOfflineNodeEntityInformation;
using ::production_msv::node_entities_proto::OfflineNodeEntityInformation;
using ::security_prodid::GoogleMachineIdentityProto;
using EntityTagIdentity = AuthzConfiguration::EntityTagIdentity;

AuthzConfiguration::AuthzConfiguration(
    const nlohmann::json& config, absl::string_view offline_node_entities_path,
    absl::string_view gmi_path)
    : offline_node_entities_path_(offline_node_entities_path),
      gmi_path_(gmi_path) {
  ReloadConfig(config, nlohmann::json());
};

std::optional<AuthzConfiguration::SpiffeIdentityMatcher>
AuthzConfiguration::CreateIdentityMatcherFromPeerSpiffeIdentity(
    const PeerSpiffeIdentity& peer) {
  AuthzConfiguration::SpiffeIdentityMatcher matcher;
  matcher.fqdn = peer.fqdn;

  const std::string spiffe_id = absl::AsciiStrToLower(peer.spiffe_id);
  // As per SPIFFE specification, SPIFFE schema and trust domain are case
  // insensitive.
  // https://github.com/spiffe/spiffe/blob/main/standards/SPIFFE-ID.md#24-spiffe-id-parsing
  constexpr std::string_view kSpiffeSchema = "spiffe://";
  if (!absl::StartsWithIgnoreCase(spiffe_id, kSpiffeSchema)) {
    return std::nullopt;
  }

  std::size_t path_start = spiffe_id.find_first_of('/', kSpiffeSchema.size());
  if (path_start == std::string_view::npos) {
    return std::nullopt;
  }

  matcher.spiffe_id_matcher.path = spiffe_id.substr(path_start);

  const std::string trust_domain = absl::AsciiStrToLower(spiffe_id.substr(
      kSpiffeSchema.size(), path_start - kSpiffeSchema.size()));

  // Trust domain consists of (issuing-role).(context).(realm).(suffix)
  // See go/zatar-prod-certificates for details.
  std::vector<std::string_view> trust_domain_pieces =
      absl::StrSplit(trust_domain, '.');

  // There should be at least 4 elements in this array for a property trust
  // domain
  if (trust_domain_pieces.size() < 4) {
    return std::nullopt;
  }

  matcher.spiffe_id_matcher.issuer = trust_domain_pieces[0];
  matcher.spiffe_id_matcher.ca_context = trust_domain_pieces[1];
  matcher.spiffe_id_matcher.realm = trust_domain_pieces[2];

  matcher.spiffe_id_matcher.trust_domain_suffix =
      absl::StrCat(".", absl::StrJoin(trust_domain_pieces.begin() + 3,
                                      trust_domain_pieces.end(), "."));

  return matcher;
}

std::optional<AuthzConfiguration::SpiffeIdentityMatcher>
AuthzConfiguration::ParseSpiffeIdentityMatcher(const nlohmann::json& config) {
  const nlohmann::json* spiffe_id_matcher =
      GetValueAsJson(config, "spiffe_id_matcher");
  if (spiffe_id_matcher == nullptr) {
    return std::nullopt;
  }
  const std::string* trust_domain_suffix =
      GetValueAsString(*spiffe_id_matcher, "trust_domain_suffix");
  if (trust_domain_suffix == nullptr) {
    LOG(WARNING) << "Warning: trust_domain_suffix isn't a string; "
                 << "|spiffe_id_matcher|=" << spiffe_id_matcher->dump(2);
    return std::nullopt;
  }

  AuthzConfiguration::SpiffeIdentityMatcher matcher;
  matcher.spiffe_id_matcher.trust_domain_suffix =
      absl::AsciiStrToLower(*trust_domain_suffix);

  const std::string* path = GetValueAsString(*spiffe_id_matcher, "path");
  if (path == nullptr) {
    return std::nullopt;
  }
  matcher.spiffe_id_matcher.path = absl::AsciiStrToLower(*path);

  const std::string* ca_context =
      GetValueAsString(*spiffe_id_matcher, "ca_context");
  if (ca_context == nullptr) {
    return std::nullopt;
  }
  matcher.spiffe_id_matcher.ca_context = absl::AsciiStrToLower(*ca_context);

  const std::string* issuer = GetValueAsString(*spiffe_id_matcher, "issuer");
  if (issuer != nullptr) {
    matcher.spiffe_id_matcher.issuer = absl::AsciiStrToLower(*issuer);
  }

  // If realm_check is enabled, get current realm from GMI
  const bool* realm_check = GetValueAsBool(*spiffe_id_matcher, "realm_check");
  if (realm_check != nullptr && *realm_check) {
    matcher.spiffe_id_matcher.realm = GetNodeRealm();
  }

  const std::string* fqdn = GetValueAsString(config, "fqdn");
  // Non-string value
  if (fqdn == nullptr && config.find("fqdn") != config.end()) {
    return std::nullopt;
  }

  if (fqdn != nullptr) {
    matcher.fqdn = *fqdn;
  }

  return matcher;
}

bool AuthzConfiguration::IsResourceOwnerExisting(
    const ResourceOwner& resource_owner) const {
  absl::MutexLock lock(&mutex_);
  return std::binary_search(resource_owners_.begin(), resource_owners_.end(),
                            resource_owner);
}

bool AuthzConfiguration::IsPeerResourceOwner(
    const PeerSpiffeIdentity& peer) const {
  std::optional<SpiffeIdentityMatcher> peer_matcher =
      CreateIdentityMatcherFromPeerSpiffeIdentity(peer);

  if (peer_matcher == std::nullopt) {
    LOG(WARNING) << "Peer spiffe id" << peer.spiffe_id << "is malformed";
    return false;
  }
  ResourceOwner peer_to_resource_owner{.identity_matcher = *peer_matcher};

  return IsResourceOwnerExisting(peer_to_resource_owner);
}

absl::StatusOr<std::string> AuthzConfiguration::GetRedfishRoleOfSpiffeMatcher(
    const SpiffeUser& spiffe_user) const {
  absl::MutexLock lock(&mutex_);

  auto it =
      std::lower_bound(spiffe_users_.begin(), spiffe_users_.end(), spiffe_user);
  // Comparing equality
  if (it == spiffe_users_.end() || *it != spiffe_user) {
    return absl::NotFoundError(
        absl::StrCat("No user with matching spiffe; spiffe_user is ",
                     spiffe_user.identity_matcher.DebugString()));
  }

  return it->redfish_role;
}

std::string AuthzConfiguration::GetPeerRedfishRole(
    const PeerSpiffeIdentity& peer) const {
  std::optional<SpiffeIdentityMatcher> peer_matcher =
      CreateIdentityMatcherFromPeerSpiffeIdentity(peer);
  if (peer_matcher == std::nullopt) {
    LOG(WARNING) << "Peer spiffe id" << peer.spiffe_id << "is malformed";
    return "";
  }

  SpiffeUser peer_to_spiffe_user{.identity_matcher = *peer_matcher,
                                 .redfish_role = ""};

  absl::StatusOr<std::string> role =
      GetRedfishRoleOfSpiffeMatcher(peer_to_spiffe_user);

  if (!role.ok()) {
    LOG(WARNING) << role.status();
    return "";
  }

  return *role;
}

uint64_t AuthzConfiguration::GetSampleRateLimit(
    const std::string& peer_role) const {
  absl::MutexLock lock(&mutex_);
  auto it = role_to_sample_rate_limit_.find(peer_role);
  if (it != role_to_sample_rate_limit_.end()) {
    return it->second;
  }
  return 0;
}

std::unordered_set<std::string> AuthzConfiguration::GetPeerRedfishPrivileges(
    const PeerSpiffeIdentity& peer) const {
  std::optional<SpiffeIdentityMatcher> peer_matcher =
      CreateIdentityMatcherFromPeerSpiffeIdentity(peer);

  if (peer_matcher == std::nullopt) {
    LOG(WARNING) << "Peer spiffe id" << peer.spiffe_id << "is malformed";
    return {"NoAuth"};
  }

  SpiffeUser peer_to_spiffe_user{.identity_matcher = *peer_matcher,
                                 .redfish_role = ""};

  absl::StatusOr<std::string> role =
      GetRedfishRoleOfSpiffeMatcher(peer_to_spiffe_user);

  if (!role.ok()) {
    LOG(WARNING) << role.status();
    return {"NoAuth"};
  }

  return GetRedfishPrivileges(*role);
}

std::unordered_set<std::string> AuthzConfiguration::GetPeerRedfishPrivileges(
    const PeerLoasIdentity& peer) const {
  std::optional<std::string> role = GetRoleByUsername(peer.username);
  if (role) {
    return GetRedfishPrivileges(*role);
  }
  role = GetRoleByGroup(peer.username);
  if (role) {
    return GetRedfishPrivileges(*role);
  }
  return {"NoAuth"};
}

absl::string_view AuthzConfiguration::GetCertNodeOwner(
    absl::string_view node_owner) const {
  if (IsHwopsState()) {
    return kHwopsRole;
  }
  return node_owner;
}

std::vector<AuthzConfiguration::SpiffeIdentityMatcher>
AuthzConfiguration::ParseSpiffeIdentityMatchersUsingOne(
    const nlohmann::json& config) const {
  const std::string* entity = GetValueAsString(config, "entity");
  if (entity == nullptr) {
    return {};
  }

  AuthorizedEntity entity_enum;
  if (!AuthorizedEntity_Parse(*entity, &entity_enum)) {
    LOG(WARNING) << "Warning: entity " << *entity
                 << " is not a valid AuthorizedEntity";
    return {};
  };
  std::vector<AuthzConfiguration::SpiffeIdentityMatcher> matchers;

  for (const std::string& entity_tag :
       GetAllEntityTagsFromAuthorizedEntity(entity_enum)) {
    AuthzConfiguration::SpiffeIdentityMatcher matcher;
    absl::StatusOr<AuthzConfiguration::EntityTagIdentity> entity_tag_identity =
        GetEntityTagIdentity(entity_tag);
    if (!entity_tag_identity.ok()) {
      continue;
    }
    matcher.fqdn =
        absl::StrCat(entity_tag_identity->hostname, ".prod.google.com");

    matcher.spiffe_id_matcher.trust_domain_suffix = ".prod.google.com";
    matcher.spiffe_id_matcher.ca_context =
        CertificateMetadataParser::kCaContextNode;

    matcher.spiffe_id_matcher.path =
        absl::StrCat("/role/", GetCertNodeOwner(entity_tag_identity->owner));

    matchers.push_back(std::move(matcher));
  }

  return matchers;
}

void AuthzConfiguration::ReloadSampleRateLimit(const nlohmann::json& config) {
  if (config.is_discarded()) {
    return;
  }
  const nlohmann::json::array_t* role_sample_rate_limits =
      GetValueAsArray(config, "role_sample_rate_limits");
  if (role_sample_rate_limits == nullptr) {
    return;
  }
  std::unordered_map<std::string, uint64_t> new_role_to_sample_rate_limit;
  for (const nlohmann::json& role_sample_rate_limit :
       *role_sample_rate_limits) {
    const std::string* role =
        GetValueAsString(role_sample_rate_limit, "redfish_role");
    if (role == nullptr) {
      continue;
    }
    const uint64_t* sample_rate_limit =
        GetValueAsUint(role_sample_rate_limit, "sample_rate_limit");
    if (sample_rate_limit == nullptr) {
      continue;
    }
    new_role_to_sample_rate_limit[*role] = *sample_rate_limit;
  }
  absl::MutexLock lock(&mutex_);
  role_to_sample_rate_limit_ = std::move(new_role_to_sample_rate_limit);
}

void AuthzConfiguration::ReloadResourceOwners(const nlohmann::json& config) {
  if (config.is_discarded()) {
    return;
  }
  const nlohmann::json::array_t* resource_owners =
      GetValueAsArray(config, "resource_owners");
  if (resource_owners == nullptr) {
    return;
  }
  std::vector<ResourceOwner> new_resource_owners;
  for (const nlohmann::json& resource_owner : *resource_owners) {
    // If the spiffe matcher is explicitly defined, parse it.
    if (resource_owner.contains("spiffe_id_matcher")) {
      ResourceOwner owner;
      std::optional<SpiffeIdentityMatcher> matcher =
          ParseSpiffeIdentityMatcher(resource_owner);
      if (matcher) {
        owner.identity_matcher = *matcher;
        new_resource_owners.push_back(std::move(owner));
      }
      continue;
    }
    // Otherwise continue with entity based parsing
    std::vector<SpiffeIdentityMatcher> matchers =
        ParseSpiffeIdentityMatchersUsingOne(resource_owner);
    for (const SpiffeIdentityMatcher& matcher : matchers) {
      ResourceOwner owner;
      owner.identity_matcher = matcher;
      new_resource_owners.push_back(std::move(owner));
    }
  }

  std::sort(new_resource_owners.begin(), new_resource_owners.end());

  absl::MutexLock lock(&mutex_);
  resource_owners_ = std::move(new_resource_owners);
}

void AuthzConfiguration::ReloadRedfishUsers(const nlohmann::json& config) {
  if (config.is_discarded()) {
    return;
  }
  const nlohmann::json::array_t* exchange_mappings =
      GetValueAsArray(config, "exchange_mappings");
  if (exchange_mappings == nullptr) {
    return;
  }
  std::vector<SpiffeUser> new_spiffe_users;
  std::unordered_map<std::string, std::string> new_username_to_role;
  std::unordered_map<std::string, std::string> new_group_to_role;
  for (const nlohmann::json& exchange_mapping : *exchange_mappings) {
    const std::string* redfish_role =
        GetValueAsString(exchange_mapping, "redfish_role");
    if (redfish_role == nullptr) {
      continue;
    }
    // If peer is explicitly defined, then create exchange mapping with peer.
    if (exchange_mapping.contains("peer")) {
      const nlohmann::json* peer = GetValueAsJson(exchange_mapping, "peer");
      if (peer == nullptr) {
        continue;
      }
      // If it is username based, parse the username and continue
      const std::string* username = GetValueAsString(*peer, "username");
      if (username != nullptr) {
        new_username_to_role[*username] = *redfish_role;
        continue;
      }
      // If it is group based, parse the group and continue
      const std::string* group = GetValueAsString(*peer, "group");
      if (group != nullptr) {
        new_group_to_role[*group] = *redfish_role;
        continue;
      }
      // Otherwise, it is spiffe based
      SpiffeUser user;
      std::optional<SpiffeIdentityMatcher> matcher =
          ParseSpiffeIdentityMatcher(*peer);
      if (!matcher) {
        continue;
      }
      user.identity_matcher = *matcher;
      user.redfish_role = *redfish_role;
      new_spiffe_users.push_back(std::move(user));
      continue;
    }

    // Otherwise, continue with entity based parsing
    std::vector<SpiffeIdentityMatcher> matchers =
        ParseSpiffeIdentityMatchersUsingOne(exchange_mapping);
    for (const SpiffeIdentityMatcher& matcher : matchers) {
      SpiffeUser user;
      user.identity_matcher = matcher;
      user.redfish_role = *redfish_role;
      // TODO(edwarddl): This is n^2. We can make it linear with refactoring of
      // vector into hashmap. Do profiling and benchmarking to see if this would
      // lead to better performance.

      // Update the redfish_role of the matcher if
      //  the user already exists.
      auto it =
          std::find(new_spiffe_users.begin(), new_spiffe_users.end(), user);
      if (it != new_spiffe_users.end()) {
        it->redfish_role = *redfish_role;
        continue;
      }

      // Add the user to the list if they aren't there
      new_spiffe_users.push_back(std::move(user));
    }
  }

  std::sort(new_spiffe_users.begin(), new_spiffe_users.end());

  absl::MutexLock lock(&mutex_);
  spiffe_users_ = std::move(new_spiffe_users);
  username_to_role_ = std::move(new_username_to_role);
  group_to_role_ = std::move(new_group_to_role);
}

static void AddRoleToPrivilegesToMap(
    const nlohmann::json& role_privileges,
    std::unordered_map<std::string, std::unordered_set<std::string>>&
        role_to_privileges) {
  if (role_privileges.is_discarded()) {
    return;
  }

  for (const nlohmann::json& role_privilege : role_privileges) {
    const std::string* role = GetValueAsString(role_privilege, "redfish_role");
    if (role == nullptr) {
      continue;
    }
    const nlohmann::json* privileges =
        GetValueAsJson(role_privilege, "privileges");
    if (privileges == nullptr) {
      continue;
    }
    std::unordered_set<std::string> privileges_set;
    const nlohmann::json::array_t* base_privileges =
        GetValueAsArray(*privileges, "base_privileges");
    if (base_privileges != nullptr) {
      for (const nlohmann::json& privilege : *base_privileges) {
        const std::string* privilege_str =
            privilege.get_ptr<const std::string*>();
        if (privilege_str == nullptr) {
          continue;
        }
        privileges_set.insert(*privilege_str);
      }
    }
    const nlohmann::json::array_t* oem_privileges =
        GetValueAsArray(*privileges, "oem_privileges");
    if (oem_privileges != nullptr) {
      for (const nlohmann::json& privilege : *oem_privileges) {
        const std::string* privilege_str =
            privilege.get_ptr<const std::string*>();
        if (privilege_str == nullptr) {
          continue;
        }
        privileges_set.insert(*privilege_str);
      }
    }

    privileges_set.insert("NoAuth");
    role_to_privileges[*role] = std::move(privileges_set);
  }
}

void AuthzConfiguration::ReloadRoleToPrivileges(
    const nlohmann::json& config, const nlohmann::json& platform_config) {
  if (config.is_discarded()) {
    return;
  }
  const nlohmann::json::array_t* role_privileges =
      GetValueAsArray(config, "role_privileges_mappings");
  std::unordered_map<std::string, std::unordered_set<std::string>>
      new_role_to_privileges;
  if (role_privileges != nullptr) {
    AddRoleToPrivilegesToMap(*role_privileges, new_role_to_privileges);
  }

  const nlohmann::json::array_t* platform_role_privileges =
      GetValueAsArray(platform_config, "role_privileges_mappings");
  if (platform_role_privileges != nullptr) {
    AddRoleToPrivilegesToMap(*platform_role_privileges, new_role_to_privileges);
  }

  absl::MutexLock lock(&mutex_);
  role_to_privileges_ = std::move(new_role_to_privileges);
}

void AuthzConfiguration::ReloadOfflineNodeEntities() {
  if (offline_node_entities_path_.empty()) {
    return;
  }

  absl::StatusOr<OfflineNodeEntityInformation> parsed_offline_node_entities =
      ReadOfflineNodeEntityInformation(offline_node_entities_path_);

  if (!parsed_offline_node_entities.ok()) {
    LOG(WARNING) << "Could not read offline node entities information from "
                 << offline_node_entities_path_ << ": "
                 << parsed_offline_node_entities.status().message();
    return;
  }

  absl::MutexLock lock(&mutex_);
  offline_node_entities_ = std::move(*parsed_offline_node_entities);
}

void AuthzConfiguration::ReloadGmi() {
  if (gmi_path_.empty()) {
    return;
  }

  std::ifstream gmi_file(gmi_path_);
  if (!gmi_file.is_open()) {
    LOG(ERROR) << "GMI file at '" << gmi_path_ << "' is missing.";
    return;
  }
  GoogleMachineIdentityProto gmi;
  if (!gmi.ParseFromIstream(&gmi_file)) {
    LOG(ERROR) << "GMI parsing failed at '" << gmi_path_ << "'";
    return;
  }

  absl::MutexLock lock(&mutex_);
  gmi_ = std::move(gmi);
}

void AuthzConfiguration::ReloadConfig(const nlohmann::json& config,
                                      const nlohmann::json& platform_config) {
  if (config.is_discarded()) {
    LOG(WARNING) << "Warning: config isn't a valid JSON";
    return;
  }

  ReloadGmi();
  ReloadOfflineNodeEntities();
  ReloadResourceOwners(config);
  ReloadRedfishUsers(config);
  ReloadRoleToPrivileges(config, platform_config);
  ReloadSampleRateLimit(config);
}

std::unordered_set<std::string> AuthzConfiguration::GetOemPrivileges() const {
  std::unordered_set<std::string> privileges;

  absl::MutexLock lock(&mutex_);
  for (const auto& [role, privilege_set] : role_to_privileges_) {
    for (const std::string& privilege : privilege_set) {
      // If the privilege is not a default redfish privilege, it is OEM.
      if (!std::binary_search(kDefaultRedfishPrivileges.begin(),
                              kDefaultRedfishPrivileges.end(), privilege)) {
        privileges.insert(privilege);
      }
    }
  }

  return privileges;
}

absl::StatusOr<EntityTagIdentity> AuthzConfiguration::GetEntityTagIdentity(
    absl::string_view entity_tag) const {
  absl::MutexLock lock(&mutex_);
  auto it =
      offline_node_entities_.resolved_config().entities().find(entity_tag);
  if (it == offline_node_entities_.resolved_config().entities().end()) {
    return absl::NotFoundError("Entity Tag Not Found");
  }
  return EntityTagIdentity{.hostname = it->second.hostname(),
                           .owner = it->second.owner()};
}

std::string AuthzConfiguration::GetMachineManagerEntityTag() const {
  absl::MutexLock lock(&mutex_);
  return offline_node_entities_.resolved_config()
      .ecclesia_machine_manager_entity();
}

std::vector<std::string> AuthzConfiguration::GetAllComputeNodeEntityTags()
    const {
  std::vector<std::string> entity_tags;
  absl::MutexLock lock(&mutex_);
  for (const auto& [entity_tag, entity] :
       offline_node_entities_.resolved_config().entities()) {
    if (entity.node_type_info().is_compute()) {
      entity_tags.push_back(entity_tag);
    }
  }

  return entity_tags;
}

std::vector<std::string> AuthzConfiguration::GetEntityTagsFromRedfishSystemId(
    absl::string_view system_id) const {
  std::vector<std::string> entity_tags;
  absl::MutexLock lock(&mutex_);
  for (const auto& [entity_tag, entity] :
       offline_node_entities_.resolved_config().entities()) {
    if (entity.redfish_system_id() == system_id) {
      entity_tags.push_back(entity_tag);
    }
  }

  return entity_tags;
}

std::vector<std::string>
AuthzConfiguration::GetAllEntityTagsFromAuthorizedEntity(
    AuthorizedEntity entity) const {
  switch (entity) {
    case AuthorizedEntity::ComputeNode:
      return GetAllComputeNodeEntityTags();
    case AuthorizedEntity::ComputeNode1:
      return GetEntityTagsFromRedfishSystemId("system1");
    case AuthorizedEntity::ComputeNode2:
      return GetEntityTagsFromRedfishSystemId("system2");
    case AuthorizedEntity::EcclesiaMachineManagerEntity:
      return {GetMachineManagerEntityTag()};
    case AuthorizedEntity::Undefined:
      return {};
    case AuthorizedEntity_INT_MIN_SENTINEL_DO_NOT_USE_:
    case AuthorizedEntity_INT_MAX_SENTINEL_DO_NOT_USE_:
      return {};
  }
  return {};
}

bool AuthzConfiguration::IsHwopsState() const {
  absl::MutexLock lock(&mutex_);
  if (gmi_.has_in_hwops_state()) {
    return gmi_.in_hwops_state();
  }
  return false;
}

std::string AuthzConfiguration::GetNodeRealm() const {
  absl::MutexLock lock(&mutex_);
  // GMI only supports 1 realm
  if (gmi_.security_realms_size() != 1) {
    LOG(ERROR) << "Machine GMI's realm is broken. All realm specific Borg jobs "
                  "will be denied.";
    return "";
  }
  return gmi_.security_realms(0);
}

}  // namespace milotic::authz