# encoding: utf-8

LibraryDetection.defer do
  named :curb

  OA_CURB_MINIMUM_VERSION = ::OneApm::VersionNumber.new("0.8.1")

  depends_on do
    defined?(::Curl) && defined?(::Curl::CURB_VERSION) &&
      ::OneApm::VersionNumber.new(::Curl::CURB_VERSION) >= OA_CURB_MINIMUM_VERSION
  end

  executes do
    OneApm::Manager.logger.info 'Installing Curb instrumentation'
    require 'one_apm/agent/cross_app/cross_app_tracing'
    require 'one_apm/support/http_clients/curb_wrappers'
  end

  executes do
    class Curl::Easy

      attr_accessor :_oa_instrumented,
                    :_oa_http_verb,
                    :_oa_header_str,
                    :_oa_original_on_header,
                    :_oa_original_on_complete,
                    :_oa_serial

      # We have to hook these three methods separately, as they don't use
      # Curl::Easy#http
      def http_head_with_oneapm(*args, &blk)
        self._oa_http_verb = :HEAD
        http_head_without_oneapm(*args, &blk)
      end
      alias_method :http_head_without_oneapm, :http_head
      alias_method :http_head, :http_head_with_oneapm

      def http_post_with_oneapm(*args, &blk)
        self._oa_http_verb = :POST
        http_post_without_oneapm(*args, &blk)
      end
      alias_method :http_post_without_oneapm, :http_post
      alias_method :http_post, :http_post_with_oneapm

      def http_put_with_oneapm(*args, &blk)
        self._oa_http_verb = :PUT
        http_put_without_oneapm(*args, &blk)
      end
      alias_method :http_put_without_oneapm, :http_put
      alias_method :http_put, :http_put_with_oneapm


      # Hook the #http method to set the verb.
      def http_with_oneapm( verb )
        self._oa_http_verb = verb.to_s.upcase
        http_without_oneapm( verb )
      end

      alias_method :http_without_oneapm, :http
      alias_method :http, :http_with_oneapm


      # Hook the #perform method to mark the request as non-parallel.
      def perform_with_oneapm
        self._oa_serial = true
        perform_without_oneapm
      end

      alias_method :perform_without_oneapm, :perform
      alias_method :perform, :perform_with_oneapm

      # We override this method in order to ensure access to header_str even
      # though we use an on_header callback
      def header_str_with_oneapm
        if self._oa_serial
          self._oa_header_str
        else
          # Since we didn't install a header callback for a non-serial request,
          # just fall back to the original implementation.
          header_str_without_oneapm
        end
      end

      alias_method :header_str_without_oneapm, :header_str
      alias_method :header_str, :header_str_with_oneapm
    end # class Curl::Easy


    class Curl::Multi
      include OneApm::Support::MethodTracer

      # Add CAT with callbacks if the request is serial
      def add_with_oneapm(curl)
        if curl.respond_to?(:_oa_serial) && curl._oa_serial
          hook_pending_request(curl) if OneApm::Manager.tl_is_execution_traced?
        end

        return add_without_oneapm( curl )
      end

      alias_method :add_without_oneapm, :add
      alias_method :add, :add_with_oneapm


      # Trace as an External/Multiple call if the first request isn't serial.
      def perform_with_oneapm(&blk)
        return perform_without_oneapm if
          self.requests.first &&
          self.requests.first.respond_to?(:_oa_serial) &&
          self.requests.first._oa_serial

        trace_execution_scoped("External/Multiple/Curb::Multi/perform") do
          perform_without_oneapm(&blk)
        end
      end

      alias_method :perform_without_oneapm, :perform
      alias_method :perform, :perform_with_oneapm


      # Instrument the specified +request+ (a Curl::Easy object) and set up cross-application
      # tracing if it's enabled.
      def hook_pending_request(request)
        wrapped_request, wrapped_response = wrap_request(request)
        state   = OneApm::TransactionState.tl_get
        t0      = Time.now
        segment = OneApm::Agent::CrossAppTracing.start_trace(state, t0, wrapped_request)

        unless request._oa_instrumented
          install_header_callback(request, wrapped_response)
          install_completion_callback(request, t0, segment, wrapped_request, wrapped_response)
          request._oa_instrumented = true
        end
      rescue => err
        OneApm::Manager.logger.error("Untrapped exception", err)
      end


      # Create request and response adapter objects for the specified +request+
      def wrap_request(request)
        return OneApm::Support::HTTPClients::CurbRequest.new(request),
               OneApm::Support::HTTPClients::CurbResponse.new(request)
      end


      # Install a callback that will record the response headers to enable
      # CAT linking
      def install_header_callback( request, wrapped_response )
        original_callback = request.on_header
        request._oa_original_on_header = original_callback
        request._oa_header_str = ''
        request.on_header do |header_data|
          wrapped_response.append_header_data( header_data )

          if original_callback
            original_callback.call( header_data )
          else
            header_data.length
          end
        end
      end

      # Install a callback that will finish the trace.
      def install_completion_callback(request, t0, segment, wrapped_request, wrapped_response)
        original_callback = request.on_complete
        request._oa_original_on_complete = original_callback
        request.on_complete do |finished_request|
          begin
            state = OneApm::TransactionState.tl_get
            OneApm::Agent::CrossAppTracing.finish_trace(state, t0, segment, wrapped_request, wrapped_response)
          ensure
            # Make sure the existing completion callback is run, and restore the
            # on_complete callback to how it was before.
            original_callback.call(finished_request) if original_callback
            remove_instrumentation_callbacks(request)
          end
        end
      end

      def remove_instrumentation_callbacks(request)
        request.on_complete(&request._oa_original_on_complete)
        request.on_header(&request._oa_original_on_header)
        request._oa_instrumented = false
      end

    end # class Curl::Multi

  end
end