module ActiveRecord::Batches

def act_on_ignored_order(error_on_ignore)
  raise_error = (error_on_ignore.nil? ? ActiveRecord.error_on_ignored_order : error_on_ignore)
  if raise_error
    raise ArgumentError.new(ORDER_IGNORE_MESSAGE)
  elsif logger
    logger.warn(ORDER_IGNORE_MESSAGE)
  end
end

def apply_finish_limit(relation, finish, order)
  relation.where(predicate_builder[primary_key, finish, order == :desc ? :gteq : :lteq])
end

def apply_limits(relation, start, finish, order)
  relation = apply_start_limit(relation, start, order) if start
  relation = apply_finish_limit(relation, finish, order) if finish
  relation
end

def apply_start_limit(relation, start, order)
  relation.where(predicate_builder[primary_key, start, order == :desc ? :lteq : :gteq])
end

def batch_order(order)
  table[primary_key].public_send(order)
end

def find_each(start: nil, finish: nil, batch_size: 1000, error_on_ignore: nil, order: :asc, &block)
  if block_given?
    find_in_batches(start: start, finish: finish, batch_size: batch_size, error_on_ignore: error_on_ignore, order: order) do |records|
      records.each(&block)
    end
  else
    enum_for(:find_each, start: start, finish: finish, batch_size: batch_size, error_on_ignore: error_on_ignore, order: order) do
      relation = self
      apply_limits(relation, start, finish, order).size
    end
  end
end

def find_in_batches(start: nil, finish: nil, batch_size: 1000, error_on_ignore: nil, order: :asc)
  relation = self
  unless block_given?
    return to_enum(:find_in_batches, start: start, finish: finish, batch_size: batch_size, error_on_ignore: error_on_ignore, order: order) do
      total = apply_limits(relation, start, finish, order).size
      (total - 1).div(batch_size) + 1
    end
  end
  in_batches(of: batch_size, start: start, finish: finish, load: true, error_on_ignore: error_on_ignore, order: order) do |batch|
    yield batch.to_a
  end
end

def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore: nil, order: :asc)
  relation = self
  unless block_given?
    return BatchEnumerator.new(of: of, start: start, finish: finish, relation: self)
  end
  unless [:asc, :desc].include?(order)
    raise ArgumentError, ":order must be :asc or :desc, got #{order.inspect}"
  end
  if arel.orders.present?
    act_on_ignored_order(error_on_ignore)
  end
  batch_limit = of
  if limit_value
    remaining   = limit_value
    batch_limit = remaining if remaining < batch_limit
  end
  relation = relation.reorder(batch_order(order)).limit(batch_limit)
  relation = apply_limits(relation, start, finish, order)
  relation.skip_query_cache! # Retaining the results in the query cache would undermine the point of batching
  batch_relation = relation
  loop do
    if load
      records = batch_relation.records
      ids = records.map(&:id)
      yielded_relation = where(primary_key => ids)
      yielded_relation.load_records(records)
    else
      ids = batch_relation.pluck(primary_key)
      yielded_relation = where(primary_key => ids)
    end
    break if ids.empty?
    primary_key_offset = ids.last
    raise ArgumentError.new("Primary key not included in the custom select clause") unless primary_key_offset
    yield yielded_relation
    break if ids.length < batch_limit
    if limit_value
      remaining -= ids.length
      if remaining == 0
        # Saves a useless iteration when the limit is a multiple of the
        # batch size.
        break
      elsif remaining < batch_limit
        relation = relation.limit(remaining)
      end
    end
    batch_relation = relation.where(
      predicate_builder[primary_key, primary_key_offset, order == :desc ? :lt : :gt]
    )
  end
end