// Copyright (c) 2017-2022 Cloudflare, Inc.
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
//     https://opensource.org/licenses/Apache-2.0

#pragma once

#include <workerd/io/outcome.capnp.h>
#include <workerd/io/trace.capnp.h>
#include <workerd/io/worker-interface.capnp.h>
#include <workerd/jsg/memory.h>
#include <workerd/util/own-util.h>

#include <kj/map.h>
#include <kj/one-of.h>
#include <kj/refcount.h>
#include <kj/string.h>
#include <kj/time.h>
#include <kj/vector.h>

#include <concepts>
#include <initializer_list>

namespace kj {
enum class HttpMethod;
class EntropySource;
}  // namespace kj

namespace workerd {

using kj::byte;
using kj::uint;

using LogLevel = rpc::Trace::Log::Level;
using ExecutionModel = rpc::Trace::ExecutionModel;

class Trace;

namespace tracing {
// A 128-bit globally unique trace identifier. This will be used for both
// external and internal tracing. Specifically, for internal tracing, this
// is used to represent tracing IDs for jaeger traces. For external tracing,
// this is used for both the trace ID and invocation ID for tail workers.
class TraceId final {
 public:
  // A null trace ID. This is only acceptable for use in tests.
  constexpr TraceId(decltype(nullptr)) {}

  // A trace ID with the given low and high values.
  constexpr TraceId(uint64_t low, uint64_t high): low(low), high(high) {}

  constexpr TraceId(const TraceId& other) = default;
  constexpr TraceId& operator=(const TraceId& other) = default;

  constexpr TraceId(TraceId&& other): low(other.low), high(other.high) {
    other.low = 0;
    other.high = 0;
  }

  constexpr TraceId& operator=(TraceId&& other) {
    low = other.low;
    high = other.high;
    other.low = 0;
    other.high = 0;
    return *this;
  }

  constexpr TraceId& operator=(decltype(nullptr)) {
    low = 0;
    high = 0;
    return *this;
  }

  constexpr bool operator==(const TraceId& other) const {
    return low == other.low && high == other.high;
  }
  constexpr bool operator==(decltype(nullptr)) const {
    return low == 0 && high == 0;
  }
  constexpr operator bool() const {
    return low || high;
  }

  operator kj::String() const {
    return toGoString();
  }

  // Replicates Jaeger go library's string serialization.
  kj::String toGoString() const;

  // Replicates Jaeger go library's protobuf serialization.
  kj::Array<byte> toProtobuf() const;

  // Replicates W3C Serialization
  kj::String toW3C() const;

  // Creates a random Trace Id, optionally using a given entropy source. If an
  // entropy source is not given, then we fallback to using BoringSSL's RAND_bytes.
  static TraceId fromEntropy(kj::Maybe<kj::EntropySource&> entropy = kj::none);

  // Replicates Jaeger go library's string serialization.
  static kj::Maybe<TraceId> fromGoString(kj::ArrayPtr<const char> s);

  // Replicates Jaeger go library's protobuf serialization.
  static kj::Maybe<TraceId> fromProtobuf(kj::ArrayPtr<const kj::byte> buf);

  // A null trace ID. This is really only acceptable for use in tests.
  static const TraceId nullId;

  inline uint64_t getLow() const {
    return low;
  }
  inline uint64_t getHigh() const {
    return high;
  }

  static TraceId fromCapnp(rpc::TraceId::Reader reader);
  void toCapnp(rpc::TraceId::Builder writer) const;

 private:
  uint64_t low = 0;
  uint64_t high = 0;
};
constexpr TraceId TraceId::nullId = nullptr;

// A 64-bit span identifier.
class SpanId final {
 public:
  // A null span ID. This is only acceptable for use in tests.
  constexpr SpanId(decltype(nullptr)): id(0) {}

  constexpr SpanId(uint64_t id): id(id) {}
  constexpr SpanId(const SpanId& other) = default;
  constexpr SpanId& operator=(const SpanId& other) = default;
  constexpr SpanId(SpanId&& other): id(other.id) {
    other.id = 0;
  }
  constexpr SpanId& operator=(SpanId&& other) {
    id = other.id;
    other.id = 0;
    return *this;
  }
  constexpr operator bool() const {
    return id != 0;
  }
  constexpr bool operator==(const SpanId& other) const {
    return id == other.id;
  }
  constexpr bool operator==(decltype(nullptr)) const {
    return id == 0;
  }

  inline operator kj::String() const {
    return toGoString();
  }

  inline operator uint64_t() const {
    return id;
  }

  kj::String toGoString() const;

  static const SpanId nullId;

  constexpr uint64_t getId() const {
    return id;
  }

  static SpanId fromEntropy(kj::Maybe<kj::EntropySource&> entropy = kj::none);

 private:
  uint64_t id;
};
constexpr SpanId SpanId::nullId = nullptr;
// Fixed spanId value to be used for tests
constexpr uint64_t staticSpanId = 0x2a2a2a2a2a2a2a2aULL;

// W3C trace flags propagated from an upstream traceparent. Wrapped in kj::Maybe
// at usage sites: kj::none means no upstream decision was made.
class TraceFlags final {
 public:
  explicit constexpr TraceFlags(uint8_t flags): flags(flags) {}

  constexpr bool isSampled() const {
    return flags & SAMPLED;
  }

  constexpr operator uint8_t() const {
    return flags;
  }

  constexpr bool operator==(const TraceFlags& other) const {
    return flags == other.flags;
  }

 private:
  // W3C trace-flags bit: the caller requested this trace be recorded.
  static constexpr uint8_t SAMPLED = 0x01;

  uint8_t flags;
};

// The InvocationSpanContext is a tuple of a trace id, invocation id, and span id.
// The trace id represents a top-level request and should be shared across all
// invocation spans and events within those spans. The invocation id identifies
// a specific worker invocation. The span id identifies a specific span within an
// invocation. Every invocation of every worker should have an InvocationSpanContext.
// That may or may not have a trigger InvocationSpanContext.
class InvocationSpanContext final {
 public:
  // The constructor is public only so kj::rc can see it and create a new instance.
  // User code should use the static factory methods or the newChild method.
  InvocationSpanContext(kj::Badge<InvocationSpanContext>,
      kj::Maybe<kj::EntropySource&> entropySource,
      TraceId traceId,
      TraceId invocationId,
      SpanId spanId,
      kj::Maybe<const InvocationSpanContext&> parentSpanContext,
      kj::Maybe<TraceFlags> traceFlags);
  // Still need a constructor to be available as long as span context is not propagated everywhere
  // we need it.
  InvocationSpanContext(
      TraceId traceId, TraceId invocationId, SpanId spanId, kj::Maybe<TraceFlags> traceFlags)
      : traceId(traceId),
        invocationId(invocationId),
        spanId(spanId),
        traceFlags(traceFlags) {};

