# Copyright 2017 OpenCensus 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.

module OpenCensus
  module Trace
    ##
    # SpanContext represents the context within which a span may be created.
    # It includes the ID of the parent trace, the ID of the parent span, and
    # sampling state.
    #
    class SpanContext
      ##
      # Internal struct that holds trace-wide data.
      #
      # @private
      #
      TraceData = Struct.new :trace_id, :trace_options, :span_map

      ##
      # Maximum integer value for a `trace_id`
      #
      # @private
      #
      MAX_TRACE_ID = 0xffffffffffffffffffffffffffffffff

      ##
      # Maximum integer value for a `span_id`
      #
      # @private
      #
      MAX_SPAN_ID = 0xffffffffffffffff

      class << self
        ##
        # Create a new root SpanContext object, given either a Trace-Context
        # header value by itself, or an entire Rack environment. If a valid
        # Trace-Context header can be obtained from either source, it is used
        # to generate the SpanContext. Otherwise, a new root context with a
        # unique `trace_id` and a root `span_id` of "" is used.
        #
        # @param [TraceContextData] trace_context The request's incoming trace
        #      context (optional)
        # @param [boolean, nil] same_process_as_parent Set to `true` if the
        #      parent span is local, `false` if it is remote, or `nil` if there
        #      is no parent span or this information is not available.
        #
        # @return [SpanContext]
        #
        def create_root trace_context: nil, same_process_as_parent: nil
          if trace_context
            trace_data = TraceData.new \
              trace_context.trace_id, trace_context.trace_options, {}
            new trace_data, nil, trace_context.span_id,
                same_process_as_parent
          else
            trace_id = rand 1..MAX_TRACE_ID
            trace_id = trace_id.to_s(16).rjust(32, "0")
            trace_data = TraceData.new trace_id, 0, {}
            new trace_data, nil, "", nil
          end
        end
      end

      ##
      # The parent of this context, or `nil` if this is a root context.
      #
      # @return [SpanContext, nil]
      #
      attr_reader :parent

      ##
      # Returns true if this is a root span context
      #
      # @return [boolean]
      #
      def root?
        parent.nil?
      end

      ##
      # The root context, which may be this context or one of its ancestors.
      #
      # @return [SpanContext]
      #
      def root
        root = self
        until (parent = root.parent).nil?
          root = parent
        end
        root
      end

      ##
      # Returns the trace context for this span.
      #
      # @return [TraceContextData]
      #
      def trace_context
        TraceContextData.new trace_id, @span_id, trace_options
      end

      ##
      # The trace ID, as a 32-character hex string.
      #
      # @return [String]
      #
      def trace_id
        @trace_data.trace_id
      end

      ##
      # The original trace options byte used to create this span context.
      #
      # @return [Integer]
      #
      def trace_options
        @trace_data.trace_options
      end

      ##
      # The span ID as a 16-character hex string, or the empty string if the
      # context refers to the root of the trace.
      #
      # @return [String]
      #
      attr_reader :span_id

      ##
      # Whether the parent of spans created by this context is local, or `nil`
      # if this context creates root spans or this information is unknown.
      #
      # @return [boolean, nil]
      #
      attr_reader :same_process_as_parent

      ##
      # Whether the context (e.g. the parent span) has been sampled. This
      # information may be used in sampling decisions for new spans.
      #
      # @return [boolean]
      #
      def sampled?
        span = this_span
        if span
          span.sampled
        else
          trace_options & 0x01 != 0
        end
      end

      ##
      # Create a new span in this context.
      # You must pass a name for the span. All other span attributes should
      # be set using the SpanBuilder methods.
      # The span will be started automatically with the current timestamp.
      # However, you are responsible for finishing the span yourself.
      #
      # @param [String] name Name of the span
      # @param [Sampler] sampler Span-scoped sampler. If not provided,
      #     defaults to the trace configuration's default sampler.
      #
      # @return [SpanBuilder] A SpanBuilder object that you can use to
      #     set span attributes and create children.
      #
      def start_span name, skip_frames: 0, sampler: nil
        child_context = create_child
        sampler ||= OpenCensus::Trace.config.default_sampler
        sampled = sampler.call span_context: self
        span = SpanBuilder.new child_context, sampled,
                               skip_frames: skip_frames + 1
        span.name = name
        span.start!
        @trace_data.span_map[child_context.span_id] = span
      end

      ##
      # Create a new span in this context.
      # You must pass a name for the span. All other span attributes should
      # be set using the SpanBuilder methods.
      #
      # The span will be started automatically with the current timestamp. The
      # SpanBuilder will then be passed to the block you provide. The span will
      # be finished automatically at the end of the block.
      #
      # @param [String] name Name of the span
      # @param [Sampler] sampler Span-scoped sampler. If not provided,
      #     defaults to the trace configuration's default sampler.
      #
      def in_span name, skip_frames: 0, sampler: nil
        span = start_span name, skip_frames: skip_frames + 1, sampler: sampler
        begin
          yield span
        ensure
          end_span span
        end
      end

      ##
      # Finish the given span, which must have been created by this span
      # context.
      #
      # @param [SpanBuilder] span The span to finish.
      #
      def end_span span
        unless span.context.parent == self
          raise "The given span was not created by this context"
        end
        span.finish!
      end

      ##
      # Returns the span that defines this context; that is, the span that is
      # the parent of spans created by this context. Returns `nil` if this
      # context is the root and doesn't correspond to an actual span, or if
      # the corresponding span is remote.
      #
      # @return [SpanBuilder, nil] The span defining this context.
      #
      def this_span
        get_span @span_id
      end

      ##
      # Builds all finished spans under this context, and returns an array of
      # built `Span` objects. Ignores any unfinished spans. The order of the
      # generated spans is undefined.
      #
      # Does not build any ancestor spans. If you want the entire span tree
      # built, call this method on the `#root` context.
      #
      # @param [Integer, nil] max_attributes The maximum number of attributes
      #     to save, or `nil` to use the config value.
      # @param [Integer, nil] max_stack_frames The maximum number of stack
      #     frames to save, or `nil` to use the config value.
      # @param [Integer, nil] max_annotations The maximum number of annotations
      #     to save, or `nil` to use the config value.
      # @param [Integer, nil] max_message_events The maximum number of message
      #     events to save, or `nil` to use the config value.
      # @param [Integer, nil] max_links The maximum number of links to save,
      #     or `nil` to use the config value.
      # @param [Integer, nil] max_string_length The maximum length in bytes for
      #     truncated strings, or `nil` to use the config value.
      #
      # @return [Array<Span>] Built Span objects
      #
      def build_contained_spans max_attributes: nil,
                                max_stack_frames: nil,
                                max_annotations: nil,
                                max_message_events: nil,
                                max_links: nil,
                                max_string_length: nil
        contained_span_builders.find_all(&:finished?).map do |sb|
          sb.to_span max_attributes: max_attributes,
                     max_stack_frames: max_stack_frames,
                     max_annotations: max_annotations,
                     max_message_events: max_message_events,
                     max_links: max_links,
                     max_string_length: max_string_length
        end
      end

      ##
      # Initialize a SpanContext object. This low-level constructor is used
      # internally only. Generally, you should create a SpanContext using the
      # `SpanContext.create_root` method.
      #
      # @private
      #
      def initialize trace_data, parent, span_id, same_process_as_parent
        @trace_data = trace_data
        @parent = parent
        @span_id = span_id
        @same_process_as_parent = same_process_as_parent
      end

      ##
      # Returns true if this context equals or is an ancestor of the given
      # context.
      #
      # @private
      # @return [boolean]
      #
      def contains? context
        until context.nil?
          return true if context == self
          context = context.parent
        end
        false
      end

      ##
      # Returns all SpanBuilder objects created by this context or any
      # descendant context. The order of the returned spans is undefined.
      #
      # @private
      # @return [Array<SpanBuilder>]
      #
      def contained_span_builders
        builders = @trace_data.span_map.values
        if root?
          builders
        else
          builders.find_all { |sb| contains? sb.context.parent }
        end
      end

      private

      ##
      # Create a child of this SpanContext, with a random unique span ID.
      #
      # @return [SpanContext] The created child context.
      #
      def create_child
        loop do
          child_span_id = rand 1..MAX_SPAN_ID
          child_span_id = child_span_id.to_s(16).rjust(16, "0")
          unless @trace_data.span_map.key? child_span_id
            return SpanContext.new @trace_data, self, child_span_id, true
          end
        end
      end

      ##
      # Get the SpanBuilder given a Span ID.
      #
      # @param [Integer] span_id the ID of the span to get
      # @return [SpanBuilder, nil] The SpanBuilder, or `nil` if ID not found
      #
      def get_span span_id
        @trace_data.span_map[span_id]
      end
    end
  end
end