# encoding: utf-8

LibraryDetection.defer do
  named :mongo

  depends_on do
    defined?(::Mongo)
  end

  depends_on do
    require 'one_apm/agent/datastores/mongo'
    OneApm::Agent::Datastores::Mongo.is_supported_version?
  end

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

  def install_mongo_instrumentation
    require 'one_apm/agent/datastores/mongo/metric_translator'
    require 'one_apm/agent/datastores/mongo/statement_formatter'

    hook_instrument_methods
    instrument_save
    instrument_ensure_index
  end

  def hook_instrument_methods
    hook_instrument_method(::Mongo::Collection)
    hook_instrument_method(::Mongo::Connection)
    hook_instrument_method(::Mongo::Cursor)
    hook_instrument_method(::Mongo::CollectionWriter) if defined?(::Mongo::CollectionWriter)
  end

  def hook_instrument_method(target_class)
    target_class.class_eval do
      include OneApm::Support::MethodTracer

      # It's key that this method eats all exceptions, as it rests between the
      # Mongo operation the user called and us returning them the data. Be safe!
      def one_apm_notice_statement(t0, payload, name)
        statement = OneApm::Agent::Datastores::Mongo::StatementFormatter.format(payload, 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, payload = nil)
        payload ||= { :collection => self.name, :database => self.db.name }
        OneApm::Agent::Datastores::Mongo::MetricTranslator.metrics_for(operation, payload)
      end

      def instrument_with_one_apm_trace(name, payload = {}, &block)
        metrics = one_apm_generate_metrics(name, payload)

        trace_execution_scoped(metrics) do
          t0 = Time.now

          result = OneApm::Manager.disable_all_tracing do
            instrument_without_one_apm_trace(name, payload, &block)
          end

          one_apm_notice_statement(t0, payload, name)
          result
        end
      end

      alias_method :instrument_without_one_apm_trace, :instrument
      alias_method :instrument, :instrument_with_one_apm_trace
    end
  end

  def instrument_save
    ::Mongo::Collection.class_eval do
      def save_with_one_apm_trace(doc, opts = {}, &block)
        metrics = one_apm_generate_metrics(:save)
        trace_execution_scoped(metrics) do
          t0 = Time.now

          result = OneApm::Manager.disable_all_tracing do
            save_without_one_apm_trace(doc, opts, &block)
          end

          one_apm_notice_statement(t0, doc, :save)
          result
        end
      end

      alias_method :save_without_one_apm_trace, :save
      alias_method :save, :save_with_one_apm_trace
    end
  end

  def instrument_ensure_index
    ::Mongo::Collection.class_eval do
      def ensure_index_with_one_apm_trace(spec, opts = {}, &block)
        metrics = one_apm_generate_metrics(:ensureIndex)
        trace_execution_scoped(metrics) do
          t0 = Time.now

          result = OneApm::Manager.disable_all_tracing do
            ensure_index_without_one_apm_trace(spec, opts, &block)
          end

          spec = case spec
                 when Array
                   Hash[spec]
                 when String, Symbol
                   { spec => 1 }
                 else
                   spec.dup
                 end

          one_apm_notice_statement(t0, spec, :ensureIndex)
          result
        end
      end

      alias_method :ensure_index_without_one_apm_trace, :ensure_index
      alias_method :ensure_index, :ensure_index_with_one_apm_trace
    end
  end
end