class RBI::Parser::TreeBuilder

def collect_dangling_comments(node)
  first_line = node.location.start_line
  last_line = node.location.end_line
  last_node_last_line = node.child_nodes.last&.location&.end_line
  first_line.upto(last_line) do |line|
    comment = @comments_by_line[line]
    next unless comment
    break if last_node_last_line && line <= last_node_last_line
    current_scope << parse_comment(comment)
    @comments_by_line.delete(line)
  end
end

def collect_orphan_comments
  last_line = nil #: Integer?
  last_node_end = @tree.nodes.last&.loc&.end_line
  @comments_by_line.each do |line, comment|
    # Associate the comment either with the header or the file or as a dangling comment at the end
    recv = if last_node_end && line >= last_node_end
      @tree
    else
      @tree.comments
    end
    # Preserve blank lines in comments
    if last_line && line > last_line + 1
      recv << BlankLine.new(loc: Loc.from_prism(@file, comment.location))
    end
    recv << parse_comment(comment)
    last_line = line
  end
end

def current_scope
  @scopes_stack.last #: as !nil # Should never be nil since we create a Tree as the root
end

def current_sigs
  sigs = @last_sigs.dup
  @last_sigs.clear
  sigs
end

def detach_comments_from_sigs(sigs)
  comments = [] #: Array[Comment]
  sigs.each do |sig|
    comments += sig.comments.dup
    sig.comments.clear
  end
  comments
end

def initialize(source, comments:, file:)
  super(source, file: file)
  @comments_by_line = comments.to_h { |c| [c.location.start_line, c] } #: Hash[Integer, Prism::Comment]
  @tree = Tree.new #: Tree
  @scopes_stack = [@tree] #: Array[Tree]
  @last_node = nil #: Prism::Node?
  @last_sigs = [] #: Array[RBI::Sig]
end

