# based on https://gist.github.com/mnutt/566725
require "active_support/core_ext/module/attr_internal"

module Searchkick
  module QueryWithInstrumentation
    def execute_search
      name = searchkick_klass ? "#{searchkick_klass.name} Search" : "Search"
      event = {
        name: name,
        query: params
      }
      ActiveSupport::Notifications.instrument("search.searchkick", event) do
        super
      end
    end
  end

  module IndexWithInstrumentation
    def store(record)
      event = {
        name: "#{record.searchkick_klass.name} Store",
        id: search_id(record)
      }
      if Searchkick.callbacks_value == :bulk
        super
      else
        ActiveSupport::Notifications.instrument("request.searchkick", event) do
          super
        end
      end
    end

    def remove(record)
      name = record && record.searchkick_klass ? "#{record.searchkick_klass.name} Remove" : "Remove"
      event = {
        name: name,
        id: search_id(record)
      }
      if Searchkick.callbacks_value == :bulk
        super
      else
        ActiveSupport::Notifications.instrument("request.searchkick", event) do
          super
        end
      end
    end

    def update_record(record, method_name)
      event = {
        name: "#{record.searchkick_klass.name} Update",
        id: search_id(record)
      }
      if Searchkick.callbacks_value == :bulk
        super
      else
        ActiveSupport::Notifications.instrument("request.searchkick", event) do
          super
        end
      end
    end

    def bulk_index(records)
      if records.any?
        event = {
          name: "#{records.first.searchkick_klass.name} Import",
          count: records.size
        }
        event[:id] = search_id(records.first) if records.size == 1
        if Searchkick.callbacks_value == :bulk
          super
        else
          ActiveSupport::Notifications.instrument("request.searchkick", event) do
            super
          end
        end
      end
    end
    alias_method :import, :bulk_index

    def bulk_update(records, *args)
      if records.any?
        event = {
          name: "#{records.first.searchkick_klass.name} Update",
          count: records.size
        }
        event[:id] = search_id(records.first) if records.size == 1
        if Searchkick.callbacks_value == :bulk
          super
        else
          ActiveSupport::Notifications.instrument("request.searchkick", event) do
            super
          end
        end
      end
    end

    def bulk_delete(records)
      if records.any?
        event = {
          name: "#{records.first.searchkick_klass.name} Delete",
          count: records.size
        }
        event[:id] = search_id(records.first) if records.size == 1
        if Searchkick.callbacks_value == :bulk
          super
        else
          ActiveSupport::Notifications.instrument("request.searchkick", event) do
            super
          end
        end
      end
    end
  end

  module IndexerWithInstrumentation
    def perform
      if Searchkick.callbacks_value == :bulk
        event = {
          name: "Bulk",
          count: queued_items.size
        }
        ActiveSupport::Notifications.instrument("request.searchkick", event) do
          super
        end
      else
        super
      end
    end
  end

  module SearchkickWithInstrumentation
    def multi_search(searches)
      event = {
        name: "Multi Search",
        body: searches.flat_map { |q| [q.params.except(:body).to_json, q.body.to_json] }.map { |v| "#{v}\n" }.join,
      }
      ActiveSupport::Notifications.instrument("multi_search.searchkick", event) do
        super
      end
    end
  end

  # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/log_subscriber.rb
  class LogSubscriber < ActiveSupport::LogSubscriber
    def self.runtime=(value)
      Thread.current[:searchkick_runtime] = value
    end

    def self.runtime
      Thread.current[:searchkick_runtime] ||= 0
    end

    def self.reset_runtime
      rt = runtime
      self.runtime = 0
      rt
    end

    def search(event)
      self.class.runtime += event.duration
      return unless logger.debug?

      payload = event.payload
      name = "#{payload[:name]} (#{event.duration.round(1)}ms)"

      index = payload[:query][:index].is_a?(Array) ? payload[:query][:index].join(",") : payload[:query][:index]
      type = payload[:query][:type]
      request_params = payload[:query].except(:index, :type, :body)

      params = []
      request_params.each do |k, v|
        params << "#{CGI.escape(k.to_s)}=#{CGI.escape(v.to_s)}"
      end

      debug "  #{color(name, YELLOW, true)}  #{index}#{type ? "/#{type.join(',')}" : ''}/_search#{params.any? ? '?' + params.join('&') : nil} #{payload[:query][:body].to_json}"
    end

    def request(event)
      self.class.runtime += event.duration
      return unless logger.debug?

      payload = event.payload
      name = "#{payload[:name]} (#{event.duration.round(1)}ms)"

      debug "  #{color(name, YELLOW, true)}  #{payload.except(:name).to_json}"
    end

    def multi_search(event)
      self.class.runtime += event.duration
      return unless logger.debug?

      payload = event.payload
      name = "#{payload[:name]} (#{event.duration.round(1)}ms)"

      debug "  #{color(name, YELLOW, true)}  _msearch #{payload[:body]}"
    end
  end

  # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/railties/controller_runtime.rb
  module ControllerRuntime
    extend ActiveSupport::Concern

    protected

    attr_internal :searchkick_runtime

    def process_action(action, *args)
      # We also need to reset the runtime before each action
      # because of queries in middleware or in cases we are streaming
      # and it won't be cleaned up by the method below.
      Searchkick::LogSubscriber.reset_runtime
      super
    end

    def cleanup_view_runtime
      searchkick_rt_before_render = Searchkick::LogSubscriber.reset_runtime
      runtime = super
      searchkick_rt_after_render = Searchkick::LogSubscriber.reset_runtime
      self.searchkick_runtime = searchkick_rt_before_render + searchkick_rt_after_render
      runtime - searchkick_rt_after_render
    end

    def append_info_to_payload(payload)
      super
      payload[:searchkick_runtime] = (searchkick_runtime || 0) + Searchkick::LogSubscriber.reset_runtime
    end

    module ClassMethods
      def log_process_action(payload)
        messages = super
        runtime = payload[:searchkick_runtime]
        messages << ("Searchkick: %.1fms" % runtime.to_f) if runtime.to_f > 0
        messages
      end
    end
  end
end

Searchkick::Query.prepend(Searchkick::QueryWithInstrumentation)
Searchkick::Index.prepend(Searchkick::IndexWithInstrumentation)
Searchkick::Indexer.prepend(Searchkick::IndexerWithInstrumentation)
Searchkick.singleton_class.prepend(Searchkick::SearchkickWithInstrumentation)
Searchkick::LogSubscriber.attach_to :searchkick
ActiveSupport.on_load(:action_controller) do
  include Searchkick::ControllerRuntime
end