# frozen_string_literal: false
require 'test/unit'
require 'tempfile'

class RubyVM
  module AbstractSyntaxTree
    class Node
      class CodePosition
        include Comparable
        attr_reader :lineno, :column
        def initialize(lineno, column)
          @lineno = lineno
          @column = column
        end

        def <=>(other)
          case
          when lineno < other.lineno
            -1
          when lineno == other.lineno
            column <=> other.column
          when lineno > other.lineno
            1
          end
        end
      end

      def beg_pos
        CodePosition.new(first_lineno, first_column)
      end

      def end_pos
        CodePosition.new(last_lineno, last_column)
      end

      alias to_s inspect
    end
  end
end

class TestAst < Test::Unit::TestCase
  class Helper
    attr_reader :errors

    def initialize(path, src: nil)
      @path = path
      @errors = []
      @debug = false
      @ast = RubyVM::AbstractSyntaxTree.parse(src) if src
    end

    def validate_range
      @errors = []
      validate_range0(ast)

      @errors.empty?
    end

    def validate_not_cared
      @errors = []
      validate_not_cared0(ast)

      @errors.empty?
    end

    def ast
      return @ast if defined?(@ast)
      @ast = RubyVM::AbstractSyntaxTree.parse_file(@path)
    end

    private

    def validate_range0(node)
      beg_pos, end_pos = node.beg_pos, node.end_pos
      children = node.children.grep(RubyVM::AbstractSyntaxTree::Node)

      return true if children.empty?
      # These NODE_D* has NODE_LIST as nd_next->nd_next whose last locations
      # we can not update when item is appended.
      return true if [:DSTR, :DXSTR, :DREGX, :DSYM].include? node.type

      min = children.map(&:beg_pos).min
      max = children.map(&:end_pos).max

      unless beg_pos <= min
        @errors << { type: :min_validation_error, min: min, beg_pos: beg_pos, node: node }
      end

      unless max <= end_pos
        @errors << { type: :max_validation_error, max: max, end_pos: end_pos, node: node }
      end

      p "#{node} => #{children}" if @debug

      children.each do |child|
        p child if @debug
        validate_range0(child)
      end
    end

    def validate_not_cared0(node)
      beg_pos, end_pos = node.beg_pos, node.end_pos
      children = node.children.grep(RubyVM::AbstractSyntaxTree::Node)

      @errors << { type: :first_lineno, node: node } if beg_pos.lineno == 0
      @errors << { type: :first_column, node: node } if beg_pos.column == -1
      @errors << { type: :last_lineno,  node: node } if end_pos.lineno == 0
      @errors << { type: :last_column,  node: node } if end_pos.column == -1

      children.each {|c| validate_not_cared0(c) }
    end
  end

  SRCDIR = File.expand_path("../../..", __FILE__)

  Dir.glob("test/**/*.rb", base: SRCDIR).each do |path|
    define_method("test_ranges:#{path}") do
      helper = Helper.new("#{SRCDIR}/#{path}")
      helper.validate_range

      assert_equal([], helper.errors)
    end
  end

  Dir.glob("test/**/*.rb", base: SRCDIR).each do |path|
    define_method("test_not_cared:#{path}") do
      helper = Helper.new("#{SRCDIR}/#{path}")
      helper.validate_not_cared

      assert_equal([], helper.errors)
    end
  end

  private def parse(src)
    EnvUtil.suppress_warning {
      RubyVM::AbstractSyntaxTree.parse(src)
    }
  end

  def test_allocate
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree::Node.allocate}
  end

  def test_parse_argument_error
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree.parse(0)}
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree.parse(nil)}
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree.parse(false)}
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree.parse(true)}
    assert_raise(TypeError) {RubyVM::AbstractSyntaxTree.parse(:foo)}
  end

  def test_column_with_long_heredoc_identifier
    term = "A"*257
    ast = parse("<<-#{term}\n""ddddddd\n#{term}\n")
    node = ast.children[2]
    assert_equal(:STR, node.type)
    assert_equal(0, node.first_column)
  end

  def test_column_of_heredoc
    node = parse("<<-SRC\nddddddd\nSRC\n").children[2]
    assert_equal(:STR, node.type)
    assert_equal(0, node.first_column)
    assert_equal(6, node.last_column)

    node = parse("<<SRC\nddddddd\nSRC\n").children[2]
    assert_equal(:STR, node.type)
    assert_equal(0, node.first_column)
    assert_equal(5, node.last_column)
  end

  def test_parse_raises_syntax_error
    assert_raise_with_message(SyntaxError, /\bend\b/) do
      RubyVM::AbstractSyntaxTree.parse("end")
    end
  end

  def test_parse_file_raises_syntax_error
    Tempfile.create(%w"test_ast .rb") do |f|
      f.puts "end"
      f.close
      assert_raise_with_message(SyntaxError, /\bend\b/) do
        RubyVM::AbstractSyntaxTree.parse_file(f.path)
      end
    end
  end

  def test_of_proc_and_method
    proc = Proc.new { 1 + 2 }
    method = self.method(__method__)

    node_proc = RubyVM::AbstractSyntaxTree.of(proc)
    node_method = RubyVM::AbstractSyntaxTree.of(method)

    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node_proc)
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node_method)

    Tempfile.create(%w"test_of .rb") do |tmp|
      tmp.print "#{<<-"begin;"}\n#{<<-'end;'}"
      begin;
        SCRIPT_LINES__ = {}
        assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(proc {|x| x}))
      end;
      tmp.close
      assert_separately(["-", tmp.path], "#{<<~"begin;"}\n#{<<~'end;'}")
      begin;
        load ARGV[0]
        assert_empty(SCRIPT_LINES__)
      end;
    end
  end

  def sample_backtrace_location
    [caller_locations(0).first, __LINE__]
  end

  def test_of_backtrace_location
    backtrace_location, lineno = sample_backtrace_location
    node = RubyVM::AbstractSyntaxTree.of(backtrace_location)
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node)
    assert_equal(lineno, node.first_lineno)
  end

  def test_of_error
    assert_raise(TypeError) { RubyVM::AbstractSyntaxTree.of("1 + 2") }
  end

  def test_of_proc_and_method_under_eval
    keep_script_lines_back = RubyVM.keep_script_lines
    RubyVM.keep_script_lines = false

    method = self.method(eval("def example_method_#{$$}; end"))
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = self.method(eval("def self.example_singleton_method_#{$$}; end"))
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = eval("proc{}")
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = self.method(eval("singleton_class.define_method(:example_define_method_#{$$}){}"))
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = self.method(eval("define_singleton_method(:example_dsm_#{$$}){}"))
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

    method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }

  ensure
    RubyVM.keep_script_lines = keep_script_lines_back
  end

  def test_of_proc_and_method_under_eval_with_keep_script_lines
    keep_script_lines_back = RubyVM.keep_script_lines
    RubyVM.keep_script_lines = true

    method = self.method(eval("def example_method_#{$$}_with_keep_script_lines; end"))
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = self.method(eval("def self.example_singleton_method_#{$$}_with_keep_script_lines; end"))
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = eval("proc{}")
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = self.method(eval("singleton_class.define_method(:example_define_method_#{$$}_with_keep_script_lines){}"))
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = self.method(eval("define_singleton_method(:example_dsm_#{$$}_with_keep_script_lines){}"))
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = eval("Class.new{def example_method_with_keep_script_lines; end}.instance_method(:example_method_with_keep_script_lines)")
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

    method = eval("Class.new{def example_method_with_keep_script_lines; end}.instance_method(:example_method_with_keep_script_lines)")
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))

  ensure
    RubyVM.keep_script_lines = keep_script_lines_back
  end

  def test_of_backtrace_location_under_eval
    keep_script_lines_back = RubyVM.keep_script_lines
    RubyVM.keep_script_lines = false

    m = Module.new do
      eval(<<-END, nil, __FILE__, __LINE__)
        def self.sample_backtrace_location
          caller_locations(0).first
        end
      END
    end
    backtrace_location = m.sample_backtrace_location
    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(backtrace_location) }

  ensure
    RubyVM.keep_script_lines = keep_script_lines_back
  end

  def test_of_backtrace_location_under_eval_with_keep_script_lines
    keep_script_lines_back = RubyVM.keep_script_lines
    RubyVM.keep_script_lines = true

    m = Module.new do
      eval(<<-END, nil, __FILE__, __LINE__)
        def self.sample_backtrace_location
          caller_locations(0).first
        end
      END
    end
    backtrace_location = m.sample_backtrace_location
    node = RubyVM::AbstractSyntaxTree.of(backtrace_location)
    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node)
    assert_equal(2, node.first_lineno)

  ensure
    RubyVM.keep_script_lines = keep_script_lines_back
  end

  def test_of_c_method
    c = Class.new { attr_reader :foo }
    assert_nil(RubyVM::AbstractSyntaxTree.of(c.instance_method(:foo)))
  end

  def test_scope_local_variables
    node = RubyVM::AbstractSyntaxTree.parse("_x = 0")
    lv, _, body = *node.children
    assert_equal([:_x], lv)
    assert_equal(:LASGN, body.type)
  end

  def test_call
    node = RubyVM::AbstractSyntaxTree.parse("nil.foo")
    _, _, body = *node.children
    assert_equal(:CALL, body.type)
    recv, mid, args = body.children
    assert_equal(:NIL, recv.type)
    assert_equal(:foo, mid)
    assert_nil(args)
  end

  def test_fcall
    node = RubyVM::AbstractSyntaxTree.parse("foo()")
    _, _, body = *node.children
    assert_equal(:FCALL, body.type)
    mid, args = body.children
    assert_equal(:foo, mid)
    assert_nil(args)
  end

  def test_vcall
    node = RubyVM::AbstractSyntaxTree.parse("foo")
    _, _, body = *node.children
    assert_equal(:VCALL, body.type)
    mid, args = body.children
    assert_equal(:foo, mid)
    assert_nil(args)
  end

  def test_defn
    node = RubyVM::AbstractSyntaxTree.parse("def a; end")
    _, _, body = *node.children
    assert_equal(:DEFN, body.type)
    mid, defn = body.children
    assert_equal(:a, mid)
    assert_equal(:SCOPE, defn.type)
    _, args, = defn.children
    assert_equal(:ARGS, args.type)
  end

  def test_defn_endless
    node = RubyVM::AbstractSyntaxTree.parse("def a = nil")
    _, _, body = *node.children
    assert_equal(:DEFN, body.type)
    mid, defn = body.children
    assert_equal(:a, mid)
    assert_equal(:SCOPE, defn.type)
    _, args, = defn.children
    assert_equal(:ARGS, args.type)
  end

  def test_defs
    node = RubyVM::AbstractSyntaxTree.parse("def a.b; end")
    _, _, body = *node.children
    assert_equal(:DEFS, body.type)
    recv, mid, defn = body.children
    assert_equal(:VCALL, recv.type)
    assert_equal(:b, mid)
    assert_equal(:SCOPE, defn.type)
    _, args, = defn.children
    assert_equal(:ARGS, args.type)
  end

  def test_defs_endless
    node = RubyVM::AbstractSyntaxTree.parse("def a.b = nil")
    _, _, body = *node.children
    assert_equal(:DEFS, body.type)
    recv, mid, defn = body.children
    assert_equal(:VCALL, recv.type)
    assert_equal(:b, mid)
    assert_equal(:SCOPE, defn.type)
    _, args, = defn.children
    assert_equal(:ARGS, args.type)
  end

  def test_dstr
    node = parse('"foo#{1}bar"')
    _, _, body = *node.children
    assert_equal(:DSTR, body.type)
    head, body = body.children
    assert_equal("foo", head)
    assert_equal(:EVSTR, body.type)
    body, = body.children
    assert_equal(:LIT, body.type)
    assert_equal([1], body.children)
  end

  def test_while
    node = RubyVM::AbstractSyntaxTree.parse('1 while qux')
    _, _, body = *node.children
    assert_equal(:WHILE, body.type)
    type1 = body.children[2]
    node = RubyVM::AbstractSyntaxTree.parse('begin 1 end while qux')
    _, _, body = *node.children
    assert_equal(:WHILE, body.type)
    type2 = body.children[2]
    assert_not_equal(type1, type2)
  end

  def test_until
    node = RubyVM::AbstractSyntaxTree.parse('1 until qux')
    _, _, body = *node.children
    assert_equal(:UNTIL, body.type)
    type1 = body.children[2]
    node = RubyVM::AbstractSyntaxTree.parse('begin 1 end until qux')
    _, _, body = *node.children
    assert_equal(:UNTIL, body.type)
    type2 = body.children[2]
    assert_not_equal(type1, type2)
  end

  def test_keyword_rest
    kwrest = lambda do |arg_str|
      node = RubyVM::AbstractSyntaxTree.parse("def a(#{arg_str}) end")
      node = node.children.last.children.last.children[1].children[-2]
      node ? node.children : node
    end

    assert_equal(nil, kwrest.call(''))
    assert_equal([nil], kwrest.call('**'))
    assert_equal(false, kwrest.call('**nil'))
    assert_equal([:a], kwrest.call('**a'))
  end

  def test_ranges_numbered_parameter
    helper = Helper.new(__FILE__, src: "1.times {_1}")
    helper.validate_range
    assert_equal([], helper.errors)
  end

  def test_op_asgn2
    node = RubyVM::AbstractSyntaxTree.parse("struct.field += foo")
    _, _, body = *node.children
    assert_equal(:OP_ASGN2, body.type)
    recv, _, mid, op, value = body.children
    assert_equal(:VCALL, recv.type)
    assert_equal(:field, mid)
    assert_equal(:+, op)
    assert_equal(:VCALL, value.type)
  end

  def test_args
    rest = 6
    node = RubyVM::AbstractSyntaxTree.parse("proc { |a| }")
    _, args = *node.children.last.children[1].children
    assert_equal(nil, args.children[rest])

    node = RubyVM::AbstractSyntaxTree.parse("proc { |a,| }")
    _, args = *node.children.last.children[1].children
    assert_equal(:NODE_SPECIAL_EXCESSIVE_COMMA, args.children[rest])

    node = RubyVM::AbstractSyntaxTree.parse("proc { |*a| }")
    _, args = *node.children.last.children[1].children
    assert_equal(:a, args.children[rest])
  end

  def test_keep_script_lines_for_parse
    node = RubyVM::AbstractSyntaxTree.parse(<<~END, keep_script_lines: true)