  KJ_DISALLOW_COPY(InvocationSpanContext);

  InvocationSpanContext(InvocationSpanContext&& other) = default;
  InvocationSpanContext& operator=(InvocationSpanContext&& other) = default;

  inline bool operator==(const InvocationSpanContext& other) const {
    return traceId == other.traceId && invocationId == other.invocationId && spanId == other.spanId;
  }

  inline const TraceId& getTraceId() const {
    return traceId;
  }

  inline const TraceId& getInvocationId() const {
    return invocationId;
  }

  inline const SpanId& getSpanId() const {
    return spanId;
  }

  inline kj::Maybe<const InvocationSpanContext&> getParent() const {
    KJ_IF_SOME(p, parentSpanContext) {
      return *p;
    }
    return kj::none;
  }

  // W3C trace flags propagated from an upstream traceparent. kj::none when
  // no upstream sampling decision exists.
  inline kj::Maybe<TraceFlags> getTraceFlags() const {
    return traceFlags;
  }

  // Creates a new child span. If the current context does not have an entropy
  // source this will assert. If isTrigger() is true then it will not have an
  // entropy source.
  InvocationSpanContext newChild() const;

  // An InvocationSpanContext is a trigger context if it has no entropy source.
  // This generally means the SpanContext was create from a capnp message and
  // represents an InvocationSpanContext that was propagated from a parent
  // or triggering context.
  bool isTrigger() const {
    return entropySource == kj::none;
  }

  // Creates a new InvocationSpanContext. If the triggerContext is given, then its
  // traceId is used as the traceId for the newly created context. Otherwise a new
  // traceId is generated. The invocationId is always generated new and the spanId
  // will be 0 with no parent span.
  static InvocationSpanContext newForInvocation(
      kj::Maybe<const InvocationSpanContext&> triggerContext = kj::none,
      kj::Maybe<kj::EntropySource&> entropySource = kj::none);

  // Creates a new InvocationSpanContext from a capnp message. The returned
  // InvocationSpanContext will not be capable of creating child spans and
  // is considered only a "trigger" span.
  static kj::Maybe<InvocationSpanContext> fromCapnp(rpc::InvocationSpanContext::Reader reader);
  void toCapnp(rpc::InvocationSpanContext::Builder writer) const;
  InvocationSpanContext clone() const;

 private:
  // If there is no entropy source, then child spans cannot be created from
  // this InvocationSpanContext.
  kj::Maybe<kj::EntropySource&> entropySource;
  TraceId traceId;
  TraceId invocationId;
  SpanId spanId;

  // The parentSpanContext can be either a direct parent or a trigger
  // context. If it is a trigger context, then it should have the same
  // traceId but a different invocationId (unless predictable mode for
  // testing is enabled). The isTrigger() should also return true.
  kj::Maybe<kj::Own<InvocationSpanContext>> parentSpanContext;

  // W3C trace flags from an upstream traceparent, propagated through the
  // invocation chain. kj::none when no upstream sampling decision was made.
  kj::Maybe<TraceFlags> traceFlags;
};

// SpanContext as used for streaming tail worker tail events. spanId is always set except for Onset
// events that don't inherit context from another invocation.
struct SpanContext {
  SpanContext(
      TraceId traceId, kj::Maybe<SpanId> spanId, kj::Maybe<TraceFlags> traceFlags = kj::none)
      : traceId(traceId),
        spanId(spanId),
        traceFlags(traceFlags) {};
  KJ_DISALLOW_COPY(SpanContext);
  SpanContext(SpanContext&& other) = default;
  SpanContext& operator=(SpanContext&& other) = default;

  inline bool operator==(const SpanContext& other) const {
    return traceId == other.traceId && spanId == other.spanId;
  }

  inline const TraceId& getTraceId() const {
    return traceId;
  }

  inline kj::Maybe<SpanId> getSpanId() const {
    return spanId;
  }

  // W3C trace flags from an upstream traceparent. kj::none when no upstream
  // sampling decision was made (e.g. subrequest propagation without a
  // worker_tracing header).
  inline kj::Maybe<TraceFlags> getTraceFlags() const {
    return traceFlags;
  }

  static SpanContext fromCapnp(rpc::SpanContext::Reader reader);
  void toCapnp(rpc::SpanContext::Builder writer) const;
  static SpanContext clone(const SpanContext& ctx) {
    return SpanContext(ctx.traceId, ctx.spanId, ctx.traceFlags);
  }

  // Parse a W3C traceparent string into a SpanContext.
  // Format: "{version}-{trace-id}-{parent-id}-{flags}"
  static kj::Maybe<SpanContext> tryFromTraceparent(kj::StringPtr traceparent);

 private:
  TraceId traceId;
  kj::Maybe<SpanId> spanId;
  kj::Maybe<TraceFlags> traceFlags;
};

kj::String KJ_STRINGIFY(const SpanId& id);
kj::String KJ_STRINGIFY(const TraceId& id);
kj::String KJ_STRINGIFY(const InvocationSpanContext& context);
kj::String KJ_STRINGIFY(const SpanContext& context);

// The various structs defined below are used in both buffered tail workers
// and streaming tail workers to report tail events.

// Describes a fetch request
struct FetchEventInfo final {
  struct Header;

  explicit FetchEventInfo(
      kj::HttpMethod method, kj::String url, kj::String cfJson, kj::Array<Header> headers);
  FetchEventInfo(rpc::Trace::FetchEventInfo::Reader reader);
  FetchEventInfo(FetchEventInfo&&) noexcept = default;
  FetchEventInfo& operator=(FetchEventInfo&&) = default;
  KJ_DISALLOW_COPY(FetchEventInfo);

