#!/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 << "#{parent.hash} -> #{child.hash};\n"
      str << dot(defs, child)
    end
  end
end

def print_passes_on(scope, passes)
  if !scope.kind_of? org.jruby.ir.IRClosure
    passes.each { |pass| pass.run(scope) } 
  end

  scope.lexical_scopes.each do |child_scope|
    print_passes_on(child_scope, passes)
  end
end

root = JRuby.parse(options[:expression])

if options[:print_source]
  puts "Source:"
  puts options[:expression]
  puts
end

if options[:print_ast]
  if options[:dot_format]
    defs = {}
    graph_section = dot(defs, root)
    puts "digraph AST {"
    puts defs.values.join('')
    puts graph_section
    puts "}"
  else
    print "AST:"
    puts indent(root.to_string) 
    puts
  end
end

if options[:print_ir]
  require 'jruby'
  runtime = JRuby::runtime
  manager = runtime.ir_manager

  manager.dry_run = true
  JRuby::IR.compiler_debug = true


  builder = runtime.is1_9 ? org.jruby.ir.IRBuilder19 : org.jruby.ir.IRBuilder

  scope = builder.new(manager).build_root(root)
  passes = manager.get_compiler_passes(scope)

  print_passes_on(scope, passes)
end

if options[:print_sexp]
  puts "SEXP:"
  puts org.jruby.ast.util.SexpMaker.create(root) 
end