# encoding: utf-8
# This file is distributed under New Relic's license terms.
# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
# frozen_string_literal: true

module NewRelic::Agent::Instrumentation
  module Memcache
    module Tracer
      SLASH = '/'
      UNKNOWN = 'unknown'
      LOCALHOST = 'localhost'
      MULTIGET_METRIC_NAME = "get_multi_request"
      MEMCACHED = "Memcached"

      def with_newrelic_tracing(operation, *args)
        segment = NewRelic::Agent::Tracer.start_datastore_segment( \
          product: MEMCACHED,
          operation: operation
        )

        begin
          NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
        ensure
          if NewRelic::Agent.config[:capture_memcache_keys]
            segment.notice_nosql_statement("#{operation} #{args.first.inspect}")
          end
          segment.finish if segment
        end
      end

      def server_for_key_with_newrelic_tracing
        yield.tap do |server|
          begin
            if txn = ::NewRelic::Agent::Tracer.current_transaction
              segment = txn.current_segment
              if ::NewRelic::Agent::Transaction::DatastoreSegment === segment
                assign_instance_to(segment, server)
              end
            end
          rescue => e
            ::NewRelic::Agent.logger.warn("Unable to set instance info on datastore segment: #{e.message}")
          end
        end
      end

      def get_multi_with_newrelic_tracing(method_name)
        segment = NewRelic::Agent::Tracer.start_segment( \
          name: "Ruby/Memcached/Dalli/#{method_name}"
        )

        begin
          NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
        ensure
          segment.finish if segment
        end
      end

      def send_multiget_with_newrelic_tracing(keys)
        segment = ::NewRelic::Agent::Tracer.start_datastore_segment( \
          product: MEMCACHED,
          operation: MULTIGET_METRIC_NAME
        )

        begin
          assign_instance_to(segment, self)
          NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
        ensure
          if ::NewRelic::Agent.config[:capture_memcache_keys]
            segment.notice_nosql_statement("#{MULTIGET_METRIC_NAME} #{keys.inspect}")
          end
          segment.finish if segment
        end
      end

      def assign_instance_to(segment, server)
        host = port_path_or_id = nil
        if server.hostname.start_with?(SLASH)
          host = LOCALHOST
          port_path_or_id = server.hostname
        else
          host = server.hostname
          port_path_or_id = server.port
        end
        segment.set_instance_info(host, port_path_or_id)
      rescue => e
        ::NewRelic::Agent.logger.debug("Failed to retrieve memcached instance info: #{e.message}")
        segment.set_instance_info(UNKNOWN, UNKNOWN)
      end
    end
  end
end