#!/usr/bin/env ruby
require 'ruby-beautify'
require 'ripper'
require 'optparse'

Options = OptionParser.new do |opts|
  opts.on("-V", "--version", "Print version") { |version| puts RBeautify::VERSION;exit 0}
  opts.on("-c", "--indent_count [COUNT]", Integer, "Count of characters to use for indenting. (default: 1)") { |count| @indent_count = count}
  opts.on("-t", "--tabs", "Use tabs for the indent character (default)") { @indent_token = "\t" }
  opts.on("-s", "--spaces", "Use spaces for the indent character") { @indent_token = " " }
  opts.banner = "Usage: print ruby into a pretty format, or break trying."
end
Options.parse!

@open_block_start = ["module", "class", "begin", "def", 'if', 'while', 'unless']
@both_block = ["else", "elsif", 'rescue']
@open_block_do  = ['do', '{']
@close_block = ['end', '}']

@open_brackets  = [:on_lparen, :on_lbracket, :on_lbrace]
@close_brackets = [:on_rparen, :on_rbracket, :on_rbrace]
@indent_token = "\t" unless @indent_token
@indent_count = 1 unless @indent_count

def flatten_content(content)
  # I realize this is slow, todo, make it not slow.
  flat_content = []
  content.split("\n").each do |l|
    w = l.split(/\s*(?<line>.*)/).last
    if w
      flat_content.push w
    else
      flat_content.push l
    end
  end

  return flat_content.join("\n")
end

def puts_indented_line(level, string)
  if string.empty? or string =~ /^\n$/
    puts
  else
    indent = (@indent_token * @indent_count) * level
    puts indent + string
  end
end

def pretty(content)
  # it's far easier if the content is 'flat', we don't have to worry about
  # pre-existing indent levels.  So lets lex that.
  begin
    lex = Ripper.lex(flatten_content(content))
  rescue
    exit 255
  end

  # quick way to get our total lines.
  total_lines = lex.last[0][0]
  line_index = 1

  indent_level = 0

  # walk through our lines.
  while (line_index <= total_lines)
    # a 'lex' copy of our line.
    line_lex = lex.select {|l| l[0][0] == line_index}
    # a string version.
    line_string = line_lex.map {|l| l.last}.join("")

    # did we just close something?  if so, lets bring it down a level.
    if closing_block?(line_lex) || closing_assignment?(line_lex)
      indent_level -=1 if indent_level > 0
    end

    # print our line, in place.
    puts_indented_line(indent_level, line_string)

    # oh, we opened something did we?  lets indent for the next run.
    if opening_block?(line_lex) || opening_assignment?(line_lex)
      indent_level +=1
    end

    line_index += 1
  end

  return nil
end

# how many times do we open in this line?
def opening_block_count(line_lex)
  line_lex.select {|l| l[1] == :on_kw && @open_block_do.include?(l[2])}.count
end

# how many times do we close?
def closing_block_count(line_lex)
  line_lex.select {|l| l[1] == :on_kw && @close_block.include?(l[2])}.count
end

# is the first word a key word?
def starts_block?(line_lex)
  line_lex.first[1] == :on_kw && @open_block_start.include?(line_lex.first[2])
end

# is the first word one of our 'both' keywords?
def both_block?(line_lex)
  return line_lex.first[1] == :on_kw && @both_block.include?(line_lex.first[2])
end

# kinda complex, we count open/close to determine if we ultimately have a
# hanging line.   Always true if it's a both_block.
def opening_block?(line_lex)
  return true if both_block? line_lex
  opens = starts_block?(line_lex) ? 1 : 0
  opens += opening_block_count line_lex
  closes = closing_block_count line_lex
  return false if opens == closes
  return true if opens > closes
end

# kinda complex, we count open/close to determine if we ultimately have close a
# hanging line.   Always true if it's a both_block.
def closing_block?(line_lex)
  return true if both_block? line_lex
  opens = starts_block?(line_lex) ? 1 : 0
  opens += opening_block_count line_lex
  closes = closing_block_count line_lex
  return false if opens == closes
  return true if opens < closes
end

# count the amount of opening assignments.
def opening_assignment_count(line_lex)
  line_lex.select {|l| @open_brackets.include? l[1]}.count
end

# count the amount of closing assignments.
def closing_assignment_count(line_lex)
  line_lex.select {|l| @close_brackets.include? l[1]}.count
end

# same trick as opening_block
def opening_assignment?(line_lex)
  opens = opening_assignment_count line_lex
  closes = closing_assignment_count line_lex
  return false if opens == closes
  return true if opens > closes
end

# ...
def closing_assignment?(line_lex)
  opens = opening_assignment_count line_lex
  closes = closing_assignment_count line_lex
  return false if opens == closes
  return true if closes > opens
end

# is this ruby file valid syntax?
def valid_syntax?(file)
  #ugly but fast, ruby returns stdout if it's ok, and stderr if not.
  #if we get anything back, we were ok, else we were not.
  r = `ruby -c #{file} 2> /dev/null`
  return false if r == ""
  return true
end

# no argument, assume we want to open STDIN.
if ARGV.empty?
  content = $stdin.read
  puts pretty content
else
  file = ARGV[0]
  if File.exist? file
    if valid_syntax? file
      content = open(file).read
      pretty content
    else
      puts content
      exit 127
    end
  else
    puts "No such file: #{file}"
  end
end