//
// Copyright 2021 gRPC authors.
//
// 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
//
//     http://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.
//

#include <grpc/support/json.h>
#include <grpc/support/port_platform.h>

#include <cstdint>
#include <memory>
#include <optional>
#include <random>
#include <string>
#include <type_traits>
#include <utility>

#include "src/core/config/core_configuration.h"
#include "src/core/credentials/transport/alts/check_gcp_environment.h"
#include "src/core/lib/channel/channel_args.h"
#include "src/core/lib/iomgr/polling_entity.h"
#include "src/core/lib/resource_quota/resource_quota.h"
#include "src/core/resolver/resolver.h"
#include "src/core/resolver/resolver_factory.h"
#include "src/core/resolver/resolver_registry.h"
#include "src/core/util/debug_location.h"
#include "src/core/util/env.h"
#include "src/core/util/gcp_metadata_query.h"
#include "src/core/util/grpc_check.h"
#include "src/core/util/json/json.h"
#include "src/core/util/json/json_writer.h"
#include "src/core/util/orphanable.h"
#include "src/core/util/ref_counted_ptr.h"
#include "src/core/util/time.h"
#include "src/core/util/uri.h"
#include "src/core/util/work_serializer.h"
#include "src/core/xds/grpc/xds_client_grpc.h"
#include "src/core/xds/xds_client/xds_bootstrap.h"
#include "absl/log/log.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/strip.h"

namespace grpc_core {

namespace {

const char* kC2PAuthority = "traffic-director-c2p.xds.googleapis.com";

class GoogleCloud2ProdResolver final : public Resolver {
 public:
  explicit GoogleCloud2ProdResolver(ResolverArgs args);

  void StartLocked() override;
  void RequestReresolutionLocked() override;
  void ResetBackoffLocked() override;
  void ShutdownLocked() override;

 private:
  void ZoneQueryDone(std::string zone);
  void IPv6QueryDone(bool ipv6_supported);
  void StartXdsResolver();
  std::shared_ptr<GrpcXdsBootstrap> ConstructBootstrap() const;

  ResourceQuotaRefPtr resource_quota_;
  std::shared_ptr<WorkSerializer> work_serializer_;
  grpc_polling_entity pollent_;
  bool using_dns_ = false;
  OrphanablePtr<Resolver> child_resolver_;
  std::string metadata_server_name_ = "metadata.google.internal.";
  bool shutdown_ = false;

  // Used to hold state between ctor and StartLocked().
  std::string xds_uri_;
  ChannelArgs args_;