  struct Header final {
    explicit Header(kj::String name, kj::String value);
    Header(rpc::Trace::FetchEventInfo::Header::Reader reader);
    Header(Header&&) noexcept = default;
    Header& operator=(Header&&) = default;
    KJ_DISALLOW_COPY(Header);

    kj::String name;
    kj::String value;

    void copyTo(rpc::Trace::FetchEventInfo::Header::Builder builder) const;
    Header clone() const;
    kj::String toString() const;

    JSG_MEMORY_INFO(Header) {
      tracker.trackField("name", name);
      tracker.trackField("value", value);
    }
  };

  kj::HttpMethod method;
  kj::String url;
  // TODO(perf): It might be more efficient to store some sort of parsed JSON result instead?
  kj::String cfJson;
  kj::Array<Header> headers;

  void copyTo(rpc::Trace::FetchEventInfo::Builder builder) const;
  FetchEventInfo clone() const;
  kj::String toString() const;
};

// Describes a jsrpc request
struct JsRpcEventInfo final {
  explicit JsRpcEventInfo(kj::String methodName);
  JsRpcEventInfo(rpc::Trace::JsRpcEventInfo::Reader reader);
  JsRpcEventInfo(JsRpcEventInfo&&) noexcept = default;
  JsRpcEventInfo& operator=(JsRpcEventInfo&&) = default;
  KJ_DISALLOW_COPY(JsRpcEventInfo);

  kj::String methodName;

  void copyTo(rpc::Trace::JsRpcEventInfo::Builder builder) const;
  JsRpcEventInfo clone() const;
  kj::String toString() const;
};

class ConnectEventInfo {
 public:
  explicit ConnectEventInfo();
  explicit ConnectEventInfo(rpc::Trace::ConnectEventInfo::Reader reader);

  void copyTo(rpc::Trace::ConnectEventInfo::Builder builder) const;
  ConnectEventInfo clone() const;
};

// Describes a scheduled request
struct ScheduledEventInfo final {
  explicit ScheduledEventInfo(double scheduledTime, kj::String cron);
  ScheduledEventInfo(rpc::Trace::ScheduledEventInfo::Reader reader);
  ScheduledEventInfo(ScheduledEventInfo&&) noexcept = default;
  ScheduledEventInfo& operator=(ScheduledEventInfo&&) = default;
  KJ_DISALLOW_COPY(ScheduledEventInfo);

  double scheduledTime;
  kj::String cron;

  void copyTo(rpc::Trace::ScheduledEventInfo::Builder builder) const;
  ScheduledEventInfo clone() const;
};

// Describes a Durable Object alarm request
struct AlarmEventInfo final {
  explicit AlarmEventInfo(kj::Date scheduledTime);
  AlarmEventInfo(rpc::Trace::AlarmEventInfo::Reader reader);
  AlarmEventInfo(AlarmEventInfo&&) noexcept = default;
  AlarmEventInfo& operator=(AlarmEventInfo&&) = default;
  KJ_DISALLOW_COPY(AlarmEventInfo);

  kj::Date scheduledTime;

  void copyTo(rpc::Trace::AlarmEventInfo::Builder builder) const;
  AlarmEventInfo clone() const;
};

// Describes a queue worker request
struct QueueEventInfo final {
  explicit QueueEventInfo(kj::String queueName, uint32_t batchSize);
  QueueEventInfo(rpc::Trace::QueueEventInfo::Reader reader);
  QueueEventInfo(QueueEventInfo&&) noexcept = default;
  QueueEventInfo& operator=(QueueEventInfo&&) = default;
  KJ_DISALLOW_COPY(QueueEventInfo);

  kj::String queueName;
  uint32_t batchSize;

  void copyTo(rpc::Trace::QueueEventInfo::Builder builder) const;
  QueueEventInfo clone() const;
};

// Describes an email request
struct EmailEventInfo final {
  explicit EmailEventInfo(kj::String mailFrom, kj::String rcptTo, uint32_t rawSize);
  EmailEventInfo(rpc::Trace::EmailEventInfo::Reader reader);
  EmailEventInfo(EmailEventInfo&&) noexcept = default;
  EmailEventInfo& operator=(EmailEventInfo&&) = default;
  KJ_DISALLOW_COPY(EmailEventInfo);

  kj::String mailFrom;
  kj::String rcptTo;
  uint32_t rawSize;

  void copyTo(rpc::Trace::EmailEventInfo::Builder builder) const;
  EmailEventInfo clone() const;
};

// Describes a buffered tail worker request
struct TracePreview final {
  explicit TracePreview(kj::String id, kj::String slug, kj::String name);
  TracePreview(rpc::Trace::TracePreviewInfo::Reader reader);
  TracePreview(TracePreview&&) noexcept = default;
  TracePreview& operator=(TracePreview&&) = default;
  KJ_DISALLOW_COPY(TracePreview);

  kj::String id;
  kj::String slug;
  kj::String name;

  void copyTo(rpc::Trace::TracePreviewInfo::Builder builder) const;
  TracePreview clone() const;
};

struct TraceEventInfo final {
  struct TraceItem;

  explicit TraceEventInfo(kj::ArrayPtr<const kj::Own<Trace>> traces);
  TraceEventInfo(kj::Array<TraceItem> traces): traces(kj::mv(traces)) {}
  TraceEventInfo(rpc::Trace::TraceEventInfo::Reader reader);
  TraceEventInfo(TraceEventInfo&&) noexcept = default;
  TraceEventInfo& operator=(TraceEventInfo&&) = default;
  KJ_DISALLOW_COPY(TraceEventInfo);

  struct TraceItem final {
    explicit TraceItem(kj::Maybe<kj::String> scriptName);
    TraceItem(rpc::Trace::TraceEventInfo::TraceItem::Reader reader);
    TraceItem(TraceItem&&) noexcept = default;
    TraceItem& operator=(TraceItem&&) = default;
    KJ_DISALLOW_COPY(TraceItem);

    kj::Maybe<kj::String> scriptName;

    void copyTo(rpc::Trace::TraceEventInfo::TraceItem::Builder builder) const;
    TraceItem clone() const;
  };

  kj::Vector<TraceItem> traces;

