# frozen_string_literal: true

require "opentelemetry/semantic_conventions"
require "opentelemetry-api"

module Aspecto
  # Aspecto OpenTelemetry Distro
  module OpenTelemetry
    module Sampler
      # OpenTelemetry sampler which implements the remote rules logic with fallback to service sampler
      class MessageProcessSampler
        def initialize(rules_sampler, service_sampler)
          @rules_sampler = rules_sampler
          @service_sampler = service_sampler
        end

        def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
          if attributes && attributes[::OpenTelemetry::SemanticConventions::Trace::MESSAGING_OPERATION] == "process"
            sampling_decision = should_sample_processing(parent_context, links, name, attributes)
            return sampling_decision if sampling_decision
          end
          @service_sampler.should_sample?(trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes)
        end

        def description
          "MessageProcessSampler"
        end

        private

        def should_sample_processing(parent_context, links, name, attributes)
          return if !links || links.empty?

          decision = link_decision(links[0])

          # sidekiq receive side is a single "process" span which is the entry for the trace
          is_root = ::OpenTelemetry::Trace.current_span(parent_context) == ::OpenTelemetry::Trace::Span::INVALID
          if is_root # rubocop:disable Style/GuardClause
            matched_rule = @rules_sampler.matching_rule(name: name, attributes: attributes)
            rule_context = matched_rule&.context
            ::OpenTelemetry::SDK::Trace::Samplers::Result.new(decision: decision, tracestate: ::OpenTelemetry::Trace.current_span(parent_context).context.tracestate) if rule_context && rule_context["messaging.sampling.inherit_from_publisher"]
          end

          # if not root, then we need to check if the entry span set "messaging.sampling.inherit_from_publisher" on the context.
          # since we currently only support sidekiq, it is not needed.
          # but in the future, it can be implemented with context.attach when the rule is applied, and context.get here
        end

        def link_decision(link)
          publisher_sampled = link.span_context.trace_flags.sampled?
          publisher_sampled ? ::OpenTelemetry::SDK::Trace::Samplers::Decision::RECORD_AND_SAMPLE : ::OpenTelemetry::SDK::Trace::Samplers::Decision::DROP
        end
      end
    end
  end
end