#!/usr/bin/ruby

require "pathname"
require "optparse"
require "parser/current"

# Always use the rubygems version of require
require "rubygems"

$options = { o: STDOUT, l: [] }

if method(:require) == method(:gem_original_require)
	# Bundler
	module Kernel
		def require(path)
			$LOAD_PATH.resolve_feature_path(path).last
		rescue LoadError => e
			raise e unless $options[:w]
			warn e
			path
		end
	end
else
	# Rubygems
	module Kernel
		def gem_original_require(path)
			$LOAD_PATH.resolve_feature_path(path).last
		rescue LoadError => e
			raise e unless $options[:w]
			warn e
			path
		end
	end
end

# opt-in to most recent AST format:
Parser::Builders::Default.emit_lambda              = true
Parser::Builders::Default.emit_procarg0            = true
Parser::Builders::Default.emit_encoding            = true
Parser::Builders::Default.emit_index               = true
Parser::Builders::Default.emit_arg_inside_procarg0 = true
Parser::Builders::Default.emit_forward_arg         = true
Parser::Builders::Default.emit_kwargs              = true
Parser::Builders::Default.emit_match_pattern       = true

op = OptionParser.new do |opts|
	opts.banner = "Usage: rld [options] file.rb ..."

	opts.accept(Pathname) do |path|
		Pathname.new(path)
	end

	opts.on("-p PATH", Pathname, "Project root") do |v|
		$options[:p] = v.realdirpath
	end

	opts.on("-L PATH", "-I PATH", Pathname, "Search directory") do |v|
		$LOAD_PATH.unshift(v)
	end

	opts.on("-o FILE", Pathname, "Output file") do |v|
		$options[:o] = v
	end

	opts.on("-l FILE", "-r FILE", Pathname, "Require this file in output") do |v|
		$options[:l] << "require #{require(v).inspect}"
	end

	opts.on("-w", "Only warn on LoadError") do
		$options[:w] = true
	end
end
op.parse!

if ARGV.empty?
	puts op
	exit 1
end

class Linker < Parser::TreeRewriter
	def on_send(node)
		super
		case node
		in [:send, [:gvar, :$: | :$LOAD_PATH], (:<< | :push | :unshift) => m, arg]
			loc = arg.location.expression
			file = Pathname.new(loc.source_buffer.name)
			path = eval(loc.source, binding, file.to_s, loc.line)
			$LOAD_PATH.public_send(m, Pathname.new(path).relative? ? "#{file.dirname}/#{path}" : path)
			remove(node.location.expression)
		in [:send, nil, :require, [:str, arg] => argnode]
			path = require(arg)
			realpath = Pathname.new(path).realdirpath
			if path.start_with?($options[:p].to_s)
				path = Pathname.new(path).relative_path_from($options[:p]).to_s
				replace(node.location.expression, "require_relative " + path.inspect)
			elsif realpath.to_s.start_with?($options[:p].to_s)
				path = realpath.relative_path_from($options[:p]).to_s
				replace(node.location.expression, "require_relative " + path.inspect)
			else
				replace(argnode.location.expression, path.inspect)
			end
		in [:send, nil, :load, [:str, arg] => argnode]
			return if arg.match?(/\A\.\.?\//)
			path = require(arg)
			replace(argnode.location.expression, path.inspect)
		in [:send, nil, :require | :load, *args]
			loc = node.location.expression
			warn "#{loc}: WARNING: ignoring non-string-literal #{loc.source}"
		else
		end
	end
end

$options[:p] ||= Pathname.new(ARGV[0]).dirname.realdirpath
$LOAD_PATH.unshift($options[:p])

$options[:l].each(&$options[:o].method(:puts))

ARGV.each do |file|
	buffer = Parser::Source::Buffer.new(file, source: Pathname.new(file).read)
	ast = Parser::CurrentRuby.new.parse(buffer)
	$options[:o].write Linker.new.rewrite(buffer, ast)
end