  void copyTo(rpc::Trace::TraceEventInfo::Builder builder) const;
  TraceEventInfo clone() const;
};

// Describes a hibernatable web socket event
struct HibernatableWebSocketEventInfo final {
  struct Message final {};
  struct Close final {
    uint16_t code;
    bool wasClean;
  };
  struct Error final {};

  using Type = kj::OneOf<Message, Close, Error>;

  explicit HibernatableWebSocketEventInfo(Type type);
  HibernatableWebSocketEventInfo(rpc::Trace::HibernatableWebSocketEventInfo::Reader reader);
  HibernatableWebSocketEventInfo(HibernatableWebSocketEventInfo&&) noexcept = default;
  HibernatableWebSocketEventInfo& operator=(HibernatableWebSocketEventInfo&&) = default;
  KJ_DISALLOW_COPY(HibernatableWebSocketEventInfo);

  Type type;

  void copyTo(rpc::Trace::HibernatableWebSocketEventInfo::Builder builder) const;
  HibernatableWebSocketEventInfo clone() const;
  static Type readFrom(rpc::Trace::HibernatableWebSocketEventInfo::Reader reader);
};

// Describes a custom event
struct CustomEventInfo final {
  explicit CustomEventInfo() {};
  CustomEventInfo(rpc::Trace::CustomEventInfo::Reader reader) {};
};

// Describes a fetch response
struct FetchResponseInfo final {
  explicit FetchResponseInfo(uint16_t statusCode);
  FetchResponseInfo(rpc::Trace::FetchResponseInfo::Reader reader);
  FetchResponseInfo(FetchResponseInfo&&) noexcept = default;
  FetchResponseInfo& operator=(FetchResponseInfo&&) = default;
  KJ_DISALLOW_COPY(FetchResponseInfo);

  uint16_t statusCode;

  void copyTo(rpc::Trace::FetchResponseInfo::Builder builder) const;
  FetchResponseInfo clone() const;
};

// Describes an event published using the node:diagnostics_channel API
struct DiagnosticChannelEvent final {
  explicit DiagnosticChannelEvent(
      kj::Date timestamp, kj::String channel, kj::Array<kj::byte> message);
  DiagnosticChannelEvent(rpc::Trace::DiagnosticChannelEvent::Reader reader);
  DiagnosticChannelEvent(DiagnosticChannelEvent&&) noexcept = default;
  KJ_DISALLOW_COPY(DiagnosticChannelEvent);

  kj::Date timestamp;
  kj::String channel;
  kj::Array<kj::byte> message;

  void copyTo(rpc::Trace::DiagnosticChannelEvent::Builder builder) const;
  DiagnosticChannelEvent clone() const;
};

// Describes a stream diagnostics event. Currently only droppedEvents is supported.
struct StreamDiagnosticsEvent final {
  explicit StreamDiagnosticsEvent(uint32_t droppedEventsCount);
  StreamDiagnosticsEvent(rpc::Trace::StreamDiagnosticsEvent::Reader reader);
  StreamDiagnosticsEvent(StreamDiagnosticsEvent&&) noexcept = default;
  KJ_DISALLOW_COPY(StreamDiagnosticsEvent);

  // The count of dropped events for the "droppedEvents" diagnostic. When we support other event
  // types, this should be replaced with a kj::OneOf<> of all the different types.
  uint32_t droppedEventsCount;

  void copyTo(rpc::Trace::StreamDiagnosticsEvent::Builder builder) const;
  StreamDiagnosticsEvent clone() const;
};

// Describes a log event
struct Log final {
  explicit Log(kj::Date timestamp, LogLevel logLevel, kj::String message);
  Log(rpc::Trace::Log::Reader reader);
  Log(Log&&) noexcept = default;
  KJ_DISALLOW_COPY(Log);
  ~Log() noexcept(false) = default;

  // Only as accurate as Worker's Date.now(), for Spectre mitigation.
  kj::Date timestamp;

  LogLevel logLevel;
  // TODO(soon): Just string for now.  Eventually, capture serialized JS objects.
  kj::String message;

  void copyTo(rpc::Trace::Log::Builder builder) const;
  Log clone() const;
};

// Describes an exception event
struct Exception final {
  explicit Exception(
      kj::Date timestamp, kj::String name, kj::String message, kj::Maybe<kj::String> stack);
  Exception(rpc::Trace::Exception::Reader reader);
  Exception(Exception&&) noexcept = default;
  KJ_DISALLOW_COPY(Exception);
  ~Exception() noexcept(false) = default;

  // Only as accurate as Worker's Date.now(), for Spectre mitigation.
  kj::Date timestamp;

  kj::String name;
  kj::String message;

  kj::Maybe<kj::String> stack;

  void copyTo(rpc::Trace::Exception::Builder builder) const;
  Exception clone() const;
};

// EventInfo types are used to describe the onset of an invocation. The FetchEventInfo
// can also be used to describe the start of a fetch subrequest.
using EventInfo = kj::OneOf<FetchEventInfo,
    JsRpcEventInfo,
    ScheduledEventInfo,
    AlarmEventInfo,
    QueueEventInfo,
    EmailEventInfo,
    TraceEventInfo,
    HibernatableWebSocketEventInfo,
    ConnectEventInfo,
    CustomEventInfo>;

EventInfo cloneEventInfo(const EventInfo& info);

template <typename T>
concept AttributeValue = kj::isSameType<kj::ConstString, T>() || kj::isSameType<bool, T>() ||
    kj::isSameType<double, T>() || kj::isSameType<int64_t, T>();

// An Attribute mark is used to add detail to a span over its lifetime.
// The Attribute struct can also be used to provide arbitrary additional
// properties for some other structs.
// Modeled after https://opentelemetry.io/docs/concepts/signals/traces/#attributes
struct Attribute final {
  using Value = kj::OneOf<kj::ConstString, bool, double, int64_t>;
  using Values = kj::Array<Value>;

  explicit Attribute(kj::ConstString name, Value&& value);
  explicit Attribute(kj::ConstString name, Values&& values);

  template <AttributeValue V>
  explicit Attribute(kj::ConstString name, kj::Array<V> vals)
      : Attribute(kj::mv(name), KJ_MAP(v, vals) { return Value(kj::mv(v)); }) {}

  template <AttributeValue V>
  explicit Attribute(kj::ConstString name, std::initializer_list<V> list)
      : Attribute(kj::mv(name), kj::heapArray<V>(list)) {}