def node_comments(node)
  comments = []
  start_line = node.location.start_line
  start_line -= 1 unless @comments_by_line.key?(start_line)
  rbs_continuation = [] #: Array[Prism::Comment]
  start_line.downto(1) do |line|
    comment = @comments_by_line[line]
    break unless comment
    text = comment.location.slice
    # If we find a RBS comment continuation `#|`, we store it until we find the start with `#:`
    if text.start_with?("#|")
      rbs_continuation << comment
      @comments_by_line.delete(line)
      next
    end
    loc = Loc.from_prism(@file, comment.location)
    # If we find the start of a RBS comment, we create a new RBSComment
    # Note that we ignore RDoc directives such as `:nodoc:`
    # See https://ruby.github.io/rdoc/RDoc/MarkupReference.html#class-RDoc::MarkupReference-label-Directives
    if text.start_with?("#:") && !(text =~ /^#:[a-z_]+:/)
      text = text.sub(/^#: ?/, "").rstrip
      # If we found continuation comments, we merge them in reverse order (since we go from bottom to top)
      rbs_continuation.reverse_each do |rbs_comment|
        continuation_text = rbs_comment.location.slice.sub(/^#\| ?/, "").strip
        continuation_loc = Loc.from_prism(@file, rbs_comment.location)
        loc = loc.join(continuation_loc)
        text = "#{text}#{continuation_text}"
      end
      rbs_continuation.clear
      comments.unshift(RBSComment.new(text, loc: loc))
    else
      # If we have unused continuation comments, we should inject them back to not lose them
      rbs_continuation.each do |rbs_comment|
        comments.unshift(parse_comment(rbs_comment))
      end
      rbs_continuation.clear
      comments.unshift(parse_comment(comment))
    end
    @comments_by_line.delete(line)
  end
  # If we have unused continuation comments, we should inject them back to not lose them
  rbs_continuation.each do |rbs_comment|
    comments.unshift(parse_comment(rbs_comment))
  end
  rbs_continuation.clear
  comments
end

def parse_comment(node)
  text = node.location.slice.sub(/^# ?/, "").rstrip
  loc = Loc.from_prism(@file, node.location)
  Comment.new(text, loc: loc)
end

def parse_params(node)
  return [] unless node.is_a?(Prism::ParametersNode)
  node_params = [
    *node.requireds,
    *node.optionals,
    *node.rest,
    *node.posts,
    *node.keywords,
    *node.keyword_rest,
    *node.block,
  ].flatten
  node_params.map do |param|
    case param
    when Prism::RequiredParameterNode
      ReqParam.new(
        param.name.to_s,
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::OptionalParameterNode
      OptParam.new(
        param.name.to_s,
        node_string!(param.value),
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::RestParameterNode
      RestParam.new(
        param.name&.to_s || "*args",
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::RequiredKeywordParameterNode
      KwParam.new(
        param.name.to_s.delete_suffix(":"),
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::OptionalKeywordParameterNode
      KwOptParam.new(
        param.name.to_s.delete_suffix(":"),
        node_string!(param.value),
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::KeywordRestParameterNode
      KwRestParam.new(
        param.name&.to_s || "**kwargs",
        loc: node_loc(param),
        comments: node_comments(param),
      )
    when Prism::BlockParameterNode
      BlockParam.new(
        param.name&.to_s || "&block",
        loc: node_loc(param),
        comments: node_comments(param),
      )
    else
      raise ParseError.new("Unexpected parameter node `#{param.class}`", node_loc(param))
    end
  end
end

def parse_send_args(node)
  args = [] #: Array[Arg]
  return args unless node.is_a?(Prism::ArgumentsNode)
  node.arguments.each do |arg|
    case arg
    when Prism::KeywordHashNode
      arg.elements.each do |assoc|
        next unless assoc.is_a?(Prism::AssocNode)
        args << KwArg.new(
          node_string!(assoc.key).delete_suffix(":"),
          node_string(assoc.value), #: as !nil
        )
      end
    else
      args << Arg.new(
        node_string(arg), #: as !nil
      )
    end
  end
  args
end

def parse_sig(node)
  builder = SigBuilder.new(@source, file: @file)
  builder.current.loc = node_loc(node)
  builder.visit_call_node(node)
  builder.current.comments = node_comments(node)
  builder.current
end

def parse_struct(node)
  send = node.value
  return unless send.is_a?(Prism::CallNode)
  return unless send.message == "new"
  recv = send.receiver
  return unless recv
  return unless node_string(recv) =~ /(::)?Struct/
  members = []
  keyword_init = false #: bool
  args = send.arguments
  if args.is_a?(Prism::ArgumentsNode)
    args.arguments.each do |arg|
      case arg
      when Prism::SymbolNode
        members << arg.value
      when Prism::KeywordHashNode
        arg.elements.each do |assoc|
          next unless assoc.is_a?(Prism::AssocNode)
          key = node_string!(assoc.key)
          val = node_string(assoc.value)
          keyword_init = val == "true" if key == "keyword_init:"
        end
      end
    end
  end
  name = case node
  when Prism::ConstantWriteNode
    node.name.to_s
  when Prism::ConstantPathWriteNode
    node_string!(node.target)
  end
  loc = node_loc(node)
  comments = node_comments(node)
  struct = Struct.new(name, members: members, keyword_init: keyword_init, loc: loc, comments: comments)
  @scopes_stack << struct
  visit(send.block)
  @scopes_stack.pop
  struct
end

def parse_tstruct_field(send)
  args = send.arguments
  return unless args.is_a?(Prism::ArgumentsNode)
  name_arg, type_arg, *rest = args.arguments
  return unless name_arg
  return unless type_arg
  name = node_string!(name_arg).delete_prefix(":")
  type = node_string!(type_arg)
  loc = node_loc(send)
  comments = node_comments(send)
  default_value = nil #: String?
  rest.each do |arg|
    next unless arg.is_a?(Prism::KeywordHashNode)
    arg.elements.each do |assoc|
      next unless assoc.is_a?(Prism::AssocNode)
      if node_string(assoc.key) == "default:"
        default_value = node_string(assoc.value)
      end
    end
  end
  current_scope << case send.message
  when "const"
    TStructConst.new(name, type, default: default_value, loc: loc, comments: comments)
  when "prop"
    TStructProp.new(name, type, default: default_value, loc: loc, comments: comments)
  else
    raise ParseError.new("Unexpected message `#{send.message}`", loc)
  end
end

def parse_visibility(name, node)
  case name
  when "public"
    Public.new(loc: node_loc(node), comments: node_comments(node))
  when "protected"
    Protected.new(loc: node_loc(node), comments: node_comments(node))
  when "private"
    Private.new(loc: node_loc(node), comments: node_comments(node))
  else
    raise ParseError.new("Unexpected visibility `#{name}`", node_loc(node))
  end
end

def separate_header_comments
  current_scope.nodes.dup.each do |child_node|
    break unless child_node.is_a?(Comment) || child_node.is_a?(BlankLine)
    current_scope.comments << child_node
    child_node.detach
  end
end

def set_root_tree_loc
  first_loc = tree.nodes.first&.loc
  last_loc = tree.nodes.last&.loc
  @tree.loc = Loc.new(
    file: @file,
    begin_line: first_loc&.begin_line || 0,
    begin_column: first_loc&.begin_column || 0,
    end_line: last_loc&.end_line || 0,
    end_column: last_loc&.end_column || 0,
  )
end

def t_enum_value?(node)
  return false unless current_scope.is_a?(TEnumBlock)
  return false unless node.is_a?(Prism::ConstantWriteNode)
  value = node.value
  return false unless value.is_a?(Prism::CallNode)
  return false unless value.message == "new"
  true
end

def type_variable_definition?(node)
  node.is_a?(Prism::CallNode) && (node.message == "type_member" || node.message == "type_template")
end

def visit_call_node(node)
  @last_node = node
  message = node.name.to_s
  case message
  when "abstract!", "sealed!", "interface!"
    current_scope << Helper.new(
      message.delete_suffix("!"),
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when "attr_reader"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    sigs = current_sigs
    comments = detach_comments_from_sigs(sigs) + node_comments(node)
    current_scope << AttrReader.new(
      *args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym },
      sigs: sigs,
      loc: node_loc(node),
      comments: comments,
    )
  when "attr_writer"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    sigs = current_sigs
    comments = detach_comments_from_sigs(sigs) + node_comments(node)
    current_scope << AttrWriter.new(
      *args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym },
      sigs: sigs,
      loc: node_loc(node),
      comments: comments,
    )
  when "attr_accessor"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    sigs = current_sigs
    comments = detach_comments_from_sigs(sigs) + node_comments(node)
    current_scope << AttrAccessor.new(
      *args.arguments.map { |arg| node_string!(arg).delete_prefix(":").to_sym },
      sigs: sigs,
      loc: node_loc(node),
      comments: comments,
    )
  when "enums"
    if node.block && node.arguments.nil?
      scope = TEnumBlock.new(loc: node_loc(node), comments: node_comments(node))
      current_scope << scope
      @scopes_stack << scope
      visit(node.block)
      @scopes_stack.pop
    else
      current_scope << Send.new(
        message,
        parse_send_args(node.arguments),
        loc: node_loc(node),
        comments: node_comments(node),
      )
    end
  when "extend"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    current_scope << Extend.new(
      *args.arguments.map { |arg| node_string!(arg) },
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when "include"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    current_scope << Include.new(
      *args.arguments.map { |arg| node_string!(arg) },
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when "mixes_in_class_methods"
    args = node.arguments
    unless args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      @last_node = nil
      return
    end
    current_scope << MixesInClassMethods.new(
      *args.arguments.map { |arg| node_string!(arg) },
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when "private", "protected", "public"
    args = node.arguments
    if args.is_a?(Prism::ArgumentsNode) && args.arguments.any?
      visit(node.arguments)
      last_node = @scopes_stack.last&.nodes&.last
      case last_node
      when Method, Attr
        last_node.visibility = parse_visibility(node.name.to_s, node)
      when Send
        current_scope << Send.new(
          message,
          parse_send_args(node.arguments),
          loc: node_loc(node),
          comments: node_comments(node),
        )
      end
    else
      current_scope << parse_visibility(node.name.to_s, node)
    end
  when "prop", "const"
    parse_tstruct_field(node)
  when "requires_ancestor"
    block = node.block
    unless block.is_a?(Prism::BlockNode)
      @last_node = nil
      return
    end
    body = block.body
    unless body.is_a?(Prism::StatementsNode)
      @last_node = nil
      return
    end
    current_scope << RequiresAncestor.new(
      node_string!(body),
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when "sig"
    unless node.receiver.nil? || self?(node.receiver) || t_sig_without_runtime?(node.receiver)
      @last_node = nil
      return
    end
    @last_sigs << parse_sig(node)
  else
    current_scope << Send.new(
      message,
      parse_send_args(node.arguments),
      loc: node_loc(node),
      comments: node_comments(node),
    )
  end
  @last_node = nil
end

def visit_class_node(node)
  @last_node = node
  superclass_name = node_string(node.superclass)
  scope = case superclass_name
  when /^(::)?T::Struct$/
    TStruct.new(
      node_string!(node.constant_path),
      loc: node_loc(node),
      comments: node_comments(node),
    )
  when /^(::)?T::Enum$/
    TEnum.new(
      node_string!(node.constant_path),
      loc: node_loc(node),
      comments: node_comments(node),
    )
  else
    Class.new(
      node_string!(node.constant_path),
      superclass_name: superclass_name,
      loc: node_loc(node),
      comments: node_comments(node),
    )
  end
  current_scope << scope
  @scopes_stack << scope
  visit(node.body)
  scope.nodes.concat(current_sigs)
  collect_dangling_comments(node)
  @scopes_stack.pop
  @last_node = nil
end

def visit_constant_assign(node)
  struct = parse_struct(node)
  current_scope << if struct
    struct
  elsif t_enum_value?(node)
    TEnumValue.new(
      case node
      when Prism::ConstantWriteNode
        node.name.to_s
      when Prism::ConstantPathWriteNode
        node_string!(node.target)
      end,
      loc: node_loc(node),
      comments: node_comments(node),
    )
  else
    adjusted_node_location = adjust_prism_location_for_heredoc(node)
    adjusted_value_location = Prism::Location.new(
      node.value.location.send(:source),
      node.value.location.start_offset,
      adjusted_node_location.end_offset - node.value.location.start_offset,
    )
    if type_variable_definition?(node.value)
      TypeMember.new(
        case node
        when Prism::ConstantWriteNode
          node.name.to_s
        when Prism::ConstantPathWriteNode
          node_string!(node.target)
        end,
        adjusted_value_location.slice,
        loc: Loc.from_prism(@file, adjusted_node_location),
        comments: node_comments(node),
      )
    else
      Const.new(
        case node
        when Prism::ConstantWriteNode
          node.name.to_s
        when Prism::ConstantPathWriteNode
          node_string!(node.target)
        end,
        adjusted_value_location.slice,
        loc: Loc.from_prism(@file, adjusted_node_location),
        comments: node_comments(node),
      )
    end
  end
end

def visit_constant_path_write_node(node)
  @last_node = node
  visit_constant_assign(node)
  @last_node = nil
end

def visit_constant_write_node(node)
  @last_node = node
  visit_constant_assign(node)
  @last_node = nil
end

def visit_def_node(node)
  @last_node = node
  # We need to collect the comments with `current_sigs_comments` _before_ visiting the parameters to make sure
  # the method comments are properly associated with the sigs and not the parameters.
  sigs = current_sigs
  comments = detach_comments_from_sigs(sigs) + node_comments(node)
  params = parse_params(node.parameters)
  current_scope << Method.new(
    node.name.to_s,
    params: params,
    sigs: sigs,
    loc: node_loc(node),
    comments: comments,
    is_singleton: !!node.receiver,
  )
  @last_node = nil
end

def visit_module_node(node)
  @last_node = node
  scope = Module.new(
    node_string!(node.constant_path),
    loc: node_loc(node),
    comments: node_comments(node),
  )
  current_scope << scope
  @scopes_stack << scope
  visit(node.body)
  scope.nodes.concat(current_sigs)
  collect_dangling_comments(node)
  @scopes_stack.pop
  @last_node = nil
end

def visit_program_node(node)
  @last_node = node
  super
  @tree.nodes.concat(current_sigs)
  collect_orphan_comments
  separate_header_comments
  set_root_tree_loc
  @last_node = nil
end

def visit_singleton_class_node(node)
  @last_node = node
  scope = SingletonClass.new(
    loc: node_loc(node),
    comments: node_comments(node),
  )
  current_scope << scope
  @scopes_stack << scope
  visit(node.body)
  scope.nodes.concat(current_sigs)
  collect_dangling_comments(node)
  @scopes_stack.pop
  @last_node = nil
end