module ActiveRecord::Calculations

def aggregate_column(column_name)
  return column_name if Arel::Expressions === column_name
  arel_column(column_name.to_s) do |name|
    Arel.sql(column_name == :all ? "*" : name)
  end
end

def all_attributes?(column_names)
  (column_names.map(&:to_s) - @klass.attribute_names - @klass.attribute_aliases.keys).empty?
end

def average(column_name)
  calculate(:average, column_name)
end

def build_count_subquery(relation, column_name, distinct)
  if column_name == :all
    column_alias = Arel.star
    relation.select_values = [ Arel.sql(FinderMethods::ONE_AS_ONE) ] unless distinct
  else
    column_alias = Arel.sql("count_column")
    relation.select_values = [ aggregate_column(column_name).as(column_alias) ]
  end
  subquery_alias = Arel.sql("subquery_for_count")
  select_value = operation_over_aggregate_column(column_alias, "count", false)
  relation.build_subquery(subquery_alias, select_value)
end

def calculate(operation, column_name)
  if has_include?(column_name)
    relation = apply_join_dependency
    if operation.to_s.downcase == "count"
      unless distinct_value || distinct_select?(column_name || select_for_count)
        relation.distinct!
        relation.select_values = [ klass.primary_key || table[Arel.star] ]
      end
      # PostgreSQL: ORDER BY expressions must appear in SELECT list when using DISTINCT
      relation.order_values = [] if group_values.empty?
    end
    relation.calculate(operation, column_name)
  else
    perform_calculation(operation, column_name)
  end
end

def count(column_name = nil)
  if block_given?
    unless column_name.nil?
      raise ArgumentError, "Column name argument is not supported when a block is passed."
    end
    super()
  else
    calculate(:count, column_name)
  end
end