  Attribute(rpc::Trace::Attribute::Reader reader);
  Attribute(Attribute&&) noexcept = default;
  Attribute& operator=(Attribute&&) = default;
  KJ_DISALLOW_COPY(Attribute);

  kj::ConstString name;
  Values value;

  void copyTo(rpc::Trace::Attribute::Builder builder) const;
  Attribute clone() const;
  kj::String toString() const;
};
using CustomInfo = kj::Array<Attribute>;
kj::String KJ_STRINGIFY(const CustomInfo& customInfo);

struct SpanOpenData {
  // Represents the data needed for a SpanOpen event
  tracing::SpanId spanId;
  tracing::SpanId parentSpanId;

  kj::ConstString operationName;
  kj::Date startTime;

  SpanOpenData(rpc::SpanOpenData::Reader reader);
  void copyTo(rpc::SpanOpenData::Builder builder) const;
  explicit SpanOpenData(tracing::SpanId spanId,
      tracing::SpanId parentSpanId,
      kj::ConstString operationName,
      kj::Date startTime)
      : spanId(spanId),
        parentSpanId(parentSpanId),
        operationName(kj::mv(operationName)),
        startTime(startTime) {}
};

struct SpanEndData {
  // Represents the data needed when closing a span, including the Attributes and SpanClose events.
  tracing::SpanId spanId;

  kj::Date endTime;
  // Should be Span::TagMap, but we can't forward-declare that.
  kj::HashMap<kj::ConstString, tracing::Attribute::Value> tags;

  SpanEndData(rpc::SpanEndData::Reader reader);
  void copyTo(rpc::SpanEndData::Builder builder) const;
  explicit SpanEndData(tracing::SpanId spanId,
      kj::Date endTime,
      kj::HashMap<kj::ConstString, tracing::Attribute::Value> tags =
          kj::HashMap<kj::ConstString, tracing::Attribute::Value>())
      : spanId(spanId),
        endTime(endTime),
        tags(kj::mv(tags)) {}
};

// A Return mark is used to mark the point at which a span operation returned
// a value. For instance, when a fetch subrequest response is received, or when
// the fetch handler returns a Response. Importantly, it does not signal that the
// span has closed, which may not happen for some period of time after the return
// mark is recorded (e.g. due to things like waitUntils or waiting to fully ready
// the response body payload, etc).
struct Return final {
  explicit Return(kj::Maybe<FetchResponseInfo> info = kj::none);
  Return(rpc::Trace::Return::Reader reader);
  Return(Return&&) noexcept = default;
  Return& operator=(Return&&) = default;
  KJ_DISALLOW_COPY(Return);

  kj::Maybe<FetchResponseInfo> info = kj::none;

  void copyTo(rpc::Trace::Return::Builder builder) const;
  Return clone() const;
};

// Mark events no longer have a corresponding type, but the term generally refers to DiagnosticChannelEvent, Exception, Log, Return, and CustomInfo events.

// Marks the opening of a child span within the streaming tail session.
struct SpanOpen final {
  // If the span represents a subrequest, then the info describes the
  // details of that subrequest.
  using Info = kj::OneOf<FetchEventInfo, JsRpcEventInfo, CustomInfo>;

  explicit SpanOpen(SpanId spanId, kj::ConstString operationName, kj::Maybe<Info> info = kj::none);
  SpanOpen(rpc::Trace::SpanOpen::Reader reader);
  SpanOpen(SpanOpen&&) noexcept = default;
  SpanOpen& operator=(SpanOpen&&) = default;
  KJ_DISALLOW_COPY(SpanOpen);

  kj::ConstString operationName;
  kj::Maybe<Info> info = kj::none;
  SpanId spanId;

  void copyTo(rpc::Trace::SpanOpen::Builder builder) const;
  SpanOpen clone() const;
  kj::String toString() const;
};

// Marks the closing of a child span within the streaming tail session.
// Once emitted, no further mark events should occur within the closed
// span.
struct SpanClose final {
  explicit SpanClose(EventOutcome outcome = EventOutcome::OK);
  SpanClose(rpc::Trace::SpanClose::Reader reader);
  SpanClose(SpanClose&&) noexcept = default;
  SpanClose& operator=(SpanClose&&) = default;
  KJ_DISALLOW_COPY(SpanClose);

  EventOutcome outcome = EventOutcome::OK;

  void copyTo(rpc::Trace::SpanClose::Builder builder) const;
  SpanClose clone() const;
  kj::String toString() const;
};

// The Onset and Outcome event types are special forms of SpanOpen and
// SpanClose that explicitly mark the start and end of the root span.
// A streaming tail session will always begin with an Onset event, and
// always end with an Outcome event.
struct Onset final {
  using Info = EventInfo;

  // Information about the worker that is being tailed.
  struct WorkerInfo final {
    ExecutionModel executionModel = ExecutionModel::STATELESS;
    kj::Maybe<kj::String> scriptName;
    kj::Maybe<kj::Own<ScriptVersion::Reader>> scriptVersion;
    kj::Maybe<TracePreview> preview;
    kj::Maybe<kj::String> dispatchNamespace;
    kj::Maybe<kj::String> scriptId;
    kj::Maybe<kj::Array<kj::String>> scriptTags;
    kj::Maybe<kj::String> entrypoint;

    WorkerInfo clone() const;
  };

  explicit Onset(
      tracing::SpanId spanId, Info&& info, WorkerInfo&& workerInfo, CustomInfo attributes);

  Onset(rpc::Trace::Onset::Reader reader);
  Onset(Onset&&) noexcept = default;
  Onset& operator=(Onset&&) = default;
  KJ_DISALLOW_COPY(Onset);

  tracing::SpanId spanId;
  Info info;
  WorkerInfo workerInfo;
  CustomInfo attributes;

  void copyTo(rpc::Trace::Onset::Builder builder) const;
  Onset clone() const;
};

// Helper functions to copy onset info to/from rpc reader
Onset::Info readOnsetInfo(const rpc::Trace::Onset::Info::Reader& info);
void writeOnsetInfo(const tracing::Onset::Info& info, rpc::Trace::Onset::Info::Builder& builder);

struct Outcome final {
  explicit Outcome(EventOutcome outcome, kj::Duration cpuTime, kj::Duration wallTime);
  Outcome(rpc::Trace::Outcome::Reader reader);
  Outcome(Outcome&&) noexcept = default;
  Outcome& operator=(Outcome&&) = default;
  KJ_DISALLOW_COPY(Outcome);

