# encoding: utf-8

module OneApm
  module Agent
    module  Instrumentation
      module ThriftHelper

        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
                      ["External/#{thost}/Thrift/#{operator_name}", "External/#{thost}/all"]
                    end
          metrics << "External/all"
          if OneApm::Transaction.recording_web_transaction?
            metrics << "External/allWeb"
          else
            metrics << "External/allOther"
          end
          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
    ::Thrift::Client.module_eval do 
      include OneApm::Agent::Instrumentation::ThriftHelper

      def send_message_with_oneapm(name, args_class, args = {})
        operations[name] = {:started_time => Time.now.to_f}
        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 = {})
        op_started = Time.now.to_f
        base, *other_metrics = metrics(name)
        result = send_oneway_message_without_oneapm(name, args_class, args)
        duration = Time.now.to_f - op_started
        OneApm::Manager.agent.stats_engine.tl_record_scoped_and_unscoped_metrics(base, other_metrics, duration)
        result
      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)
        op = operator(result_klass)
        op_started = started_time(op)
        base, *other_metrics = metrics(op)
        result = receive_message_without_oneapm(result_klass)
        duration = Time.now.to_f - op_started
        OneApm::Manager.agent.stats_engine.tl_record_scoped_and_unscoped_metrics(base, other_metrics, duration)
        result
      end

      alias :receive_message_without_oneapm :receive_message
      alias :receive_message :receive_message_with_oneapm
    end
  end
end