module Marginalia::ActiveRecordInstrumentation

def self.included(instrumented_class)
  instrumented_class.class_eval do
    if instrumented_class.method_defined?(:execute)
      alias_method :execute_without_marginalia, :execute
      alias_method :execute, :execute_with_marginalia
    end
    if instrumented_class.private_method_defined?(:execute_and_clear)
      alias_method :execute_and_clear_without_marginalia, :execute_and_clear
      alias_method :execute_and_clear, :execute_and_clear_with_marginalia
    else
      is_mysql2 = defined?(ActiveRecord::ConnectionAdapters::Mysql2Adapter) &&
        ActiveRecord::ConnectionAdapters::Mysql2Adapter == instrumented_class
      # Dont instrument exec_query on mysql2 as it calls execute internally
      unless is_mysql2
        if instrumented_class.method_defined?(:exec_query)
          alias_method :exec_query_without_marginalia, :exec_query
          alias_method :exec_query, :exec_query_with_marginalia
        end
      end
      is_postgres = defined?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) &&
        ActiveRecord::ConnectionAdapters::PostgreSQLAdapter == instrumented_class
      # Instrument exec_delete and exec_update since they don't call
      # execute internally
      if is_postgres
        if instrumented_class.method_defined?(:exec_delete)
          alias_method :exec_delete_without_marginalia, :exec_delete
          alias_method :exec_delete, :exec_delete_with_marginalia
        end
        if instrumented_class.method_defined?(:exec_update)
          alias_method :exec_update_without_marginalia, :exec_update
          alias_method :exec_update, :exec_update_with_marginalia
        end
      end
    end
  end
end

def annotate_sql(sql)
  Marginalia::Comment.update_adapter!(self)
  comment = Marginalia::Comment.construct_comment
  if comment.present? && !sql.include?(comment)
    sql = if Marginalia::Comment.prepend_comment
      "/*#{comment}*/ #{sql}"
    else
      "#{sql} /*#{comment}*/"
    end
  end
  inline_comment = Marginalia::Comment.construct_inline_comment
  if inline_comment.present? && !sql.include?(inline_comment)
    sql = if Marginalia::Comment.prepend_comment
      "/*#{inline_comment}*/ #{sql}"
    else
      "#{sql} /*#{inline_comment}*/"
    end
  end
  sql
end

def exec_delete_with_marginalia(sql, *args)
  exec_delete_without_marginalia(annotate_sql(sql), *args)
end

def exec_query_with_marginalia(sql, *args, **options)
  options[:prepare] ||= false
  exec_query_without_marginalia(annotate_sql(sql), *args, **options)
end

def exec_update_with_marginalia(sql, *args)
  exec_update_without_marginalia(annotate_sql(sql), *args)
end

def execute_and_clear_with_marginalia(sql, *args, &block)
  execute_and_clear_without_marginalia(annotate_sql(sql), *args, &block)
end

def execute_with_marginalia(sql, *args)
  execute_without_marginalia(annotate_sql(sql), *args)
end