1.times do
  2.times do
  end
end
__END__
dummy
    END

    expected = [
      "1.times do\n",
      "  2.times do\n",
      "  end\n",
      "end\n",
      "__END__\n",
    ]
    assert_equal(expected, node.script_lines)

    expected =
      "1.times do\n" +
      "  2.times do\n" +
      "  end\n" +
      "end"
    assert_equal(expected, node.source)

    expected =
             "do\n" +
      "  2.times do\n" +
      "  end\n" +
      "end"
    assert_equal(expected, node.children.last.children.last.source)

    expected =
        "2.times do\n" +
      "  end"
    assert_equal(expected, node.children.last.children.last.children.last.source)
  end

  def test_keep_script_lines_for_of
    proc = Proc.new { 1 + 2 }
    method = self.method(__method__)

    node_proc = RubyVM::AbstractSyntaxTree.of(proc, keep_script_lines: true)
    node_method = RubyVM::AbstractSyntaxTree.of(method, keep_script_lines: true)

    assert_equal("{ 1 + 2 }", node_proc.source)
    assert_equal("def test_keep_script_lines_for_of\n", node_method.source.lines.first)
  end

  def test_e_option
    assert_in_out_err(["-e", "def foo; end; pp RubyVM::AbstractSyntaxTree.of(method(:foo)).type"],
                      "", [":SCOPE"], [])
  end
end