class SyntaxTree::CallChainFormatter

def self.chained?(node)
  return false if ENV["STREE_FAST_FORMAT"]
  case node
  when CallNode
    !node.receiver.nil?
  when MethodAddBlock
    call = node.call
    call.is_a?(CallNode) && !call.receiver.nil?
  else
    false
  end
end

def attach_directly?(node)
  case node.receiver
  when ArrayLiteral, HashLiteral, Heredoc, IfNode, UnlessNode,
       XStringLiteral
    true
  else
    false
  end
end

def format(q)
  children = [node]
  threshold = 3
  # First, walk down the chain until we get to the point where we're not
  # longer at a chainable node.
  loop do
    case (child = children.last)
    when CallNode
      case (receiver = child.receiver)
      when CallNode
        if receiver.receiver.nil?
          break
        else
          children << receiver
        end
      when MethodAddBlock
        if (call = receiver.call).is_a?(CallNode) && !call.receiver.nil?
          children << receiver
        else
          break
        end
      else
        break
      end
    when MethodAddBlock
      if (call = child.call).is_a?(CallNode) && !call.receiver.nil?
        children << call
      else
        break
      end
    else
      break
    end
  end
  # Here, we have very specialized behavior where if we're within a sig
  # block, then we're going to assume we're creating a Sorbet type
  # signature. In that case, we really want the threshold to be lowered so
  # that we create method chains off of any two method calls within the
  # block. For more details, see
  # https://github.com/prettier/plugin-ruby/issues/863.
  parents = q.parents.take(4)
  if (parent = parents[2])
    # If we're at a block with the `do` keywords, then we want to go one
    # more level up. This is because do blocks have BodyStmt nodes instead
    # of just Statements nodes.
    parent = parents[3] if parent.is_a?(BlockNode) && parent.keywords?
    if parent.is_a?(MethodAddBlock) &&
         (call = parent.call).is_a?(CallNode) && call.message.value == "sig"
      threshold = 2
    end
  end
  if children.length >= threshold
    q.group do
      q
        .if_break { format_chain(q, children) }
        .if_flat { node.format_contents(q) }
    end
  else
    node.format_contents(q)
  end
end

def format_chain(q, children)
  # We're going to have some specialized behavior for if it's an entire
  # chain of calls without arguments except for the last one. This is common
  # enough in Ruby source code that it's worth the extra complexity here.
  empty_except_last =
    children
      .drop(1)
      .all? { |child| child.is_a?(CallNode) && child.arguments.nil? }
  # Here, we're going to add all of the children onto the stack of the
  # formatter so it's as if we had descending normally into them. This is
  # necessary so they can check their parents as normal.
  q.stack.concat(children)
  q.format(children.last.receiver) if children.last.receiver
  q.group do
    if attach_directly?(children.last)
      format_child(q, children.pop)
      q.stack.pop
    end
    q.indent do
      # We track another variable that checks if you need to move the
      # operator to the previous line in case there are trailing comments
      # and a trailing operator.
      skip_operator = false
      while (child = children.pop)
        if child.is_a?(CallNode)
          if (receiver = child.receiver).is_a?(CallNode) &&
               (receiver.message != :call) &&
               (receiver.message.value == "where") &&
               (child.message != :call && child.message.value == "not")
            # This is very specialized behavior wherein we group
            # .where.not calls together because it looks better. For more
            # information, see
            # https://github.com/prettier/plugin-ruby/issues/862.
          else
            # If we're at a Call node and not a MethodAddBlock node in the
            # chain then we're going to add a newline so it indents
            # properly.
            q.breakable_empty
          end
        end
        format_child(
          q,
          child,
          skip_comments: children.empty?,
          skip_operator: skip_operator,
          skip_attached: empty_except_last && children.empty?
        )
        # If the parent call node has a comment on the message then we need
        # to print the operator trailing in order to keep it working.
        last_child = children.last
        if last_child.is_a?(CallNode) && last_child.message != :call &&
             last_child.message.comments.any? && last_child.operator
          q.format(CallOperatorFormatter.new(last_child.operator))
          skip_operator = true
        else
          skip_operator = false
        end
        # Pop off the formatter's stack so that it aligns with what would
        # have happened if we had been formatting normally.
        q.stack.pop
      end
    end
  end
  if empty_except_last
    case node
    when CallNode
      node.format_arguments(q)
    when MethodAddBlock
      q.format(node.block)
    end
  end
end

def format_child(
  q,
  child,
  skip_comments: false,
  skip_operator: false,
  skip_attached: false
)
  # First, format the actual contents of the child.
  case child
  when CallNode
    q.group do
      if !skip_operator && child.operator
        q.format(CallOperatorFormatter.new(child.operator))
      end
      q.format(child.message) if child.message != :call
      child.format_arguments(q) unless skip_attached
    end
  when MethodAddBlock
    q.format(child.block) unless skip_attached
  end
  # If there are any comments on this node then we need to explicitly print
  # them out here since we're bypassing the normal comment printing.
  if child.comments.any? && !skip_comments
    child.comments.each do |comment|
      comment.inline? ? q.text(" ") : q.breakable_space
      comment.format(q)
    end
    q.break_parent
  end
end

def initialize(node)
  @node = node
end