  OrphanablePtr<GcpMetadataQuery> zone_query_;
  std::string zone_;
};

//
// GoogleCloud2ProdResolver
//

bool XdsBootstrapConfigured() {
  return GetEnv("GRPC_XDS_BOOTSTRAP").has_value() ||
         GetEnv("GRPC_XDS_BOOTSTRAP_CONFIG").has_value();
}

GoogleCloud2ProdResolver::GoogleCloud2ProdResolver(ResolverArgs args)
    : resource_quota_(args.args.GetObjectRef<ResourceQuota>()),
      work_serializer_(std::move(args.work_serializer)),
      pollent_(grpc_polling_entity_create_from_pollset_set(args.pollset_set)),
      args_(std::move(args.args)) {
  absl::string_view name_to_resolve = absl::StripPrefix(args.uri.path(), "/");
  // If we're not running on GCP, we can't use DirectPath, so delegate
  // to the DNS resolver.
  const bool test_only_pretend_running_on_gcp =
      args_.GetBool("grpc.testing.google_c2p_resolver_pretend_running_on_gcp")
          .value_or(false);
  const bool running_on_gcp =
      test_only_pretend_running_on_gcp || grpc_alts_is_running_on_gcp();
  const bool federation_enabled = XdsFederationEnabled();
  if (!running_on_gcp ||
      // If the client is already using xDS and federation is not enabled,
      // we can't use it here, because they may be talking to a completely
      // different xDS server than we want to.
      // TODO(roth): When we remove xDS federation env var protection,
      // remove this constraint.
      (!federation_enabled && XdsBootstrapConfigured())) {
    using_dns_ = true;
    child_resolver_ =
        CoreConfiguration::Get().resolver_registry().CreateResolver(
            absl::StrCat("dns:", name_to_resolve), args_, args.pollset_set,
            work_serializer_, std::move(args.result_handler));
    GRPC_CHECK(child_resolver_ != nullptr);
    return;
  }
  // Maybe override metadata server name for testing
  std::optional<std::string> test_only_metadata_server_override =
      args_.GetOwnedString(
          "grpc.testing.google_c2p_resolver_metadata_server_override");
  if (test_only_metadata_server_override.has_value() &&
      !test_only_metadata_server_override->empty()) {
    metadata_server_name_ = std::move(*test_only_metadata_server_override);
  }
  // Create xds resolver.
  xds_uri_ = federation_enabled
                 ? absl::StrCat("xds://", kC2PAuthority, "/", name_to_resolve)
                 : absl::StrCat("xds:", name_to_resolve);
  child_resolver_ = CoreConfiguration::Get().resolver_registry().CreateResolver(
      xds_uri_, args_, args.pollset_set, work_serializer_,
      std::move(args.result_handler));
  GRPC_CHECK(child_resolver_ != nullptr);
}

void GoogleCloud2ProdResolver::StartLocked() {
  if (using_dns_) {
    child_resolver_->StartLocked();
    return;
  }
  // Using xDS.  Start metadata server queries.
  zone_query_ = MakeOrphanable<GcpMetadataQuery>(
      metadata_server_name_, std::string(GcpMetadataQuery::kZoneAttribute),
      &pollent_,
      [resolver = RefAsSubclass<GoogleCloud2ProdResolver>()](
          std::string /* attribute */,
          absl::StatusOr<std::string> result) mutable {
        resolver->work_serializer_->Run([resolver,
                                         result = std::move(result)]() mutable {
          resolver->ZoneQueryDone(result.ok() ? std::move(result).value() : "");
        });
      },
      Duration::Seconds(10));
}

void GoogleCloud2ProdResolver::RequestReresolutionLocked() {
  if (child_resolver_ != nullptr) {
    child_resolver_->RequestReresolutionLocked();
  }
}

void GoogleCloud2ProdResolver::ResetBackoffLocked() {
  if (child_resolver_ != nullptr) {
    child_resolver_->ResetBackoffLocked();
  }
}

void GoogleCloud2ProdResolver::ShutdownLocked() {
  shutdown_ = true;
  zone_query_.reset();
  child_resolver_.reset();
}

void GoogleCloud2ProdResolver::ZoneQueryDone(std::string zone) {
  zone_query_.reset();
  zone_ = std::move(zone);
  StartXdsResolver();
}

std::shared_ptr<GrpcXdsBootstrap> GoogleCloud2ProdResolver::ConstructBootstrap()
    const {
  // Construct bootstrap JSON.
  std::random_device rd;
  std::mt19937 mt(rd());
  std::uniform_int_distribution<uint64_t> dist(1, UINT64_MAX);
  Json::Object node = {
      {"id", Json::FromString(absl::StrCat("C2P-", dist(mt)))},
  };
  if (!zone_.empty()) {
    node["locality"] = Json::FromObject({
        {"zone", Json::FromString(zone_)},
    });
  };
  // Enable dualstack endpoints from TD
  node["metadata"] = Json::FromObject({
      {"TRAFFICDIRECTOR_DIRECTPATH_C2P_IPV6_CAPABLE", Json::FromBool(true)},
  });
  // Allow the TD server uri to be overridden for testing purposes.
  auto override_server =
      GetEnv("GRPC_TEST_ONLY_GOOGLE_C2P_RESOLVER_TRAFFIC_DIRECTOR_URI");
  const char* server_uri =
      override_server.has_value() && !override_server->empty()
          ? override_server->c_str()
          : "directpath-pa.googleapis.com";
  Json xds_server = Json::FromArray({
      Json::FromObject({
          {"server_uri", Json::FromString(server_uri)},
          {"channel_creds",
           Json::FromArray({
               Json::FromObject({
                   {"type", Json::FromString("google_default")},
               }),
           })},
          {"server_features",
           Json::FromArray({Json::FromString("ignore_resource_deletion")})},
      }),
  });
  Json bootstrap = Json::FromObject({
      {"xds_servers", xds_server},
      {"authorities",
       Json::FromObject({
           {kC2PAuthority, Json::FromObject({
                               {"xds_servers", std::move(xds_server)},
                           })},
       })},
      {"node", Json::FromObject(std::move(node))},
  });
  return GrpcXdsBootstrap::Create(JsonDump(bootstrap)).value();
}

void GoogleCloud2ProdResolver::StartXdsResolver() {
  if (shutdown_) return;
  // Create XdsClient here with a custom bootstrap for C2P.  When the
  // xds resolver calls GrpcXdsClient::GetOrCreate(), it will get back
  // the same instance rather than creating one with the normal
  // bootstrap config from env vars.
  static NoDestruct<std::shared_ptr<GrpcXdsBootstrap>> bootstrap(
      ConstructBootstrap());
  auto xds_client = GrpcXdsClient::GetOrCreate(
      xds_uri_, args_, "google-c2p resolver", *bootstrap);
  args_ = ChannelArgs();  // Don't need the args anymore.
  // Now start xDS resolver.
  child_resolver_->StartLocked();
}

//
// Factory
//

class GoogleCloud2ProdResolverFactory final : public ResolverFactory {
 public:
  absl::string_view scheme() const override { return "google-c2p"; }

  bool IsValidUri(const URI& uri) const override {
    if (GPR_UNLIKELY(!uri.authority().empty())) {
      LOG(ERROR) << "google-c2p URI scheme does not support authorities";
      return false;
    }
    return true;
  }

  OrphanablePtr<Resolver> CreateResolver(ResolverArgs args) const override {
    if (!IsValidUri(args.uri)) return nullptr;
    return MakeOrphanable<GoogleCloud2ProdResolver>(std::move(args));
  }
};

// TODO(apolcyn): remove this class after user code has updated to the
// stable "google-c2p" URI scheme.
class ExperimentalGoogleCloud2ProdResolverFactory final
    : public ResolverFactory {
 public:
  absl::string_view scheme() const override {
    return "google-c2p-experimental";
  }

  bool IsValidUri(const URI& uri) const override {
    if (GPR_UNLIKELY(!uri.authority().empty())) {
      LOG(ERROR) << "google-c2p-experimental URI scheme does not support "
                    "authorities";
      return false;
    }
    return true;
  }

  OrphanablePtr<Resolver> CreateResolver(ResolverArgs args) const override {
    if (!IsValidUri(args.uri)) return nullptr;
    return MakeOrphanable<GoogleCloud2ProdResolver>(std::move(args));
  }
};

}  // namespace

void RegisterCloud2ProdResolver(CoreConfiguration::Builder* builder) {
  builder->resolver_registry()->RegisterResolverFactory(
      std::make_unique<GoogleCloud2ProdResolverFactory>());
  builder->resolver_registry()->RegisterResolverFactory(
      std::make_unique<ExperimentalGoogleCloud2ProdResolverFactory>());
}

}  // namespace grpc_core
