module ActiveRecord::Calculations::ClassMethods

def average(column_name, options = {})
  calculate(:avg, column_name, options)
end

def calculate(operation, column_name, options = {})
  validate_calculation_options(operation, options)
  column_name     = options[:select] if options[:select]
  column_name     = '*' if column_name == :all
  column          = column_for column_name
  catch :invalid_query do
    if options[:group]
      return execute_grouped_calculation(operation, column_name, column, options)
    else
      return execute_simple_calculation(operation, column_name, column, options)
    end
  end
  0
end

def column_alias_for(*keys)
  table_name = keys.join(' ')
  table_name.downcase!
  table_name.gsub!(/\*/, 'all')
  table_name.gsub!(/\W+/, ' ')
  table_name.strip!
  table_name.gsub!(/ +/, '_')
  connection.table_alias_for(table_name)
end

def column_for(field)
  field_name = field.to_s.split('.').last
  columns.detect { |c| c.name.to_s == field_name }
end

def construct_calculation_sql(operation, column_name, options) #:nodoc:
  operation = operation.to_s.downcase
  options = options.symbolize_keys
  scope           = scope(:find)
  merged_includes = merge_includes(scope ? scope[:include] : [], options[:include])
  aggregate_alias = column_alias_for(operation, column_name)
  column_name     = "#{connection.quote_table_name(table_name)}.#{column_name}" if column_names.include?(column_name.to_s)
  if operation == 'count'
    if merged_includes.any?
      options[:distinct] = true
      column_name = options[:select] || [connection.quote_table_name(table_name), primary_key] * '.'
    end
    if options[:distinct]
      use_workaround = !connection.supports_count_distinct?
    end
  end
  if options[:distinct] && column_name.to_s !~ /\s*DISTINCT\s+/i
    distinct = 'DISTINCT ' 
  end
  sql = "SELECT #{operation}(#{distinct}#{column_name}) AS #{aggregate_alias}"
  # A (slower) workaround if we're using a backend, like sqlite, that doesn't support COUNT DISTINCT.
  sql = "SELECT COUNT(*) AS #{aggregate_alias}" if use_workaround
  options[:group_fields].each_index{|i| sql << ", #{options[:group_fields][i]} AS #{options[:group_aliases][i]}" } if options[:group]
  if options[:from]
    sql << " FROM #{options[:from]} "
  elsif scope && scope[:from] && !use_workaround
    sql << " FROM #{scope[:from]} "
  else
    sql << " FROM (SELECT #{distinct}#{column_name}" if use_workaround
    sql << " FROM #{connection.quote_table_name(table_name)} "
  end
  joins = ""
  add_joins!(joins, options[:joins], scope)
  if merged_includes.any?
    join_dependency = ActiveRecord::Associations::ClassMethods::JoinDependency.new(self, merged_includes, joins)
    sql << join_dependency.join_associations.collect{|join| join.association_join }.join
  end
  sql << joins unless joins.blank?
  add_conditions!(sql, options[:conditions], scope)
  add_limited_ids_condition!(sql, options, join_dependency) if join_dependency && !using_limitable_reflections?(join_dependency.reflections) && ((scope && scope[:limit]) || options[:limit])
  if options[:group]
    group_key = connection.adapter_name == 'FrontBase' ?  :group_aliases : :group_fields
    sql << " GROUP BY #{options[group_key].join(',')} "
  end
  if options[:group] && options[:having]
    having = sanitize_sql_for_conditions(options[:having])
    # FrontBase requires identifiers in the HAVING clause and chokes on function calls
    if connection.adapter_name == 'FrontBase'
      having.downcase!
      having.gsub!(/#{operation}\s*\(\s*#{column_name}\s*\)/, aggregate_alias)
    end
    sql << " HAVING #{having} "
  end
  sql << " ORDER BY #{options[:order]} "       if options[:order]
  add_limit!(sql, options, scope)
  sql << ") #{aggregate_alias}_subquery" if use_workaround
  sql
end

def construct_count_options_from_args(*args)
  options     = {}
  column_name = :all
  
  # We need to handle
  #   count()
  #   count(:column_name=:all)
  #   count(options={})
  #   count(column_name=:all, options={})
  case args.size
  when 1
    args[0].is_a?(Hash) ? options = args[0] : column_name = args[0]
  when 2
    column_name, options = args
  else
    raise ArgumentError, "Unexpected parameters passed to count(): #{args.inspect}"
  end if args.size > 0
  
  [column_name, options]
end

def count(*args)
  calculate(:count, *construct_count_options_from_args(*args))
end

def execute_grouped_calculation(operation, column_name, column, options) #:nodoc:
  group_attr    = options[:group]
  association   = reflect_on_association(group_attr.to_s.to_sym)
  associated    = association && association.macro == :belongs_to # only count belongs_to associations
  group_fields  = Array(associated ? association.primary_key_name : group_attr)
  group_aliases = []
  group_columns = {}
  
  group_fields.each do |field|
    group_aliases << column_alias_for(field)
    group_columns[column_alias_for(field)] = column_for(field)
  end
  sql             = construct_calculation_sql(operation, column_name, options.merge(:group_fields => group_fields, :group_aliases => group_aliases))
  calculated_data = connection.select_all(sql)
  aggregate_alias = column_alias_for(operation, column_name)
  if association
    key_ids     = calculated_data.collect { |row| row[group_aliases.first] }
    key_records = association.klass.base_class.find(key_ids)
    key_records = key_records.inject({}) { |hsh, r| hsh.merge(r.id => r) }
  end
  calculated_data.inject(ActiveSupport::OrderedHash.new) do |all, row|
    key   = group_aliases.map{|group_alias| type_cast_calculated_value(row[group_alias], group_columns[group_alias])}
    key   = key.first if key.size == 1
    key   = key_records[key] if associated
    value = row[aggregate_alias]
    all[key] = type_cast_calculated_value(value, column, operation)
    all
  end
end

def execute_simple_calculation(operation, column_name, column, options) #:nodoc:
  value = connection.select_value(construct_calculation_sql(operation, column_name, options))
  type_cast_calculated_value(value, column, operation)
end

def maximum(column_name, options = {})
  calculate(:max, column_name, options)
end

def minimum(column_name, options = {})
  calculate(:min, column_name, options)
end

def sum(column_name, options = {})
  calculate(:sum, column_name, options)
end

def type_cast_calculated_value(value, column, operation = nil)
  if value.is_a?(String) || value.nil?
    case operation.to_s.downcase
    when 'count' then value.to_i
    when 'sum'   then type_cast_using_column(value || '0', column)
    when 'avg' then value.try(:to_d)
    else type_cast_using_column(value, column)
    end
  else
    value
  end
end

def type_cast_using_column(value, column)
  column ? column.type_cast(value) : value
end

def validate_calculation_options(operation, options = {})
  options.assert_valid_keys(CALCULATIONS_OPTIONS)
end