class PSTree

# File lib/pstree.rb, line 18
def initialize(root_pid = nil, charset: 'UTF-8')
  @charset  = charset.to_s.upcase
  @root_pid = root_pid.to_i
end

# File lib/pstree.rb, line 31
def children_not_zero(last)
  if @charset == 'UTF-8'
    last ? '├─ ' : '│  '
  else
    last ? '+- ' : '|  '
  end
end

# File lib/pstree.rb, line 23
def children_zero(last)
  if @charset == 'UTF-8'
    last ? '└─ ' : '   '
  else
    last ? '`- ' : '   '
  end
end

# File lib/pstree.rb, line 55
def each(&block)
  build
  recurse @root_pid, &block
  self
end

# File lib/pstree.rb, line 39
def to_s
  build
  result = ''
  recurse @root_pid,
    -> children, last {
      if children.zero?
        result << children_zero(last)
      else
        result << children_not_zero(last)
      end
    } do |ps|
    result << ps.to_s << "\n"
  end
  result
end

# File lib/pstree.rb, line 63
def build
  @child_count = [ 0 ]
  @process = {}
  @pstree = Hash.new { |h,k| h[k] = [] }
  pid = @root_pid.zero? ? nil : @root_pid
  psoutput = `/bin/ps axww -o ppid,pid,user,command`
  psoutput.each_line do |line|
    next if line !~ /^\s*\d+/
    line.strip!
    ps = ProcStruct.new(*line.split(/\s+/, 4))
    @process[ps.pid] = ps
    @pstree[ps.ppid] << ps
  end
end

# File lib/pstree.rb, line 78
def recurse(pid, shift_callback = nil, level = 0, in_root = false, &node_callback)
  in_root = in_root || @root_pid == 0 || @root_pid == pid
  @child_count[level] = @pstree[pid].size
  for l in 0...level
    shift_callback and shift_callback.call(@child_count[l], l == level - 1)
  end
  node_callback and process = @process[pid] and node_callback.call(process)
  if @pstree.key?(pid)
    @child_count[level] = @pstree[pid].size - 1
    @pstree[pid].each do |ps|
      recurse(ps.pid, shift_callback, level + 1, in_root, &node_callback)
      @child_count[level] -= 1
    end
  end
  @pstree.delete pid
end