  EventOutcome outcome = EventOutcome::OK;
  kj::Duration cpuTime;
  kj::Duration wallTime;

  void copyTo(rpc::Trace::Outcome::Builder builder) const;
  Outcome clone() const;
  kj::String toString() const;
};

// A streaming tail worker receives a series of Tail Events. Tail events always
// occur within an InvocationSpanContext. The first TailEvent delivered to a
// streaming tail session is always an Onset. The final TailEvent delivered is
// always an Outcome. Between those can be any number of SpanOpen, SpanClose,
// and Mark events. Every SpanOpen *must* be associated with a SpanClose unless
// the stream was abruptly terminated.
// A future version may add support for Link events again.
struct TailEvent final {
  using Event = kj::OneOf<Onset,
      Outcome,
      SpanOpen,
      SpanClose,
      DiagnosticChannelEvent,
      Exception,
      Log,
      StreamDiagnosticsEvent,
      Return,
      CustomInfo>;

  explicit TailEvent(SpanContext context,
      TraceId invocationId,
      kj::Date timestamp,
      kj::uint sequence,
      Event&& event);
  TailEvent(TraceId traceId,
      TraceId invocationId,
      kj::Maybe<SpanId> spanId,
      kj::Date timestamp,
      kj::uint sequence,
      Event&& event,
      kj::Maybe<TraceFlags> traceFlags = kj::none);
  TailEvent(rpc::Trace::TailEvent::Reader reader);
  TailEvent(TailEvent&&) = default;
  TailEvent& operator=(TailEvent&&) = default;
  KJ_DISALLOW_COPY(TailEvent);

  // The span context this event is associated with.
  SpanContext spanContext;
  TraceId invocationId;

  kj::Date timestamp;  // Unix epoch, Spectre-mitigated resolution
  kj::uint sequence;

  Event event;

  void copyTo(rpc::Trace::TailEvent::Builder builder) const;
  TailEvent clone() const;
};

kj::String KJ_STRINGIFY(const tracing::TailEvent::Event& event);

}  // namespace tracing

enum class PipelineLogLevel {
  // WARNING: This must be kept in sync with PipelineDef::LogLevel (which is not in the OSS
  //   release).
  NONE,
  FULL
};

// TODO(someday): See if we can merge similar code concepts...  Trace fills a role similar to
// MetricsCollector::Reporter::StageEvent, and Tracer fills a role similar to
// MetricsCollector::Request.  Currently, the major differences are:
//
//   - MetricsCollector::Request uses its destructor to measure a IoContext's wall time, so
//     it needs to live exactly as long as its IoContext.  Tracer currently needs to live as
//     long as both the IoContext and those of any subrequests.
//   - Due to the difference in lifetimes, results of each become available in a different order,
//     and intermediate values can be freed at different times.
//   - Request builds a vector of results, while Tracer builds a tree.

// TODO(cleanup) - worth separating into immutable Trace vs. mutable TraceBuilder?

// Collects trace information about the handling of a worker/pipeline fetch event.
class Trace final: public kj::Refcounted {
 public:
  explicit Trace(kj::Maybe<kj::String> stableId,
      kj::Maybe<kj::String> scriptName,
      kj::Maybe<kj::Own<ScriptVersion::Reader>> scriptVersion,
      kj::Maybe<kj::String> dispatchNamespace,
      kj::Maybe<kj::String> scriptId,
      kj::Array<kj::String> scriptTags,
      kj::Maybe<kj::String> entrypoint,
      ExecutionModel executionModel,
      kj::Maybe<kj::String> durableObjectId = kj::none,
      kj::Maybe<tracing::TracePreview> preview = kj::none);
  Trace(rpc::Trace::Reader reader);
  ~Trace() noexcept(false);
  KJ_DISALLOW_COPY_AND_MOVE(Trace);

  // Empty for toplevel worker.
  kj::Maybe<kj::String> stableId;

  // We treat the origin value as "unset".
  kj::Date eventTimestamp = kj::UNIX_EPOCH;

  kj::Maybe<tracing::EventInfo> eventInfo;

  // TODO(someday): Work out what sort of information we may want to convey about the parent
  // trace, if any.

  kj::Maybe<kj::String> scriptName;
  kj::Maybe<kj::Own<ScriptVersion::Reader>> scriptVersion;
  kj::Maybe<kj::String> dispatchNamespace;
  kj::Maybe<kj::String> scriptId;
  kj::Array<kj::String> scriptTags;
  kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes;
  kj::Maybe<kj::String> entrypoint;
  kj::Maybe<tracing::TracePreview> preview;
  kj::Maybe<kj::String> durableObjectId;

  kj::Vector<tracing::Log> logs;
  // A request's trace can have multiple exceptions due to separate request/waitUntil tasks.
  kj::Vector<tracing::Exception> exceptions;

  kj::Vector<tracing::DiagnosticChannelEvent> diagnosticChannelEvents;

  ExecutionModel executionModel;
  EventOutcome outcome = EventOutcome::UNKNOWN;

  kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo;

  kj::Duration cpuTime;
  kj::Duration wallTime;

  bool truncated = false;
  bool exceededLogLimit = false;
  bool exceededExceptionLimit = false;
  bool exceededDiagnosticChannelEventLimit = false;
  // Trace data is recorded outside of the JS heap.  To avoid DoS, we keep an estimate of trace
  // data size, and we stop recording if too much is used.
  size_t bytesUsed = 0;

  // Copy content from this trace into `builder`.
  void copyTo(rpc::Trace::Builder builder) const;

