# encoding: utf-8

module OneApm
  module Agent
    module Instrumentation
      class MongoCommandSubscriber

        OA_MONGODB = 'MongoDB'.freeze
        OA_STARTED = 'STARTED'.freeze
        OA_SUCCEEDED = 'SUCCEEDED'.freeze
        OA_FAILED = 'FAILED'.freeze

        def started(event)
          begin
            return unless OneApm::Manager.tl_is_execution_traced?
            operations[event.operation_id] = event
          rescue Exception => e
            log_operations_error(OA_STARTED, e)
          end
        end

        def succeeded(event)
          operator OA_SUCCEEDED, event
        end

        def failed(event)
          operator OA_FAILED, event
        end

        private

        def operations
          @operations ||= {}
        end

        def generate_metrics(event)
          @product ||= begin
            host = event.address.host rescue nil
            port = event.address.port rescue nil
            OneApm::Agent::Datastore.oneapm_product(OA_MONGODB, host, port)
          end
          OneApm::Agent::Datastore::MetricHelper.metrics_for(@product, event.command_name, event.command.values.first)
        end

        def operator(operator_type, event)
          begin
            state = OneApm::TransactionState.tl_get
            return unless state.is_execution_traced?
            stared_event = operations.delete(event.operation_id)
            one_apm_record_metrics(stared_event, event.duration)
            one_apm_notice_statement(stared_event, operator_type, event.duration)
            one_apm_notice_sql(state, stared_event, operator_type, event.duration)
          rescue => e
            log_operations_error(operator_type, e)
          end
        end

        def one_apm_record_metrics(event, duration)         
          base, *other_metrics = generate_metrics(event)
          OneApm::Manager.agent.stats_engine.tl_record_scoped_and_unscoped_metrics(
            base, other_metrics, duration
          )
        end

        def one_apm_notice_sql(state, event, status, duration)
          stack = state.traced_method_stack
          base, *other_metrics = generate_metrics(event)

          started_time = Time.now.to_f
          frame = stack.push_frame(state, :mongo_tracer, started_time - duration)

          builder = state.transaction_sample_builder
          format_sql = OneApm::Agent::Datastore::Mongo::CommandFormatter.format_sql(event, status)
          OneApm::Manager.agent.transaction_sampler.send(:notice_extra_data, builder, format_sql, duration, :sql)

          stack.pop_frame(state, frame, base, started_time)
        end

        def one_apm_notice_statement(event, status, duration)
          statement = OneApm::Agent::Datastore::Mongo::CommandFormatter.format(event, status)
          if statement
             OneApm::Manager.agent.transaction_sampler.notice_nosql_statement(statement, duration)
          end
          rescue => e
            OneApm::Manager.logger.debug("Exception during Mongo statement gathering", e)
        end

        def log_operations_error(event_type, error)
          OneApm::Manager.logger.error("Error during MongoDB #{event_type} event:")
          OneApm::Manager.logger.log_exception(:error, error)
        end
  
      end
    end
  end
end

LibraryDetection.defer do
  named :mongo2

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

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

  def install_mongo_command_subscriber
    require 'one_apm/agent/datastore/metric_helper'
    require 'one_apm/agent/datastore/mongo/command_formatter'
    ::Mongo::Monitoring::Global.subscribe(
      ::Mongo::Monitoring::COMMAND,
      OneApm::Agent::Instrumentation::MongoCommandSubscriber.new
    )
  end
end