# encoding: utf-8

module OneApm
  module Agent
    module  Instrumentation
      module ThriftHelper

        attr_accessor :segment

        def rpc_request
          @rpc_request || {}
        end

        def rpc_request=(rpc_request)
          @rpc_request = rpc_request
        end

        def rpc_reset!
          @rpc_request = nil
          @segment = nil
        end

        def operator result_klass
          namespaces = result_klass.to_s.split('::')
          operator_name = 'unknown'
          if namespaces.last =~ /_result/
            operator_name = namespaces.last.sub('_result', '').downcase
          end
          operator_name
        end

        def thost
         @thost ||= @iprot.instance_variable_get("@trans") \
          .instance_variable_get("@transport") \
          .instance_variable_get("@host") rescue nil
        end

        def metrics operator_name
          metrics = if thost.nil?
                      ["External/Thrift/#{operator_name}"]
                    else
                      #changed from */host/Thrift to */Thrift/host
                      ["External/Thrift/#{thost}/#{operator_name}", "External/#{thost}/all"]
                    end
          metrics << "External/all"
          metrics
        end

        def started_time op
          _op_ = operations.delete(op)
          (_op_ && _op_[:started_time]) or Time.now.to_f
        end

        def operations
          @operations ||= {}
        end
      end
    end
  end
end

LibraryDetection.defer do
  named :thrift

  depends_on do
    defined?(::Thrift) && defined?(::Thrift::Client)
  end

  executes do
    OneApm::Manager.logger.info 'Installing Thrift instrumentation'
  end

  executes do
    require 'one_apm/agent/cross_app/cross_app_tracing_message'

    ::Thrift::Client.module_eval do 
      include OneApm::Agent::Instrumentation::ThriftHelper

      def send_message_with_oneapm(name, args_class, args = {})
        state   = OneApm::TransactionState.tl_get
        t0      = Time.now
        self.segment = OneApm::Agent::CrossAppTracingMessage.start_trace(state, t0, rpc_request)
        operations[name] = { :started_time => t0 }
        send_message_without_oneapm(name, args_class, args)
      end
      alias :send_message_without_oneapm :send_message
      alias :send_message :send_message_with_oneapm

      def send_oneway_message_with_oneapm(name, args_class, args = {})
        begin
          state   = OneApm::TransactionState.tl_get
          t0      = Time.now
          segment = OneApm::Agent::CrossAppTracingMessage.start_trace(state, t0, rpc_request)
          result  = send_oneway_message_without_oneapm(name, args_class, args)
          OneApm::Agent::CrossAppTracingMessage.finish_trace(state, t0, segment, rpc_request, metrics(name))
          result
         rescue => e
          OneApm::Manager.logger.error "Thrift receive_message error: #{e}"
        ensure
          rpc_reset!
          result
        end
      end
      alias :send_oneway_message_without_oneapm :send_oneway_message
      alias :send_oneway_message :send_oneway_message_with_oneapm

      def receive_message_with_oneapm(result_klass)
        begin
          op = operator(result_klass)
          op_started = started_time(op)
          result = receive_message_without_oneapm(result_klass)

          state = OneApm::TransactionState.tl_get
          
          OneApm::Agent::CrossAppTracingMessage.finish_trace(state, op_started, segment, rpc_request, metrics(op))
          result
        rescue => e
          OneApm::Manager.logger.error "Thrift receive_message error: #{e}"
        ensure
          rpc_reset!
          result
        end
      end

      alias :receive_message_without_oneapm :receive_message
      alias :receive_message :receive_message_with_oneapm
    end
  end
  executes do
    require 'one_apm/support/http_clients/thrift_tracer'
    [::Thrift::BinaryProtocol, ::Thrift::CompactProtocol, ::Thrift::JsonProtocol].each do |protocol_classs|
      protocol_classs.send :include, OneApm::Support::HTTPClients::ThriftTracer
    end
  end
end