pretty_print_sql.rb

require 'ostruct'
require 'anbt-sql-formatter/formatter'

module PrettyPrintSql
  module ActiveRecord
    def ppsql
      to_sql.tap {|sql| ::PrettyPrintSql.pp(sql) }
    end
  end

  module_function

  def pp(str, out = $>, theme = nil)
    out << pretty_print(str.dup, theme: theme, corlorize: out.tty?) << "\n"
  end

  def pretty_print(str, theme: nil, corlorize: false)
    @sql_formatter ||= sql_formatter
    sql = @sql_formatter.format(str)
    return sql unless corlorize && @terminal_formatter

    public_send(@terminal_formatter, sql, theme)
  end

  def config
    @config ||= OpenStruct.new().tap do |config|
      config.default_theme = 'github' if @terminal_formatter == :rouge
    end
  end

  def themes
    @terminal_formatter == :rouge ? Rouge::Theme.registry.keys : []
  end

  def load_terminal_formatter!(formatter = nil)
    @terminal_formatter = nil
    load_terminal_formatter(formatter)
  end

  private
  module_function

  def sql_formatter
    rule = AnbtSql::Rule.new
    rule.keyword = AnbtSql::Rule::KEYWORD_UPPER_CASE
    rule.kw_minus1_indent_nl_x_plus1_indent += ['INNER JOIN', 'LEFT OUTER JOIN']
    rule.kw_nl_x += ['ON']
    rule.kw_nl_x_plus1_indent -= ['ON']
    rule.kw_multi_words += ['INNER JOIN', 'LEFT OUTER JOIN']
    AnbtSql::Formatter.new(rule)
  end

  def coderay(sql, _theme)
    @duo ||= CodeRay::Duo[:sql, :terminal]
    @duo.encode(sql)
  end

  def rouge(sql, theme)
    @rouge_lexer ||= Rouge::Lexer.find('sql').new
    formatter =
      if theme.nil? || theme == config.default_theme
        @rouge_formatter ||= rouge_formatter(config.default_theme)
      else
        rouge_formatter(theme)
      end
    formatter.format(@rouge_lexer.lex(sql))
  end

  def rouge_formatter(theme)
    Rouge::Formatter.find('terminal256')
      .new(Rouge::Theme.find(theme).new)
  end

  def detect_loaded_terminal_formatter
    if defined?(::CodeRay)
      :coderay
    elsif defined?(::Rouge)
      :rouge
    else
      nil
    end
  end

  def load_terminal_formatter(formatter = nil)
    return if @terminal_formatter

    if formatter.nil?
      @terminal_formatter = detect_loaded_terminal_formatter
      return if @terminal_formatter
    end

    {
      rouge: 'rouge',
      coderay: 'coderay',
    }.each do |name, lib|
      next if name && name != formatter
      begin
        require lib
        @terminal_formatter = name
        break
      rescue LoadError
      end
    end
  end

  load_terminal_formatter
end

if defined?(::ActiveRecord)
  ::ActiveRecord::Relation.include(PrettyPrintSql::ActiveRecord)
end