#!/usr/bin/env jruby

require 'optparse'

opts = {}
options = {
  :print_source => false,
  :print_sexp   => false,
  :print_ast    => true,
  :print_ir     => false,
  :dot_format   => false
}

OptionParser.new do |opts|
  opts.banner = "Usage: #{$0} [options]"

  opts.on('-d', '--dot', 'Display as dot data') do |h|
    options[:dot_format] = true
  end
  
  opts.on('-h', '--help', 'Display this help') do |h|
    puts opts
    exit true
  end

  opts.on('-i', '--ir', 'Dump all IR passes without executing') do |h|
    options[:print_ir] = true
  end
  
  opts.on('-s', '--sexp', 'Display the S-Expression for the AST') do |t|
    options[:print_sexp] = true
  end
  
  opts.on('--source', 'Display the source') do |s|
    options[:print_source] = true
  end
  
  opts.on('--no-ast', 'Do not print out the AST for this (only useful with -s)') do |a|
    options[:print_ast] = false
  end
  
  opts.on('-e exp', '--expression') do |e|
    options[:expression] = e
  end
  
end.parse!

if ARGV.length > 1
  abort "You may only specify one script (see --help)"
elsif ARGV.length == 1
  if options[:expression]
    abort "-e and a script is not a valid combination (see --help)"
  end
  options[:expression] = File.read(ARGV.shift)
elsif ! options.has_key?(:expression)
  abort "No script specified (see --help)"
end

require 'jruby'

$indent_string = "  "

def indexes(string, lindex, rindex)
  lindex = string.index("(", lindex) if lindex != nil
  rindex = string.index(")", rindex) if rindex != nil
  return lindex, rindex
end

def indent(string)
  depth = -1

  lindex, rindex = indexes(string, 0, 0)

  while (lindex != nil || rindex != nil)
    if (lindex != nil && lindex < rindex)
      depth += 1
      string[lindex, 1] = "\n#{$indent_string * depth}"
    else
      depth -= 1
      string[rindex, 1] = "\n"
    end

    lindex, rindex = indexes(string, lindex, rindex)
  end
  string.gsub(/,\s*$/, '').squeeze("\n")
end

def dot_label(node)
  extra = case node
          when org.jruby.ast.StrNode then
            ": '#{node.value}'"
          when org.jruby.ast.FixnumNode then
            ": #{node.value}"
          when org.jruby.ast.FloatNode then
            ": #{node.value}"
          when org.jruby.ast.types.INameNode then
            ": #{node.name}"
          else
            ""
          end

  "#{short_name(node)}#{extra}"
end

def short_name(node)
  node.class.name.split("::")[-1].gsub("Node", "")
end

def dot_node_def(node)
  %Q{#{node.hash} [label="#{dot_label(node)}"];\n}
end

def dot(defs, parent)
  defs[parent.hash] = dot_node_def(parent)

  "".tap do |str|
    parent.child_nodes.each do |child|
      str