def distinct_select?(column_name)
  column_name.is_a?(::String) && /\bDISTINCT[\s(]/i.match?(column_name)
end

def execute_grouped_calculation(operation, column_name, distinct) # :nodoc:
  group_fields = group_values
  group_fields = group_fields.uniq if group_fields.size > 1
  if group_fields.size == 1 && group_fields.first.respond_to?(:to_sym)
    association  = klass._reflect_on_association(group_fields.first)
    associated   = association && association.belongs_to? # only count belongs_to associations
    group_fields = Array(association.foreign_key) if associated
  end
  group_fields = arel_columns(group_fields)
  column_alias_tracker = ColumnAliasTracker.new(connection)
  group_aliases = group_fields.map { |field|
    field = connection.visitor.compile(field) if Arel.arel_node?(field)
    column_alias_tracker.alias_for(field.to_s.downcase)
  }
  group_columns = group_aliases.zip(group_fields)
  column = aggregate_column(column_name)
  column_alias = column_alias_tracker.alias_for("#{operation} #{column_name.to_s.downcase}")
  select_value = operation_over_aggregate_column(column, operation, distinct)
  select_value.as(connection.quote_column_name(column_alias))
  select_values = [select_value]
  select_values += self.select_values unless having_clause.empty?
  select_values.concat group_columns.map { |aliaz, field|
    aliaz = connection.quote_column_name(aliaz)
    if field.respond_to?(:as)
      field.as(aliaz)
    else
      "#{field} AS #{aliaz}"
    end
  }
  relation = except(:group).distinct!(false)
  relation.group_values  = group_fields
  relation.select_values = select_values
  calculated_data = skip_query_cache_if_necessary { @klass.connection.select_all(relation.arel, "#{@klass.name} #{operation.capitalize}") }
  if association
    key_ids     = calculated_data.collect { |row| row[group_aliases.first] }
    key_records = association.klass.base_class.where(association.klass.base_class.primary_key => key_ids)
    key_records = key_records.index_by(&:id)
  end
  key_types = group_columns.each_with_object({}) do |(aliaz, col_name), types|
    types[aliaz] = col_name.try(:type_caster) ||
      type_for(col_name) do
        calculated_data.column_types.fetch(aliaz, Type.default_value)
      end
  end
  hash_rows = calculated_data.cast_values(key_types).map! do |row|
    calculated_data.columns.each_with_object({}).with_index do |(col_name, hash), i|
      hash[col_name] = row[i]
    end
  end
  if operation != "count"
    type = column.try(:type_caster) ||
      lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
    type = type.subtype if Enum::EnumType === type
  end
  hash_rows.each_with_object({}) do |row, result|
    key = group_aliases.map { |aliaz| row[aliaz] }
    key = key.first if key.size == 1
    key = key_records[key] if associated
    result[key] = type_cast_calculated_value(row[column_alias], operation, type)
  end
end

def execute_simple_calculation(operation, column_name, distinct) # :nodoc:
  if operation == "count" && (column_name == :all && distinct || has_limit_or_offset?)
    # Shortcut when limit is zero.
    return 0 if limit_value == 0
    query_builder = build_count_subquery(spawn, column_name, distinct)
  else
    # PostgreSQL doesn't like ORDER BY when there are no GROUP BY
    relation = unscope(:order).distinct!(false)
    column = aggregate_column(column_name)
    select_value = operation_over_aggregate_column(column, operation, distinct)
    select_value.distinct = true if operation == "sum" && distinct
    relation.select_values = [select_value]
    query_builder = relation.arel
  end
  result = skip_query_cache_if_necessary { @klass.connection.select_all(query_builder, "#{@klass.name} #{operation.capitalize}") }
  if operation != "count"
    type = column.try(:type_caster) ||
      lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
    type = type.subtype if Enum::EnumType === type
  end
  type_cast_calculated_value(result.cast_values.first, operation, type)
end

def has_include?(column_name)
  eager_loading? || (includes_values.present? && column_name && column_name != :all)
end

def ids
  pluck primary_key
end

def lookup_cast_type_from_join_dependencies(name, join_dependencies = build_join_dependencies)
  each_join_dependencies(join_dependencies) do |join|
    type = join.base_klass.attribute_types.fetch(name, nil)
    return type if type
  end
  nil
end

def maximum(column_name)
  calculate(:maximum, column_name)
end

def minimum(column_name)
  calculate(:minimum, column_name)
end

def operation_over_aggregate_column(column, operation, distinct)
  operation == "count" ? column.count(distinct) : column.public_send(operation)
end

def perform_calculation(operation, column_name)
  operation = operation.to_s.downcase
  # If #count is used with #distinct (i.e. `relation.distinct.count`) it is
  # considered distinct.
  distinct = distinct_value
  if operation == "count"
    column_name ||= select_for_count
    if column_name == :all
      if !distinct
        distinct = distinct_select?(select_for_count) if group_values.empty?
      elsif group_values.any? || select_values.empty? && order_values.empty?
        column_name = primary_key
      end
    elsif distinct_select?(column_name)
      distinct = nil
    end
  end
  if group_values.any?
    execute_grouped_calculation(operation, column_name, distinct)
  else
    execute_simple_calculation(operation, column_name, distinct)
  end
end

def pick(*column_names)
  if loaded? && all_attributes?(column_names)
    return records.pick(*column_names)
  end
  limit(1).pluck(*column_names).first
end

def pluck(*column_names)
  if loaded? && all_attributes?(column_names)
    return records.pluck(*column_names)
  end
  if has_include?(column_names.first)
    relation = apply_join_dependency
    relation.pluck(*column_names)
  else
    klass.disallow_raw_sql!(column_names)
    columns = arel_columns(column_names)
    relation = spawn
    relation.select_values = columns
    result = skip_query_cache_if_necessary do
      if where_clause.contradiction?
        ActiveRecord::Result.empty
      else
        klass.connection.select_all(relation.arel, "#{klass.name} Pluck")
      end
    end
    type_cast_pluck_values(result, columns)
  end
end

def select_for_count
  if select_values.present?
    return select_values.first if select_values.one?
    select_values.join(", ")
  else
    :all
  end
end

def sum(identity_or_column = nil, &block)
  if block_given?
    values = map(&block)
    if identity_or_column.nil? && (values.first.is_a?(Numeric) || values.first(1) == [] || values.first.respond_to?(:coerce))
      identity_or_column = 0
    end
    if identity_or_column.nil?
      ActiveSupport::Deprecation.warn(<<-MSG.squish)
        Rails 7.0 has deprecated Enumerable.sum in favor of Ruby's native implementation available since 2.4.
        Sum of non-numeric elements requires an initial argument.
      MSG
      values.inject(:+) || 0
    else
      values.sum(identity_or_column)
    end
  else
    calculate(:sum, identity_or_column)
  end
end

def type_cast_calculated_value(value, operation, type)
  case operation
  when "count"
    value.to_i
  when "sum"
    type.deserialize(value || 0)
  when "average"
    case type.type
    when :integer, :decimal
      value&.to_d
    else
      type.deserialize(value)
    end
  else # "minimum", "maximum"
    type.deserialize(value)
  end
end

def type_cast_pluck_values(result, columns)
  cast_types = if result.columns.size != columns.size
    klass.attribute_types
  else
    join_dependencies = nil
    columns.map.with_index do |column, i|
      column.try(:type_caster) ||
        klass.attribute_types.fetch(name = result.columns[i]) do
          join_dependencies ||= build_join_dependencies
          lookup_cast_type_from_join_dependencies(name, join_dependencies) ||
            result.column_types[name] || Type.default_value
        end
    end
  end
  result.cast_values(cast_types)
end

def type_for(field, &block)
  field_name = field.respond_to?(:name) ? field.name.to_s : field.to_s.split(".").last
  @klass.type_for_attribute(field_name, &block)
end