  // Adds all content from `reader` to this `Trace`. (Typically this trace is empty before the
  // call.)  Also applies filtering to the trace as if it were recorded with the given
  // pipelineLogLevel.
  void mergeFrom(rpc::Trace::Reader reader, PipelineLogLevel pipelineLogLevel);
};

// =======================================================================================

// Helper function used when setting "truncated_script_id" tags. Truncates the scriptId to 10
// characters.
inline kj::String truncateScriptId(kj::StringPtr id) {
  auto truncatedId = id.first(kj::min(id.size(), 10));
  return kj::str(truncatedId);
}

// =======================================================================================
// Span tracing
//
// TODO(cleanup): (Streaming) tail workers now have access to span tracing as well, but with that
// the trace worker and span interfaces are still mostly independent of each other; separate span
// tracing into a separate header.

class SpanBuilder;
class SpanObserver;

struct Span {
  // Represents a trace span. `Span` objects are delivered to `SpanObserver`s for recording. To
  // create a `Span`, use a `SpanBuilder`.

 public:
  using TagValue = tracing::Attribute::Value;
  // TODO(someday): Support binary bytes, too.
  using TagMap = kj::HashMap<kj::ConstString, TagValue>;
  using Tag = TagMap::Entry;

  struct Log {
    kj::Date timestamp;
    Tag tag;
  };

  kj::ConstString operationName;
  kj::Date startTime;
  kj::Date endTime;
  TagMap tags;
  kj::Vector<Log> logs;

  // We set an arbitrary (-ish) cap on log messages for safety. If we drop logs because of this,
  // we report how many in a final "dropped_logs" log.
  //
  // At the risk of being too clever, I chose a limit that is one below a power of two so that
  // we'll typically have space for one last element available for the "dropped_logs" log without
  // needing to grow the vector.
  static constexpr auto MAX_LOGS = 1023;
  uint droppedLogs = 0;

  explicit Span(kj::ConstString operationName, kj::Date startTime)
      : operationName(kj::mv(operationName)),
        startTime(startTime),
        endTime(startTime) {}
};

// Utility functions for handling span tags.
void serializeTagValue(rpc::TagValue::Builder builder, const Span::TagValue& value);
Span::TagValue deserializeTagValue(rpc::TagValue::Reader value);

// Clone function for span tags, avoids memory allocation for string literals and non-string values.
Span::TagValue spanTagClone(const Span::TagValue& tag);

// An opaque token which can be used to create child spans of some parent. This is typically
// passed down from a caller to a callee when the caller wants to allow the callee to create
// spans for itself that show up as children of the caller's span, but the caller does not
// want to give the callee any other ability to modify the parent span.
class SpanParent {
 public:
  SpanParent(SpanBuilder& builder);

  // Make a SpanParent that causes children not to be reported anywhere.
  SpanParent(decltype(nullptr)) {}

  SpanParent(kj::Maybe<kj::Own<SpanObserver>> observer): observer(kj::mv(observer)) {}

  SpanParent(SpanParent&& other) = default;
  SpanParent& operator=(SpanParent&& other) = default;
  KJ_DISALLOW_COPY(SpanParent);

  SpanParent addRef();

  // Create a new child span.
  //
  // `operationName` should be a string literal with infinite lifetime.
  [[nodiscard]] SpanBuilder newChild(
      kj::ConstString operationName, kj::Maybe<kj::Date> startTime = kj::none);

  // Useful to skip unnecessary code when not observed.
  bool isObserved() {
    return observer != kj::none;
  }

  // Get the underlying SpanObserver representing the parent span.
  //
  // This is needed in particular when making outbound network requests that must be annotated with
  // trace IDs in a way that is specific to the trace back-end being used. The caller must downcast
  // the `SpanObserver` to the expected observer type in order to extract the trace ID.
  kj::Maybe<SpanObserver&> getObserver() {
    return observer;
  }

  // Return the serializable identity of this span for cross-boundary propagation.
  kj::Maybe<tracing::SpanContext> toSpanContext();

  // Returns the observer's spanId, or SpanId::nullId if there is none.
  tracing::SpanId getSpanId();

 private:
  kj::Maybe<kj::Own<SpanObserver>> observer;
};

// Interface for writing a span. Essentially, this is a mutable interface to a `Span` object,
// given only to the code which is meant to create the span, whereas code that merely collects
// and reports spans gets the `Span` type.
//
// The reason we use a separate builder type rather than rely on constness is so that the methods
// can be no-ops when there is no observer, avoiding unnecessary allocations. To allow for this,
// SpanBuilder is designed to be write-only -- you cannot read back the content. Only the
// observer (if there is one) receives the content.
class SpanBuilder {
 public:
  // Create a new top-level span that will report to the given observer. If the observer is null,
  // no data is collected.
  //
  // `operationName` should be a string literal with infinite lifetime, or somehow otherwise be
  // attached to the observer observing this span.
  explicit SpanBuilder(kj::Maybe<kj::Own<SpanObserver>> observer,
      kj::ConstString operationName,
      kj::Maybe<kj::Date> startTime = kj::none);

  // Make a SpanBuilder that ignores all calls. (Useful if you want to assign it later.)
  SpanBuilder(decltype(nullptr)) {}

  SpanBuilder(SpanBuilder&& other) = default;
  SpanBuilder& operator=(SpanBuilder&& other);  // ends the existing span and starts a new one
  KJ_DISALLOW_COPY(SpanBuilder);

  ~SpanBuilder() noexcept(false);

  // Finishes and submits the span. This is done implicitly by the destructor, but sometimes it's
  // useful to be able to submit early. The SpanBuilder ignores all further method calls after this
  // is invoked.
  void end();

  // Useful to skip unnecessary code when not observed.
  bool isObserved() {
    return observer != kj::none;
  }

  // Get the underlying SpanObserver representing the span.
  //
  // This is needed in particular when making outbound network requests that must be annotated with
  // trace IDs in a way that is specific to the trace back-end being used. The caller must downcast
  // the `SpanObserver` to the expected observer type in order to extract the trace ID.
  kj::Maybe<SpanObserver&> getObserver() {
    return observer;
  }

  // Create a new child span.
  //
  // `operationName` should be a string literal with infinite lifetime.
  [[nodiscard]] SpanBuilder newChild(
      kj::ConstString operationName, kj::Maybe<kj::Date> startTime = kj::none);

  // Change the operation name from what was specified at span creation.
  //
  // `operationName` should be a string literal with infinite lifetime.
  void setOperationName(kj::ConstString operationName);

