class ElasticAPM::TraceContext

def apply_headers
  yield 'Traceparent', traceparent.to_header
  if tracestate && !tracestate.to_header.empty?
    yield 'Tracestate', tracestate.to_header
  end
  return unless ElasticAPM.agent.config.use_elastic_traceparent_header
  yield 'Elastic-Apm-Traceparent', traceparent.to_header
end

def child
  dup.tap do |tc|
    tc.traceparent = tc.traceparent.child
  end
end

def initialize(
  traceparent: nil,
  tracestate: nil
)
  @traceparent = traceparent || Traceparent.new
  @tracestate = tracestate || Tracestate.new
end

def parse(env: nil, metadata: nil)
  unless env || metadata
    raise ArgumentError, 'TraceContext expects env:, metadata: ' \
      'or single argument header string'
  end
  if env
    trace_context_from_env(env)
  elsif metadata
    trace_context_from_metadata(metadata)
  end
end

def trace_context_from_env(env)
  return unless (
    header =
      env['HTTP_ELASTIC_APM_TRACEPARENT'] || env['HTTP_TRACEPARENT']
  )
  parent = TraceContext::Traceparent.parse(header)
  state =
    if (header = env['HTTP_TRACESTATE'])
      TraceContext::Tracestate.parse(header)
    end
  new(traceparent: parent, tracestate: state)
end

def trace_context_from_metadata(metadata)
  return unless (header = metadata['elastic-apm-traceparent'] ||
    metadata['traceparent'])
  parent = TraceContext::Traceparent.parse(header)
  state =
    if (header = metadata['tracestate'])
      TraceContext::Tracestate.parse(header)
    end
  new(traceparent: parent, tracestate: state)
end