# encoding: utf-8

LibraryDetection.defer do
  named :mongo2

  depends_on do
    require 'one_apm/agent/datastore/mongo'
    defined?(::Mongo) &&  OneApm::Agent::Datastore::Mongo.is_version2? && !OneApm::Manager.config[:disable_mongo]
  end

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

  def install_mongo_instrumentation
    require 'one_apm/agent/datastore/mongo/metric_translator'
    require 'one_apm/agent/datastore/mongo/statement_formatter'
    hook_instrument_method_for_collection
    hook_instrument_method_for_view
  end

  def hook_instrument_method_for_collection
    methods = [:create,:drop, :find, :indexes, :insert_one, :insert_many, :bulk_write, :parallel_scan]
    instrument_methods_with_oneapm(::Mongo::Collection, methods)
  end


  def hook_instrument_method_for_view
    methods = [:count, :distinct, :map_reduce,:find_one_and_delete, :find_one_and_replace, :find_one_and_update, 
      :delete_many, :delete_one, :replace_one, :update_many, :update_one]
    instrument_methods_with_oneapm(::Mongo::Collection::View, methods)
  end


  def instrument_methods_with_oneapm klass, methods
    klass.class_eval do 
      include OneApm::Support::MethodTracer

      def payload_for_oneapm
        collection = self.name rescue self.collection.name
        @payload_for_oneapm ||= { :collection => collection, :database => self.database.name }
      end

      def one_apm_notice_statement(t0, name)
        statement = OneApm::Agent::Datastore::Mongo::StatementFormatter.format(payload_for_oneapm, name)
        if statement
          OneApm::Manager.agent.transaction_sampler.notice_nosql_statement(statement, (Time.now - t0).to_f)
        end
      rescue => e
        OneApm::Manager.logger.debug("Exception during Mongo statement gathering", e)
      end

      def one_apm_generate_metrics(operation)
        OneApm::Agent::Datastore::Mongo::MetricTranslator.metrics_for(operation, payload_for_oneapm)
      end
    end

    methods.each do |method_name|
      next unless klass.method_defined?(method_name)

      klass.class_eval do 

        method_name_without_oneapm = :"#{method_name}_without_oneapm"
        method_name_with_oneapm = :"#{method_name}_with_oneapm"

        define_method method_name_with_oneapm do |*args, &block|
          metrics = one_apm_generate_metrics(method_name)

          trace_execution_scoped(metrics) do
            t0 = Time.now
            result = OneApm::Manager.disable_all_tracing do
              send method_name_without_oneapm, *args, &block
            end
            one_apm_notice_statement(t0, method_name)
            result
          end
         
        end
        alias_method method_name_without_oneapm, method_name.to_sym
        alias_method method_name.to_sym, method_name_with_oneapm
      end
    end

  end
 
end