#pragma once

#include <string>

#include <utility>

#include <vector>

#include "envoy/common/time.h"

#include "envoy/tracing/tracer.h"

#include "source/common/common/empty_string.h"

#include "source/common/common/hex.h"

#include "source/common/common/random_generator.h"

#include "source/common/http/codes.h"

#include "source/common/protobuf/utility.h"

#include "source/common/tracing/common_values.h"

#include "source/common/tracing/trace_context_impl.h"

#include "source/extensions/tracers/xray/daemon_broker.h"

#include "source/extensions/tracers/xray/sampling_strategy.h"

#include "source/extensions/tracers/xray/xray_configuration.h"

#include "absl/container/flat_hash_map.h"

#include "absl/strings/string_view.h"

namespace Envoy {

namespace Extensions {

namespace Tracers {

namespace XRay {

constexpr absl::string_view SpanClientIp = "client_ip";

constexpr absl::string_view SpanXForwardedFor = "x_forwarded_for";

constexpr absl::string_view Subsegment = "subsegment";

const Tracing::TraceContextHandler& xRayTraceHeader();

const Tracing::TraceContextHandler& xForwardedForHeader();

class Span : public Tracing::Span, Logger::Loggable<Logger::Id::config> {

public:

  /**

   * Creates a new Span.

   *

   * @param time_source A time source to get the span end time

   * @param random random generator for generating unique child span ids

   * @param broker Facilitates communication with the X-Ray daemon.

   */

  Span(TimeSource& time_source, Random::RandomGenerator& random, DaemonBroker& broker)

  /**

   * Sets the Span's trace ID.

   */

  /**

   * Gets the Span's trace ID.

   */

  /**

   * Completes the current span, serialize it and send it to the X-Ray daemon.

   */

  void finishSpan() override;

  /**

   * Sets the current operation name on the Span.

   * This information will be included in the X-Ray span's metadata.

   */

  /**

   * Sets the current direction on the Span.

   * This information will be included in the X-Ray span's annotation.

   */

  /**

   * Sets the name of the Span.

   */

  /**

   * Sets the origin of the Span.

   */

  /**

   * Gets the origin of the Span.

   */

  /**

   * Adds a key-value pair to either the Span's annotations or metadata.

   * An allowlist of keys are added to the annotations, everything else is added to the metadata.

   */

  void setTag(absl::string_view name, absl::string_view value) override;

  /**

   * Sets the ID of the parent segment. This is different from the Trace ID.

   * The parent ID is used if the request originated from an instrumented application.

   * For more information see:

   * https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-tracingheader

   */

  /**

   * Sets the type of the Span. In X-Ray, an independent subsegment has a type of "subsegment".

   * https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html#api-segmentdocuments-subsegments

   */

  /**

   * Sets the aws metadata field of the Span.

   */

  /*

   * Adds to the http request annotation field of the Span.

   */

  /*

   * Check if key is set in http request annotation field of a Span.

   */

  /*

   * Adds to the http response annotation field of the Span.

   */

  /**

   * Sets the recording start time of the traced operation/request.

   */

  /**

   * Marks the span as either "sampled" or "not-sampled".

   * By default, Spans are "sampled".

   * This is handy in cases where the sampling decision has already been determined either by Envoy

   * or by a downstream service.

   */

  // X-Ray tracer does not use the sampling decision from Envoy anyway.

  /**

   * Sets the server error as true for the traced operation/request.

   */

  /**

   * Sets the http response status code for the traced operation/request.

   */

  /**

   * Adds X-Ray trace header to the set of outgoing headers.

   */

  void injectContext(Tracing::TraceContext& trace_context,

                     const Tracing::UpstreamContext&) override;

  /**

   * Gets the start time of this Span.

   */

  /**

   * Gets this Span's ID.

   */

  /**

   * Gets this Span's parent ID.

   */

  /**

   * Gets this Span's direction.

   */

  /**

   * Gets this Span's type.

   */

  /**

   * Gets this Span's name.

   */

  /**

   * Determines whether this span is sampled.

   */

  /**

   * Determines if a server error occurred (response status code was 5XX Server Error).

   */

  /**

   * Determines if a client error occurred (response status code was 4XX Client Error).

   */

  /**

   * Determines if a request was throttled (response status code was 429 Too Many Requests).

   */

  /**

   * Not used by X-Ray because the Spans are "logged" (serialized) to the X-Ray daemon.

   */

  // X-Ray doesn't support baggage, so noop these OpenTracing functions.

  // TODO(#34412): This method is unimplemented for X-Ray.

  /**

   * Creates a child span.

   * In X-Ray terms this creates a sub-segment and sets its parent ID to the current span's ID.

   * @param operation_name The span of the child span.

   * @param start_time The time at which this child span started.

   */

  Tracing::SpanPtr spawnChild(const Tracing::Config&, const std::string& operation_name,

                              Envoy::SystemTime start_time) override;

private:

  Envoy::TimeSource& time_source_;

  Random::RandomGenerator& random_;

  DaemonBroker& broker_;

  Envoy::SystemTime start_time_;

  std::string operation_name_;

  std::string direction_;

  std::string id_;

  std::string trace_id_;

  std::string parent_segment_id_;

  std::string name_;

  std::string origin_;

  std::string type_;

  absl::flat_hash_map<std::string, Protobuf::Value> aws_metadata_;

  absl::flat_hash_map<std::string, Protobuf::Value> http_request_annotations_;

  absl::flat_hash_map<std::string, Protobuf::Value> http_response_annotations_;

  absl::flat_hash_map<std::string, std::string> custom_annotations_;

  bool server_error_{false};

  uint64_t response_status_code_{0};

  bool sampled_{true};

};

using SpanPtr = std::unique_ptr<Span>;

class Tracer {

public:

  Tracer(absl::string_view segment_name, absl::string_view origin,

         const absl::flat_hash_map<std::string, Protobuf::Value>& aws_metadata,

         DaemonBrokerPtr daemon_broker, TimeSource& time_source, Random::RandomGenerator& random)

  /**

   * Starts a tracing span for X-Ray

   */

  Tracing::SpanPtr startSpan(const Tracing::Config&, const std::string& operation_name,

                             Envoy::SystemTime start_time,

                             const absl::optional<XRayHeader>& xray_header,

                             const absl::optional<absl::string_view> client_ip);

  /**

   * Creates a Span that is marked as not-sampled.

   * This is useful when the sampling decision is done in Envoy's X-Ray and we want to avoid

   * overruling that decision in the upstream service in case that service itself uses X-Ray for

   * tracing. Also at the same time if X-Ray header is set then preserve its value.

   */

  XRay::SpanPtr createNonSampledSpan(const absl::optional<XRayHeader>& xray_header) const;

private:

  const std::string segment_name_;

  const std::string origin_;

  const absl::flat_hash_map<std::string, Protobuf::Value> aws_metadata_;

  const DaemonBrokerPtr daemon_broker_;

  Envoy::TimeSource& time_source_;

  Random::RandomGenerator& random_;

};

using TracerPtr = std::unique_ptr<Tracer>;

} // namespace XRay

} // namespace Tracers

} // namespace Extensions

} // namespace Envoy