# encoding: utf-8
require 'one_apm/agent/database/active_record_helper'
require 'one_apm/inst/support/evented_subscriber'

# Listen for ActiveSupport::Notifications events for ActiveRecord query
# events.  Write metric data, transaction trace segments and slow sql
# nodes for each event.
module OneApm
  module Agent
    module Instrumentation
      class ActiveRecordSubscriber < EventedSubscriber
        OA_CACHED_QUERY_NAME = 'CACHE'.freeze unless defined? OA_CACHED_QUERY_NAME

        def start(name, id, payload)
          return if payload[:name] == OA_CACHED_QUERY_NAME
          return unless OneApm::Manager.tl_is_execution_traced?
          super
        rescue => e
          log_notification_error(e, name, 'start')
        end

        def finish(name, id, payload)
          return if payload[:name] == OA_CACHED_QUERY_NAME
          state = OneApm::TransactionState.tl_get
          return unless state.is_execution_traced?
          event = pop_event(id)
          record_metrics(event)
          notice_sql(state, event)
        rescue => e
          log_notification_error(e, name, 'finish')
        end

        def get_explain_plan( config, query )
          connection = OneApm::Agent::Database.get_connection(config) do
            ::ActiveRecord::Base.send("#{config[:adapter]}_connection",
                                      config)
          end
          if connection && connection.respond_to?(:execute)
            return connection.execute("EXPLAIN #{query}")
          end
        end

        def notice_sql(state, event)
          stack  = state.traced_method_stack
          config = active_record_config_for_event(event)

          base, *other_metrics = generate_metrics(event)

          # enter transaction trace segment
          frame = stack.push_frame(state, :active_record, event.time)

          OneApm::Manager.agent.transaction_sampler \
            .notice_sql(event.payload[:sql], config,
                        Helper.milliseconds_to_seconds(event.duration),
                        state, &method(:get_explain_plan))

          OneApm::Manager.agent.sql_sampler \
            .notice_sql(event.payload[:sql], base, config,
                        Helper.milliseconds_to_seconds(event.duration),
                        state, &method(:get_explain_plan))

          # exit transaction trace segment
          stack.pop_frame(state, frame, base, event.end)
        end

        def generate_metrics(event)
          @config  ||= active_record_config_for_event(event)
          @product ||= ActiveRecordHelper.database_info(@config)
          operation = operation(event)
          model_name = ActiveRecordHelper.model_for_name(event.payload[:name])
          ActiveRecordHelper.metric_for(@product, operation, model_name)
        end

        def record_metrics(event)
          base, *other_metrics = generate_metrics(event)

          OneApm::Manager.agent.stats_engine.tl_record_scoped_and_unscoped_metrics(
            base, other_metrics,
            Helper.milliseconds_to_seconds(event.duration)
          )
        end

        def operation(event)
          ActiveRecordHelper.operator_for_name(event.payload[:name]) ||
            ActiveRecordHelper.operator_for_sql(OneApm::Helper.correctly_encoded(event.payload[:sql]))
        end

        def active_record_config_for_event(event)
          return unless event.payload[:connection_id]

          connections = ::ActiveRecord::Base.connection_handler.connection_pool_list.map { |handler| handler.connections }.flatten
          connection = connections.detect { |cnxn| cnxn.object_id == event.payload[:connection_id] }

          connection.instance_variable_get(:@config) if connection
        end
      end
    end
  end
end