  using TagValue = Span::TagValue;
  // `key` must point to memory that will remain valid all the way until this span's data is
  // serialized.
  // Allow setting tags with an extended set of types to elide string allocations when we have a
  // string literal or are not being observed. We include String/LiteralStringConst here to avoid
  // having to manually cast them to ConstString each time.
  using TagInitValue = kj::OneOf<kj::StringPtr,
      kj::String,
      kj::LiteralStringConst,
      kj::ConstString,
      bool,
      double,
      int64_t>;

  void setTag(kj::ConstString key, TagInitValue value);

  // `key` must point to memory that will remain valid all the way until this span's data is
  // serialized.
  //
  // The differences between this and `setTag()` is that logs are timestamped and may have
  // duplicate keys.
  void addLog(kj::Date timestamp, kj::ConstString key, TagValue value);

 private:
  kj::Maybe<kj::Own<SpanObserver>> observer;
  // The under-construction span, or null if the span has ended.
  kj::Maybe<Span> span;

  friend class SpanParent;
};

// Abstract interface for observing trace spans reported by the runtime. Different
// implementations might support different tracing back-ends, e.g. Trace Workers, Jaeger, or
// whatever infrastructure you prefer to use for this.
//
// A new SpanObserver is created at the start of each Span. The SpanBuilder drives the observer
// through its lifecycle: onOpen() is called when the span is created, onClose() when the span
// ends, and onUpdateName() if the operation name changes between open and close.
class SpanObserver: public kj::Refcounted {
 public:
  // Allocate a new child span.
  //
  // Note that children can be created long after a span has completed.
  [[nodiscard]] virtual kj::Own<SpanObserver> newChild() = 0;

  // Allocate a child for a span initiated directly by user JavaScript (via
  // `ctx.tracing.enterSpan`). Allows implementations to apply different policies than for
  // runtime-issued spans (notably, edgeworker bypasses its operation-name allowlist here).
  [[nodiscard]] virtual kj::Own<SpanObserver> newChildFromUserCode() {
    return newChild();
  }

  // Called when the span is opened. Delivers the initial operation name and start time.
  // Called exactly once, before any other lifecycle method.
  virtual void onOpen(kj::ConstString operationName, kj::Date startTime) = 0;

  // Called when the span is closed. Delivers the end time, tags, and logs.
  // Called exactly once per observer, after onOpen(). Tags and logs are moved from the span;
  // the observer takes ownership.
  virtual void onClose(kj::Date endTime, Span::TagMap&& tags, kj::Vector<Span::Log>&& logs) = 0;

  // Called when the operation name is changed after the span was opened (via
  // SpanBuilder::setOperationName()). Observers that eagerly stream the open event should handle
  // this; others may simply update their buffered state. Default implementation is a no-op.
  virtual void onUpdateName(kj::ConstString operationName) {}

  // The current time to be provided for the span. For user tracing, we will override this to
  // provide I/O time. This *requires* that spans are only created when an IOContext is available
  // (usually it is difficult to violate this assumption, but care must be taken that the observer
  // isn't used directly to create a span before the IoContext has been constructed (previously this
  // was a case with a top-level span owned by the WorkerTracer itself).
  virtual kj::Date getTime() {
    return kj::systemPreciseCalendarClock().now();
  }

  // Return the serializable identity of this span for cross-boundary propagation.
  // Returns kj::none if this observer doesn't carry identity.
  virtual kj::Maybe<tracing::SpanContext> toSpanContext() {
    return kj::none;
  }

  // Returns this observer's spanId, or SpanId::nullId if it has no identity.
  virtual tracing::SpanId getSpanId() {
    return tracing::SpanId::nullId;
  }
};

inline kj::Maybe<tracing::SpanContext> SpanParent::toSpanContext() {
  KJ_IF_SOME(obs, observer) {
    return obs->toSpanContext();
  }
  return kj::none;
}

inline tracing::SpanId SpanParent::getSpanId() {
  KJ_IF_SOME(obs, observer) {
    return obs->getSpanId();
  }
  return tracing::SpanId::nullId;
}

inline SpanParent::SpanParent(SpanBuilder& builder): observer(mapAddRef(builder.observer)) {}

inline SpanParent SpanParent::addRef() {
  return SpanParent(mapAddRef(observer));
}

inline SpanBuilder SpanParent::newChild(
    kj::ConstString operationName, kj::Maybe<kj::Date> startTime) {
  return SpanBuilder(observer.map([](kj::Own<SpanObserver>& obs) { return obs->newChild(); }),
      kj::mv(operationName), startTime);
}

inline SpanBuilder SpanBuilder::newChild(
    kj::ConstString operationName, kj::Maybe<kj::Date> startTime) {
  return SpanBuilder(observer.map([](kj::Own<SpanObserver>& obs) { return obs->newChild(); }),
      kj::mv(operationName), startTime);
}

// TraceContext to keep track of user tracing/existing tracing better
class TraceContext {
 public:
  TraceContext(): span(nullptr), userSpan(nullptr) {}
  TraceContext(SpanBuilder span, SpanBuilder userSpan)
      : span(kj::mv(span)),
        userSpan(kj::mv(userSpan)) {}
  TraceContext(TraceContext&& other) = default;
  TraceContext& operator=(TraceContext&& other) = default;
  KJ_DISALLOW_COPY(TraceContext);

  // Set a tag on both the internal span and user span.
  void setTag(kj::ConstString key, SpanBuilder::TagInitValue value);
  bool isObserved() {
    return span.isObserved() || userSpan.isObserved();
  }
  SpanParent getInternalSpanParent() {
    return SpanParent(span);
  }

  SpanParent getUserSpanParent() {
    return SpanParent(userSpan);
  }

 private:
  SpanBuilder span;
  SpanBuilder userSpan;
};

// RAII object that measures the time duration over its lifetime. It tags this duration onto a
// given request span using a specified tag name. Ideal for automatically tracking and logging
// execution times within a scoped block.
class ScopedDurationTagger {
 public:
  explicit ScopedDurationTagger(
      SpanBuilder& span, kj::ConstString key, const kj::MonotonicClock& timer);
  ~ScopedDurationTagger() noexcept(false);
  KJ_DISALLOW_COPY_AND_MOVE(ScopedDurationTagger);

 private:
  SpanBuilder& span;
  kj::ConstString key;
  const kj::MonotonicClock& timer;
  const kj::TimePoint startTime;
};

}  // namespace workerd