require 'benchmark'
require 'active_record'
require 'action_controller'
require 'sql_logging/logged_query'

class ActiveRecord::ConnectionAdapters::AbstractAdapter
  @@stats_queries = @@stats_bytes = @@stats_rows = @@stats_execs = 0
  @@top_sql_queries = {}

  def self.get_stats
    { :queries => @@stats_queries,
      :rows => @@stats_rows,
      :bytes => @@stats_bytes,
      :execs => @@stats_execs }
  end
  
  def self.update_stats(values)
    return if values[:name].blank? || values[:name] =~ / Columns$/
    
    @@stats_queries += values[:queries] || 0
    @@stats_rows += values[:rows] || 0
    @@stats_bytes += values[:bytes] || 0
    @@stats_execs += values[:execs] || 0

    unless ActionController::Base.show_top_sql_queries == false
      backtrace = values[:backtrace].join("\n    ")
      key = "#{values[:name]}:#{backtrace}"
      unless query = @@top_sql_queries[key]
        query = LoggedQuery.new(values[:sql], values[:name], backtrace)
        @@top_sql_queries[key] = query
      end
      query.log_query(values[:rows] || 0, values[:bytes] || 0, values[:msec])
    end
  end
  
  def self.get_top_queries
    @@top_sql_queries
  end

  def self.reset_stats
    @@stats_queries = @@stats_bytes = @@stats_rows = @@stats_execs = 0
    @@top_sql_queries = {}
  end

  def record_query_stats(sql, name, msec, rows)
    bytes = 0
    rows.each do |row|
      row.each do |key, value|
        bytes += key.length if key
        bytes += value.length if value
      end
    end
    ntuples = if rows.respond_to?(:length)
      rows.length
    elsif rows.respond_to?(:ntuples)
      rows.ntuples
    end
    self.class.update_stats(:name => name, :sql => sql, :backtrace => clean_backtrace, :queries => 1, :rows => ntuples, :bytes => bytes, :msec => msec)
    @logger.info sprintf("%d rows, %.1f KB in %dms", ntuples, bytes.to_f / 1.kilobytes, msec)
  end

  def execute_with_stats(sql, name = nil)
    res = nil
    elapsed = Benchmark.measure do
      res = execute_without_stats(sql, name)
    end
    msec = (elapsed.real * 1000).to_i
    if res && res.respond_to?(:result)
      record_query_stats(sql, msec, res.result)
    else
      self.class.update_stats(:name => name, :sql => sql, :backtrace => clean_backtrace, :execs => 1, :msec => msec)
    end
    res
  end
end

class ActionController::Base
  @@show_top_sql_queries = :total_time
  @@top_sql_queries_shown = 10

  def self.show_top_sql_queries
    @@show_top_sql_queries
  end
  
  def self.show_top_sql_queries=(value)
    unless [ false, :rows, :queries, :bytes, :total_time, :median_time ].include?(value)
      raise ArgumentError, "show_top_sql_queries must be one of false, :rows, :queries, :bytes, :total_time or :median_time"
    end
    
    @@show_top_sql_queries = value
  end

  cattr_accessor :top_sql_queries_shown
  
  def perform_action_with_db_stats_reset
    ActiveRecord::ConnectionAdapters::AbstractAdapter.reset_stats
    perform_action_without_db_stats_reset
  end

  alias_method_chain :perform_action, :db_stats_reset

  def active_record_runtime(*args)
    stats = ActiveRecord::ConnectionAdapters::AbstractAdapter.get_stats
    runtime = "#{super}, #{stats[:queries]} selects, #{sprintf("%.1f KB", stats[:bytes].to_f / 1.kilobytes)}, #{stats[:execs]} execs"
    
    unless @@show_top_sql_queries == false
      runtime << "\nTop #{@@top_sql_queries_shown} SQL executions:"
      queries = ActiveRecord::ConnectionAdapters::AbstractAdapter.get_top_queries
      sorted_keys = queries.keys.sort_by { |k| queries[k][@@show_top_sql_queries] }.reverse
      sorted_keys.slice(0...@@top_sql_queries_shown).each do |key|
        runtime << "\n  Executed #{queries[key].queries} times in #{queries[key].total_time}ms " +
          "(#{queries[key].min_time}/#{queries[key].median_time}/#{queries[key].max_time}ms min/median/max), " +
          "returning #{queries[key][:rows]} rows " +
          "(#{sprintf("%.1f KB", queries[key].bytes.to_f / 1.kilobytes)}):\n" +
          "    #{queries[key].name}\n" +
          "    First exec was: #{queries[key].sql}\n" +
          "    #{queries[key].backtrace}"
      end
    end
    
    runtime
  end
end

begin
  require "sql_logging/#{ActiveRecord::Base.connection.adapter_name.downcase}"
rescue LoadError => e
  ActiveRecord::Base.logger.warn "SQL logging extensions are not available for #{ActiveRecord::Base.connection.adapter_name}, functionality will be limited"
end