# encoding: utf-8

require 'sequel' unless defined?( Sequel )
require 'oneapm_rpm' unless defined?( OneApm )
require 'one_apm/agent/database/active_record_helper'

module Sequel
  module OneApmInstrumentation
    include OneApm::Support::MethodTracer,
            OneApm::Agent::Instrumentation::ActiveRecordHelper

    # Instrument all queries that go through #execute_query.
    def log_yield(sql, args=nil) #THREAD_LOCAL_ACCESS
      state = OneApm::TransactionState.tl_get
      return super unless state.is_execution_traced?

      t0 = Time.now
      rval = super
      t1 = Time.now

      begin
        duration = t1 - t0
        record_metrics(sql, args, duration)
        notice_sql(state, sql, args, t0, t1)
      rescue => err
        OneApm::Agent.logger.debug "while recording metrics for Sequel", err
      end

      return rval
    end

    # Record metrics for the specified +sql+ and +args+ using the specified
    # +duration+.
    def record_metrics(sql, args, duration) #THREAD_LOCAL_ACCESS
      primary_metric = primary_metric_for(sql, args)
      engine         = OneApm::Agent.instance.stats_engine

      metrics = rollup_metrics_for(primary_metric)
      metrics << remote_service_metric(*self.opts.values_at(:adapter, :host)) if self.opts.key?(:adapter)

      engine.tl_record_scoped_and_unscoped_metrics(primary_metric, metrics, duration)
    end

    THREAD_SAFE_CONNECTION_POOL_CLASSES = [
      (defined?(::Sequel::ThreadedConnectionPool) && ::Sequel::ThreadedConnectionPool),
    ].compact.freeze

    # Record the given +sql+ within a new frame, using the given +start+ and
    # +finish+ times.
    def notice_sql(state, sql, args, start, finish)
      metric   = primary_metric_for(sql, args)
      agent    = OneApm::Agent.instance
      duration = finish - start
      stack    = state.traced_method_stack

      begin
        frame = stack.push_frame(state, :sequel, start)
        explainer = Proc.new do |*|
          if THREAD_SAFE_CONNECTION_POOL_CLASSES.include?(self.pool.class)
            self[ sql ].explain
          else
            OneApm::Agent.logger.log_once(:info, :sequel_explain_skipped, "Not running SQL explains because Sequel is not in recognized multi-threaded mode")
            nil
          end
        end
        agent.transaction_sampler.notice_sql(sql, self.opts, duration, state, &explainer)
        agent.sql_sampler.notice_sql(sql, metric, self.opts, duration, state, &explainer)
      ensure
        stack.pop_frame(state, frame, metric, finish)
      end
    end


    # Derive a primary database metric for the specified +sql+.
    def primary_metric_for(sql, _)
      return metric_for_sql(OneApm::Helper.correctly_encoded(sql))
    end

  end # module OneApmInstrumentation

  OneApm::Agent.logger.debug "Registering the :oneapm_instrumentation extension."
  Database.register_extension(:oneapm_instrumentation, OneApmInstrumentation)

end # module Sequel